add modelID support to model builder InProcess model (#1580)

* add modelID support to model builder InProcess model

* fix format

* import pkg only if model id is specified

* fix format and remove unused imports

* add inference_spec to triton inProcess mode

* fix format

* fix format using black

* fix app.py

* Implement stop_server for fast api, remove model_server param from in_process mode

* fix tests and format

* rebase

* remove Inference spec support arg in DJL builder, add fastapi reqs to pyproject

* make changes to not support JS in in_process mode

* fix code style errors

* fix flake8 warnings

* fix tgi unit tests to build for TGI usecase instead of JS

Author: pravali96

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
   "boto3>=1.34.142,<2.0",
   "cloudpickle==2.2.1",
   "docker",
+  "fastapi",
   "google-pasta",
   "importlib-metadata>=1.4.0,<7.0",
   "jsonschema",
@@ -54,6 +55,7 @@ dependencies = [
   "tblib>=1.7.0,<4",
   "tqdm",
   "urllib3>=1.26.8,<3.0.0",
+  "uvicorn"
 ]
 
 [project.scripts]

src/sagemaker/serve/builder/djl_builder.py

@@ -78,7 +78,6 @@ 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
@@ -263,12 +262,7 @@ def _build_for_hf_djl(self):
 
         _create_dir_structure(self.model_path)
         if not hasattr(self, "pysdk_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.env_vars.update({"HF_MODEL_ID": self.model})
 
             self.hf_model_config = _get_model_config_properties_from_hf(
                 self.env_vars.get("HF_MODEL_ID"), self.env_vars.get("HF_TOKEN")

src/sagemaker/serve/builder/model_builder.py

@@ -87,7 +87,10 @@
     _extract_speculative_draft_model_provider,
     _jumpstart_speculative_decoding,
 )
-from sagemaker.serve.utils.predictors import _get_local_mode_predictor, InProcessModePredictor
+from sagemaker.serve.utils.predictors import (
+    _get_local_mode_predictor,
+    _get_in_process_mode_predictor,
+)
 from sagemaker.serve.utils.hardware_detector import (
     _get_gpu_info,
     _get_gpu_info_fallback,
@@ -435,11 +438,11 @@ def _prepare_for_mode(
             # init the InProcessMode object
             self.modes[str(Mode.IN_PROCESS)] = InProcessMode(
                 inference_spec=self.inference_spec,
+                model=self.model,
                 schema_builder=self.schema_builder,
                 session=self.sagemaker_session,
                 model_path=self.model_path,
                 env_vars=self.env_vars,
-                model_server=self.model_server,
             )
             self.modes[str(Mode.IN_PROCESS)].prepare()
             return None
@@ -575,7 +578,7 @@ def _model_builder_deploy_wrapper(
         if self.mode == Mode.IN_PROCESS:
             serializer, deserializer = self._get_client_translators()
 
-            predictor = InProcessModePredictor(
+            predictor = _get_in_process_mode_predictor(
                 self.modes[str(Mode.IN_PROCESS)], serializer, deserializer
             )
 
@@ -597,6 +600,7 @@ def _model_builder_deploy_wrapper(
                 self.image_uri, container_timeout_in_second, self.secret_key, predictor
             )
             return predictor
+
         if self.mode == Mode.SAGEMAKER_ENDPOINT:
             # Validate parameters
             # Instance type and instance count parameter validation is done based on deployment type
@@ -650,16 +654,17 @@ def _build_for_torchserve(self) -> Type[Model]:
         """Build the model for torchserve"""
         self._save_model_inference_spec()
 
-        self._auto_detect_container()
+        if self.mode != Mode.IN_PROCESS:
+            self._auto_detect_container()
 
-        self.secret_key = prepare_for_torchserve(
-            model_path=self.model_path,
-            shared_libs=self.shared_libs,
-            dependencies=self.dependencies,
-            session=self.sagemaker_session,
-            image_uri=self.image_uri,
-            inference_spec=self.inference_spec,
-        )
+            self.secret_key = prepare_for_torchserve(
+                model_path=self.model_path,
+                shared_libs=self.shared_libs,
+                dependencies=self.dependencies,
+                session=self.sagemaker_session,
+                image_uri=self.image_uri,
+                inference_spec=self.inference_spec,
+            )
 
         self._prepare_for_mode()
         self.model = self._create_model()
@@ -854,6 +859,7 @@ def build(  # pylint: disable=R0911
         Returns:
             Type[Model]: A deployable ``Model`` object.
         """
+        from sagemaker.modules.train.model_trainer import ModelTrainer
 
         self.modes = dict()
 
@@ -908,10 +914,25 @@ def build(  # pylint: disable=R0911
 
         if isinstance(self.model, str):
             model_task = None
-            if self._is_jumpstart_model_id() or self._use_jumpstart_equivalent():
+
+            if self._is_jumpstart_model_id():
+                if self.mode == Mode.IN_PROCESS:
+                    raise ValueError(
+                        f"{self.mode} is not supported for Jumpstart models. "
+                        "Please use LOCAL_CONTAINER mode to deploy a Jumpstart model"
+                        " on your local machine."
+                    )
                 self.model_hub = ModelHub.JUMPSTART
+                logger.debug("Building for Jumpstart model Id...")
                 self.built_model = self._build_for_jumpstart()
                 return self.built_model
+
+            if self.mode != Mode.IN_PROCESS:
+                if self._use_jumpstart_equivalent():
+                    self.model_hub = ModelHub.JUMPSTART
+                    logger.debug("Building for Jumpstart equiavalent model Id...")
+                    self.built_model = self._build_for_jumpstart()
+                    return self.built_model
             self.model_hub = ModelHub.HUGGINGFACE
 
             if self.model_metadata:
@@ -931,7 +952,7 @@ def build(  # pylint: disable=R0911
                 if model_task == "text-generation":
                     self.built_model = self._build_for_tgi()
                     return self.built_model
-                if model_task == "sentence-similarity":
+                if model_task in ["sentence-similarity", "feature-extraction"]:
                     self.built_model = self._build_for_tei()
                     return self.built_model
                 elif self._can_fit_on_single_gpu():
@@ -951,16 +972,6 @@ 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
-            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 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/tei_builder.py

@@ -26,13 +26,14 @@
 )
 from sagemaker.serve.model_server.tgi.prepare import _create_dir_structure
 from sagemaker.serve.utils.optimize_utils import _is_optimized
-from sagemaker.serve.utils.predictors import TeiLocalModePredictor
+from sagemaker.serve.utils.predictors import InProcessModePredictor, TeiLocalModePredictor
 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.base_predictor import PredictorBase
 
 logger = logging.getLogger(__name__)
+LOCAL_MODES = [Mode.LOCAL_CONTAINER, Mode.IN_PROCESS]
 
 _CODE_FOLDER = "code"
 
@@ -141,6 +142,17 @@ def _tei_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:
+            self._prepare_for_mode()
+            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")
 
@@ -222,7 +234,7 @@ def _build_for_hf_tei(self):
 
         self.pysdk_model = self._create_tei_model()
 
-        if self.mode == Mode.LOCAL_CONTAINER:
+        if self.mode in LOCAL_MODES:
             self._prepare_for_mode()
 
         return self.pysdk_model

src/sagemaker/serve/builder/tgi_builder.py

@@ -49,13 +49,14 @@
     _get_gpu_info_fallback,
 )
 from sagemaker.serve.model_server.tgi.prepare import _create_dir_structure
-from sagemaker.serve.utils.predictors import TgiLocalModePredictor
+from sagemaker.serve.utils.predictors import TgiLocalModePredictor, InProcessModePredictor
 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.base_predictor import PredictorBase
 
 logger = logging.getLogger(__name__)
+LOCAL_MODES = [Mode.LOCAL_CONTAINER, Mode.IN_PROCESS]
 
 _CODE_FOLDER = "code"
 _INVALID_SAMPLE_DATA_EX = (
@@ -176,6 +177,17 @@ def _tgi_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")
 
@@ -280,7 +292,7 @@ def _build_for_hf_tgi(self):
             ] = _default_max_new_tokens
         self.pysdk_model = self._create_tgi_model()
 
-        if self.mode == Mode.LOCAL_CONTAINER:
+        if self.mode in LOCAL_MODES:
             self._prepare_for_mode()
 
         return self.pysdk_model

src/sagemaker/serve/mode/in_process_mode.py

@@ -4,14 +4,13 @@
 
 from pathlib import Path
 import logging
-from typing import Dict, Type
+from typing import Dict, Type, Optional
 import time
 from datetime import datetime, timedelta
 
 from sagemaker.base_predictor import PredictorBase
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 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.in_process_model_server.in_process_server import InProcessServing
 from sagemaker.session import Session
@@ -26,8 +25,8 @@ class InProcessMode(InProcessServing):
 
     def __init__(
         self,
-        model_server: ModelServer,
-        inference_spec: Type[InferenceSpec],
+        model: Optional[str],
+        inference_spec: Optional[InferenceSpec],
         schema_builder: Type[SchemaBuilder],
         session: Session,
         model_path: str = None,
@@ -36,12 +35,12 @@ def __init__(
         # pylint: disable=bad-super-call
         super().__init__()
 
+        self.model = model
         self.inference_spec = inference_spec
         self.model_path = model_path
         self.env_vars = env_vars
         self.session = session
         self.schema_builder = schema_builder
-        self.model_server = model_server
         self._ping_local_server = None
 
     def load(self, model_path: str = None):
@@ -61,18 +60,15 @@ def create_server(
         self,
         predictor: PredictorBase,
     ):
-        """Creating the server and checking ping health."""
-        logger.info("Waiting for model server %s to start up...", self.model_server)
-
-        if self.model_server == ModelServer.MMS:
-            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()
+        """Creating the fast api server and checking ping health."""
+
+        logger.info("Waiting for fastapi server to start up...")
+
+        logger.warning("Note: This is not a standard model server.")
+        logger.warning("The model is being hosted directly on the FastAPI server.")
+
+        self._ping_local_server = self._deep_ping
+        self._start_serving()
 
         # allow some time for server to be ready.
         time.sleep(1)