Add in_process mode support for DJL and TorchServe servers (#1570)

* add in-process mode for DJL server

* fix format

* add inference_spec as a member of DJL

* add the validations for model server

* fix typo

* fix test assertion

* add unit-testing

* have a common server for inprocess mode

* fix failing tests

* add support to torchserve

* fix tests to include torchserve servers

* use custom inference_spec code instead of HF pipelines

* fix tests for app.py

* fix unit test failure

* fix format

* use schema_builder for serialization and deserialization

* remove task field

* remove unused import

Author: pravali96

src/sagemaker/serve/builder/djl_builder.py

@@ -47,14 +47,15 @@
 from sagemaker.serve.model_server.djl_serving.prepare import (
     _create_dir_structure,
 )
-from sagemaker.serve.utils.predictors import DjlLocalModePredictor
+from sagemaker.serve.utils.predictors import InProcessModePredictor, DjlLocalModePredictor
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
 from sagemaker.djl_inference.model import DJLModel
 from sagemaker.base_predictor import PredictorBase
 
 logger = logging.getLogger(__name__)
+LOCAL_MODES = [Mode.LOCAL_CONTAINER, Mode.IN_PROCESS]
 
 # Match JumpStart DJL entrypoint format
 _CODE_FOLDER = "code"
@@ -77,6 +78,7 @@ def __init__(self):
         self.mode = None
         self.model_server = None
         self.image_uri = None
+        self.inference_spec = None
         self._is_custom_image_uri = False
         self.image_config = None
         self.vpc_config = None
@@ -96,11 +98,11 @@ def __init__(self):
 
     @abstractmethod
     def _prepare_for_mode(self):
-        """Placeholder docstring"""
+        """Abstract method"""
 
     @abstractmethod
     def _get_client_translators(self):
-        """Placeholder docstring"""
+        """Abstract method"""
 
     def _is_djl(self):
         """Placeholder docstring"""
@@ -146,7 +148,7 @@ def _create_djl_model(self) -> Type[Model]:
 
     @_capture_telemetry("djl.deploy")
     def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
-        """Placeholder docstring"""
+        """Returns predictor depending on local mode or endpoint mode"""
         timeout = kwargs.get("model_data_download_timeout")
         if timeout:
             self.env_vars.update({"MODEL_LOADING_TIMEOUT": str(timeout)})
@@ -189,6 +191,18 @@ def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
 
         serializer = self.schema_builder.input_serializer
         deserializer = self.schema_builder._output_deserializer
+
+        if self.mode == Mode.IN_PROCESS:
+
+            predictor = InProcessModePredictor(
+                self.modes[str(Mode.IN_PROCESS)], serializer, deserializer
+            )
+
+            self.modes[str(Mode.IN_PROCESS)].create_server(
+                predictor,
+            )
+            return predictor
+
         if self.mode == Mode.LOCAL_CONTAINER:
             timeout = kwargs.get("model_data_download_timeout")
 
@@ -249,9 +263,15 @@ def _build_for_hf_djl(self):
 
         _create_dir_structure(self.model_path)
         if not hasattr(self, "pysdk_model"):
-            self.env_vars.update({"HF_MODEL_ID": self.model})
+            if self.inference_spec is not None:
+                self.env_vars.update({"HF_MODEL_ID": self.inference_spec.get_model()})
+            else:
+                self.env_vars.update({"HF_MODEL_ID": self.model})
+
+            logger.info(self.env_vars)
+
             self.hf_model_config = _get_model_config_properties_from_hf(
-                self.model, self.env_vars.get("HF_TOKEN")
+                self.env_vars.get("HF_MODEL_ID"), self.env_vars.get("HF_TOKEN")
             )
             default_djl_configurations, _default_max_new_tokens = _get_default_djl_configurations(
                 self.model, self.hf_model_config, self.schema_builder
@@ -260,9 +280,10 @@ def _build_for_hf_djl(self):
             self.schema_builder.sample_input["parameters"][
                 "max_new_tokens"
             ] = _default_max_new_tokens
+
         self.pysdk_model = self._create_djl_model()
 
-        if self.mode == Mode.LOCAL_CONTAINER:
+        if self.mode in LOCAL_MODES:
             self._prepare_for_mode()
 
         return self.pysdk_model
@@ -451,7 +472,6 @@ def _build_for_djl(self):
         """Placeholder docstring"""
         self._validate_djl_serving_sample_data()
         self.secret_key = None
-
         self.pysdk_model = self._build_for_hf_djl()
         self.pysdk_model.tune = self._tune_for_hf_djl
         if self.role_arn:

src/sagemaker/serve/builder/model_builder.py

@@ -81,7 +81,7 @@
     _extract_speculative_draft_model_provider,
     _jumpstart_speculative_decoding,
 )
-from sagemaker.serve.utils.predictors import _get_local_mode_predictor
+from sagemaker.serve.utils.predictors import _get_local_mode_predictor, InProcessModePredictor
 from sagemaker.serve.utils.hardware_detector import (
     _get_gpu_info,
     _get_gpu_info_fallback,
@@ -566,6 +566,18 @@ def _model_builder_deploy_wrapper(
         if mode and mode != self.mode:
             self._overwrite_mode_in_deploy(overwrite_mode=mode)
 
+        if self.mode == Mode.IN_PROCESS:
+            serializer, deserializer = self._get_client_translators()
+
+            predictor = InProcessModePredictor(
+                self.modes[str(Mode.IN_PROCESS)], serializer, deserializer
+            )
+
+            self.modes[str(Mode.IN_PROCESS)].create_server(
+                predictor,
+            )
+            return predictor
+
         if self.mode == Mode.LOCAL_CONTAINER:
             serializer, deserializer = self._get_client_translators()
             predictor = _get_local_mode_predictor(
@@ -919,11 +931,16 @@ def build(  # pylint: disable=R0911
 
     def _build_validations(self):
         """Validations needed for model server overrides, or auto-detection or fallback"""
-        if self.mode == Mode.IN_PROCESS and self.model_server is not ModelServer.MMS:
+        if (
+            self.mode == Mode.IN_PROCESS
+            and self.model_server is not ModelServer.MMS
+            and self.model_server is not ModelServer.DJL_SERVING
+            and self.model_server is not ModelServer.TORCHSERVE
+        ):
             raise ValueError(
-                "IN_PROCESS mode is only supported for MMS/Transformers server in beta release."
+                "IN_PROCESS mode is only supported for the following servers "
+                "in beta release: MMS/Transformers, TORCHSERVE, DJL_SERVING server"
             )
-
         if self.inference_spec and self.model:
             raise ValueError("Can only set one of the following: model, inference_spec.")
 

src/sagemaker/serve/builder/transformers_builder.py

@@ -38,7 +38,7 @@
 from sagemaker.serve.utils.optimize_utils import _is_optimized
 from sagemaker.serve.utils.predictors import (
     TransformersLocalModePredictor,
-    TransformersInProcessModePredictor,
+    InProcessModePredictor,
 )
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.mode.function_pointers import Mode
@@ -237,7 +237,7 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
         if self.mode == Mode.IN_PROCESS:
             timeout = kwargs.get("model_data_download_timeout")
 
-            predictor = TransformersInProcessModePredictor(
+            predictor = InProcessModePredictor(
                 self.modes[str(Mode.IN_PROCESS)], serializer, deserializer
             )
 

src/sagemaker/serve/mode/in_process_mode.py

@@ -13,20 +13,15 @@
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.utils.exceptions import InProcessDeepPingException
-from sagemaker.serve.model_server.multi_model_server.server import InProcessMultiModelServer
+from sagemaker.serve.model_server.in_process_model_server.in_process_server import InProcessServing
 from sagemaker.session import Session
 
 logger = logging.getLogger(__name__)
 
-_PING_HEALTH_CHECK_FAIL_MSG = (
-    "Ping health check did not pass. "
-    + "Please increase container_timeout_seconds or review your inference code."
-)
+_PING_HEALTH_CHECK_FAIL_MSG = "Ping health check did not pass. Please review your inference code."
 
 
-class InProcessMode(
-    InProcessMultiModelServer,
-):
+class InProcessMode(InProcessServing):
     """A class that holds methods to deploy model to a container in process environment"""
 
     def __init__(
@@ -70,7 +65,13 @@ def create_server(
         logger.info("Waiting for model server %s to start up...", self.model_server)
 
         if self.model_server == ModelServer.MMS:
-            self._ping_local_server = self._multi_model_server_deep_ping
+            self._ping_local_server = self._deep_ping
+            self._start_serving()
+        elif self.model_server == ModelServer.DJL_SERVING:
+            self._ping_local_server = self._deep_ping
+            self._start_serving()
+        elif self.model_server == ModelServer.TORCHSERVE:
+            self._ping_local_server = self._deep_ping
             self._start_serving()
 
         # allow some time for server to be ready.

src/sagemaker/serve/app.py renamed to src/sagemaker/serve/model_server/in_process_model_server/app.py

@@ -3,10 +3,13 @@
 from __future__ import absolute_import
 
 import asyncio
+import io
 import logging
 import threading
 from typing import Optional
 
+from sagemaker.serve.spec.inference_spec import InferenceSpec
+from sagemaker.serve.builder.schema_builder import SchemaBuilder
 
 logger = logging.getLogger(__name__)
 
@@ -17,45 +20,44 @@
     logger.error("Unable to import uvicorn, check if uvicorn is installed.")
 
 
-try:
-    from transformers import pipeline
-except ImportError:
-    logger.error("Unable to import transformers, check if transformers is installed.")
-
-
 try:
     from fastapi import FastAPI, Request, APIRouter
 except ImportError:
     logger.error("Unable to import fastapi, check if fastapi is installed.")
 
 
 class InProcessServer:
-    """Placeholder docstring"""
+    """Generic In-Process Server for Serving Models using InferenceSpec"""
 
-    def __init__(self, model_id: Optional[str] = None, task: Optional[str] = None):
+    def __init__(
+        self,
+        inference_spec: Optional[InferenceSpec] = None,
+        schema_builder: Optional[SchemaBuilder] = None,
+    ):
         self._thread = None
         self._loop = None
         self._stop_event = asyncio.Event()
         self._router = APIRouter()
-        self._model_id = model_id
-        self._task = task
         self.server = None
         self.port = None
         self.host = None
-        # TODO: Pick up device automatically.
-        self._generator = pipeline(task, model=model_id, device="cpu")
-
-        # pylint: disable=unused-variable
-        @self._router.post("/generate")
-        async def generate_text(prompt: Request):
-            """Placeholder docstring"""
-            str_prompt = await prompt.json()
-            str_prompt = str_prompt["inputs"] if "inputs" in str_prompt else str_prompt
-
-            generated_text = self._generator(
-                str_prompt, max_length=30, num_return_sequences=1, truncation=True
+        self.inference_spec = inference_spec
+        self.schema_builder = schema_builder
+        self._load_model = self.inference_spec.load(model_dir=None)
+
+        @self._router.post("/invoke")
+        async def invoke(request: Request):
+            """Generate text based on the provided prompt"""
+
+            request_header = request.headers
+            request_body = await request.body()
+            content_type = request_header.get("Content-Type", None)
+            input_data = schema_builder.input_deserializer.deserialize(
+                io.BytesIO(request_body), content_type[0]
             )
-            return generated_text
+            logger.debug(f"Received request: {input_data}")
+            response = self.inference_spec.invoke(input_data, self._load_model)
+            return response
 
         self._create_server()
 

src/sagemaker/serve/model_server/in_process_model_server/in_process_server.py

@@ -0,0 +1,60 @@
+"""Module for In_process Serving"""
+
+from __future__ import absolute_import
+
+import requests
+import logging
+from sagemaker.serve.utils.exceptions import LocalModelInvocationException
+from sagemaker.base_predictor import PredictorBase
+
+logger = logging.getLogger(__name__)
+
+
+class InProcessServing:
+    """In Process Mode server instance"""
+
+    def _start_serving(self):
+        """Initializes the start of the server"""
+        from sagemaker.serve.model_server.in_process_model_server.app import InProcessServer
+
+        self.server = InProcessServer(
+            inference_spec=self.inference_spec, schema_builder=self.schema_builder
+        )
+        self.server.start_server()
+
+    def _stop_serving(self):
+        """Stops the server"""
+        self.server.stop_server()
+
+    def _invoke_serving(self, request: object, content_type: str, accept: str):
+        """Placeholder docstring"""
+        try:
+            response = requests.post(
+                f"http://{self.server.host}:{self.server.port}/invoke",
+                data=request,
+                headers={"Content-Type": content_type, "Accept": accept},
+                timeout=600,
+            )
+            response.raise_for_status()
+
+            return response.content
+        except Exception as e:
+            if "Connection refused" in str(e):
+                raise Exception(
+                    "Unable to send request to the local server: Connection refused."
+                ) from e
+            raise Exception("Unable to send request to the local container server %s", str(e))
+
+    def _deep_ping(self, predictor: PredictorBase):
+        """Sends a deep ping to ensure prediction"""
+        healthy = False
+        response = None
+        try:
+            response = predictor.predict(self.schema_builder.sample_input)
+            healthy = response is not None
+            # pylint: disable=broad-except
+        except Exception as e:
+            if "422 Client Error: Unprocessable Entity for url" in str(e):
+                raise LocalModelInvocationException(str(e))
+
+        return healthy, response