Introduce truss server passthrough for OpenAI methods (#1364)

Author: nnarayen

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.60rc004"
+version = "0.9.60rc005"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/templates/server/model_wrapper.py

@@ -7,7 +7,7 @@
 import sys
 from http import HTTPStatus
 from pathlib import Path
-from typing import Dict, Optional, Union
+from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Optional, Union
 
 import pydantic
 import uvicorn
@@ -17,7 +17,7 @@
 from fastapi import Depends, FastAPI, HTTPException, Request
 from fastapi.responses import ORJSONResponse, StreamingResponse
 from fastapi.routing import APIRoute as FastAPIRoute
-from model_wrapper import InputType, ModelWrapper, OutputType
+from model_wrapper import MODEL_BASENAME, MethodName, ModelWrapper
 from opentelemetry import propagate as otel_propagate
 from opentelemetry import trace
 from opentelemetry.sdk import trace as sdk_trace
@@ -38,6 +38,9 @@
 TIMEOUT_GRACEFUL_SHUTDOWN = 120
 INFERENCE_SERVER_FAILED_FILE = Path("~/inference_server_crashed.txt").expanduser()
 
+if TYPE_CHECKING:
+    from model_wrapper import InputType, MethodDescriptor, OutputType
+
 
 async def parse_body(request: Request) -> bytes:
     """
@@ -63,7 +66,7 @@ def __init__(self, model: ModelWrapper, tracer: sdk_trace.Tracer) -> None:
         self._model = model
         self._tracer = tracer
 
-    def _safe_lookup_model(self, model_name: str) -> ModelWrapper:
+    def _safe_lookup_model(self, model_name: str = MODEL_BASENAME) -> ModelWrapper:
         if model_name != self._model.name:
             raise errors.ModelMissingError(model_name)
         return self._model
@@ -116,7 +119,7 @@ async def _parse_body(
         body_raw: bytes,
         truss_schema: Optional[TrussSchema],
         span: trace.Span,
-    ) -> InputType:
+    ) -> "InputType":
         if self.is_binary(request):
             with tracing.section_as_event(span, "binary-deserialize"):
                 inputs = serialization.truss_msgpack_deserialize(body_raw)
@@ -148,36 +151,38 @@ async def _parse_body(
 
         return inputs
 
-    async def predict(
-        self, model_name: str, request: Request, body_raw: bytes = Depends(parse_body)
+    async def _execute_request(
+        self,
+        model: ModelWrapper,
+        method: Callable[["InputType", Request], Awaitable["OutputType"]],
+        method_name: MethodName,
+        request: Request,
+        body_raw: bytes,
     ) -> Response:
         """
-        This method calls the user-provided predict method
+        Executes a predictive endpoint
         """
         if await request.is_disconnected():
-            msg = "Client disconnected. Skipping `predict`."
+            msg = f"Client disconnected. Skipping `{method_name}`."
             logging.info(msg)
             raise ClientDisconnect(msg)
 
-        model: ModelWrapper = self._safe_lookup_model(model_name)
-
         self.check_healthy(model)
         trace_ctx = otel_propagate.extract(request.headers) or None
         # This is the top-level span in the truss-server, so we set the context here.
         # Nested spans "inherit" context automatically.
         with self._tracer.start_as_current_span(
-            "predict-endpoint", context=trace_ctx
+            f"{method_name}-endpoint", context=trace_ctx
         ) as span:
-            inputs: Optional[InputType]
+            inputs: Optional["InputType"]
             if model.model_descriptor.skip_input_parsing:
                 inputs = None
             else:
                 inputs = await self._parse_body(
                     request, body_raw, model.model_descriptor.truss_schema, span
                 )
-            # Calls ModelWrapper which runs: preprocess, predict, postprocess.
             with tracing.section_as_event(span, "model-call"):
-                result: OutputType = await model(inputs, request)
+                result: "OutputType" = await method(inputs, request)
 
             # In the case that the model returns a Generator object, return a
             # StreamingResponse instead.
@@ -190,8 +195,59 @@ async def predict(
                 return result
             return self._serialize_result(result, self.is_binary(request), span)
 
+    async def chat_completions(
+        self, request: Request, body_raw: bytes = Depends(parse_body)
+    ) -> Response:
+        model = self._safe_lookup_model()
+        self._raise_if_not_supported(
+            MethodName.CHAT_COMPLETIONS, model.model_descriptor.chat_completions
+        )
+
+        return await self._execute_request(
+            model=model,
+            method=model.chat_completions,
+            method_name=MethodName.CHAT_COMPLETIONS,
+            request=request,
+            body_raw=body_raw,
+        )
+
+    def _raise_if_not_supported(
+        self, method_name: MethodName, descriptor: Optional["MethodDescriptor"]
+    ):
+        if not descriptor:
+            raise HTTPException(status_code=404, detail=f"{method_name} not supported.")
+
+    async def completions(
+        self, request: Request, body_raw: bytes = Depends(parse_body)
+    ) -> Response:
+        model = self._safe_lookup_model()
+        self._raise_if_not_supported(
+            MethodName.COMPLETIONS, model.model_descriptor.completions
+        )
+
+        return await self._execute_request(
+            model=model,
+            method=model.completions,
+            method_name=MethodName.COMPLETIONS,
+            request=request,
+            body_raw=body_raw,
+        )
+
+    async def predict(
+        self, model_name: str, request: Request, body_raw: bytes = Depends(parse_body)
+    ) -> Response:
+        model = self._safe_lookup_model(model_name)
+
+        return await self._execute_request(
+            model=model,
+            method=model,  # We overwrote __call__ on ModelWrapper
+            method_name=MethodName.PREDICT,
+            request=request,
+            body_raw=body_raw,
+        )
+
     def _serialize_result(
-        self, result: OutputType, is_binary: bool, span: trace.Span
+        self, result: "OutputType", is_binary: bool, span: trace.Span
     ) -> Response:
         response_headers = {}
         if is_binary:
@@ -338,6 +394,19 @@ def create_application(self):
                     methods=["POST"],
                     tags=["V1"],
                 ),
+                # OpenAI Spec
+                FastAPIRoute(
+                    r"/v1/chat/completions",
+                    self._endpoints.chat_completions,
+                    methods=["POST"],
+                    tags=["V1"],
+                ),
+                FastAPIRoute(
+                    r"/v1/completions",
+                    self._endpoints.completions,
+                    methods=["POST"],
+                    tags=["V1"],
+                ),
                 # Endpoint aliases for Sagemaker hosting
                 FastAPIRoute(r"/ping", self._endpoints.invocations_ready),
                 FastAPIRoute(

truss/templates/server/truss_server.py

@@ -553,6 +553,11 @@ def trt_llm_truss_container_fs(tmp_path, test_data_path):
     return _build_truss_fs(test_data_path / "test_trt_llm_truss", tmp_path)
 
 
+@pytest.fixture
+def open_ai_container_fs(tmp_path, test_data_path):
+    return _build_truss_fs(test_data_path / "test_openai", tmp_path)
+
+
 @pytest.fixture
 def truss_control_container_fs(tmp_path, test_data_path):
     test_truss_dir = test_data_path / "test_truss"

truss/tests/conftest.py

@@ -190,6 +190,30 @@ async def mock_predict(return_value, request: Request):
             assert resp == expected_predict_response
 
 
+@pytest.mark.anyio
+async def test_open_ai_completion_endpoints(open_ai_container_fs, helpers):
+    app_path = open_ai_container_fs / "app"
+    with _clear_model_load_modules(), helpers.sys_paths(app_path), _change_directory(
+        app_path
+    ):
+        model_wrapper_module = importlib.import_module("model_wrapper")
+        model_wrapper_class = getattr(model_wrapper_module, "ModelWrapper")
+        config = yaml.safe_load((app_path / "config.yaml").read_text())
+
+        model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
+        model_wrapper.load()
+
+        mock_req = MagicMock(spec=Request)
+        predict_resp = await model_wrapper.predict({}, mock_req)
+        assert predict_resp == "predict"
+
+        completions_resp = await model_wrapper.completions({}, mock_req)
+        assert completions_resp == "completions"
+
+        chat_completions_resp = await model_wrapper.chat_completions({}, mock_req)
+        assert chat_completions_resp == "chat_completions"
+
+
 @contextmanager
 def _change_directory(new_directory: Path):
     original_directory = os.getcwd()

truss/tests/templates/server/test_model_wrapper.py

@@ -0,0 +1,12 @@
+environment_variables: {}
+external_package_dirs: []
+model_metadata: {}
+model_name: Test OpenAI Compatibility
+python_version: py39
+resources:
+  accelerator: null
+  cpu: '1'
+  memory: 2Gi
+  use_gpu: false
+secrets: {}
+system_packages: []

truss/tests/test_data/test_openai/__init__.py

@@ -0,0 +1,15 @@
+from typing import Dict
+
+
+class Model:
+    def __init__(self, **kwargs):
+        pass
+
+    def chat_completions(self, input: Dict) -> str:
+        return "chat_completions"
+
+    def completions(self, input: Dict) -> str:
+        return "completions"
+
+    def predict(self, input: Dict) -> str:
+        return "predict"

truss/tests/test_data/test_openai/config.yaml

@@ -33,6 +33,8 @@
 
 DEFAULT_LOG_ERROR = "Internal Server Error"
 PREDICT_URL = "http://localhost:8090/v1/models/model:predict"
+COMPLETIONS_URL = "http://localhost:8090/v1/completions"
+CHAT_COMPLETIONS_URL = "http://localhost:8090/v1/chat/completions"
 
 
 @pytest.fixture
@@ -1714,3 +1716,105 @@ def make_request(consume_chunks, timeout, task_id):
 
         result = make_request(True, timeout=0.55, task_id=4)
         print(f"Final chunks: {result}")
+
+
+@pytest.mark.integration
+def test_custom_openai_endpoints():
+    """
+    Test a Truss that exposes an OpenAI compatible endpoint.
+    """
+    model = """
+    from typing import Dict
+
+    class Model:
+        def __init__(self):
+            pass
+
+        def load(self):
+            self._predict_count = 0
+            self._completions_count = 0
+
+        async def predict(self, inputs: Dict) -> int:
+            self._predict_count += inputs["increment"]
+            return self._predict_count
+
+        async def completions(self, inputs: Dict) -> int:
+            self._completions_count += inputs["increment"]
+            return self._completions_count
+    """
+    with ensure_kill_all(), _temp_truss(model) as tr:
+        tr.docker_run(local_port=8090, detach=True, wait_for_server_ready=True)
+
+        response = requests.post(PREDICT_URL, json={"increment": 1})
+        assert response.status_code == 200
+        assert response.json() == 1
+
+        response = requests.post(COMPLETIONS_URL, json={"increment": 2})
+        assert response.status_code == 200
+        assert response.json() == 2
+
+        response = requests.post(CHAT_COMPLETIONS_URL, json={"increment": 3})
+        assert response.status_code == 404
+
+
+@pytest.mark.integration
+def test_postprocess_async_generator_streaming():
+    """
+    Test a Truss that exposes an OpenAI compatible endpoint.
+    """
+    model = """
+    from typing import Dict, List, Generator
+
+    class Model:
+        def __init__(self):
+            pass
+
+        def load(self):
+            pass
+
+        async def predict(self, inputs: Dict) -> List[str]:
+            nums: List[int] = inputs["nums"]
+            return nums
+
+        async def postprocess(self, nums: List[str]) -> Generator[str, None, None]:
+            for num in nums:
+                yield num
+    """
+    with ensure_kill_all(), _temp_truss(model) as tr:
+        tr.docker_run(local_port=8090, detach=True, wait_for_server_ready=True)
+
+        response = requests.post(PREDICT_URL, json={"nums": ["1", "2"]}, stream=True)
+        assert response.headers.get("transfer-encoding") == "chunked"
+        assert [
+            byte_string.decode() for byte_string in list(response.iter_content())
+        ] == ["1", "2"]
+
+
+@pytest.mark.integration
+def test_preprocess_async_generator():
+    """
+    Test a Truss that exposes an OpenAI compatible endpoint.
+    """
+    model = """
+    from typing import Dict, List, AsyncGenerator
+
+    class Model:
+        def __init__(self):
+            pass
+
+        def load(self):
+            pass
+
+        async def preprocess(self, inputs: Dict) -> AsyncGenerator[str, None]:
+            for num in inputs["nums"]:
+                yield num
+
+        async def predict(self, nums: AsyncGenerator[str, None]) -> List[str]:
+            return [num async for num in nums]
+    """
+    with ensure_kill_all(), _temp_truss(model) as tr:
+        tr.docker_run(local_port=8090, detach=True, wait_for_server_ready=True)
+
+        response = requests.post(PREDICT_URL, json={"nums": ["1", "2"]})
+        assert response.status_code == 200
+        assert response.json() == ["1", "2"]