fix: model builder limited container support for endpoint mode. (#4683)

makungaj1 · Jonathan Makunga · web-flow · commit 8002d7f92c42 · 2024-05-15T19:18:16.000-07:00
* Allow ModelBuilder's endpoint mode for Jumpstart models packaged with containers other than TGI and DJL

* increase coverage

* Add JS Support for MMS Serving

* Add JS Support for MMS Serving

* Unit tests

* Refactoring

* Refactoring

* Refactoring

---------

Co-authored-by: Jonathan Makunga &lt;makung@amazon.com&gt;
diff --git a/src/sagemaker/serve/builder/jumpstart_builder.py b/src/sagemaker/serve/builder/jumpstart_builder.py
@@ -23,6 +23,7 @@
 from sagemaker import model_uris
 from sagemaker.serve.model_server.djl_serving.prepare import prepare_djl_js_resources
 from sagemaker.serve.model_server.djl_serving.utils import _get_admissible_tensor_parallel_degrees
+from sagemaker.serve.model_server.multi_model_server.prepare import prepare_mms_js_resources
 from sagemaker.serve.model_server.tgi.prepare import prepare_tgi_js_resources, _create_dir_structure
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.exceptions import (
@@ -35,6 +36,7 @@
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
     TgiLocalModePredictor,
+    TransformersLocalModePredictor,
 )
 from sagemaker.serve.utils.local_hardware import (
     _get_nb_instance,
@@ -90,6 +92,7 @@ def __init__(self):
         self.existing_properties = None
         self.prepared_for_tgi = None
         self.prepared_for_djl = None
+        self.prepared_for_mms = None
         self.schema_builder = None
         self.nb_instance_type = None
         self.ram_usage_model_load = None
@@ -137,7 +140,11 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
 
             if overwrite_mode == Mode.SAGEMAKER_ENDPOINT:
                 self.mode = self.pysdk_model.mode = Mode.SAGEMAKER_ENDPOINT
-                if not hasattr(self, "prepared_for_djl") or not hasattr(self, "prepared_for_tgi"):
+                if (
+                    not hasattr(self, "prepared_for_djl")
+                    or not hasattr(self, "prepared_for_tgi")
+                    or not hasattr(self, "prepared_for_mms")
+                ):
                     self.pysdk_model.model_data, env = self._prepare_for_mode()
             elif overwrite_mode == Mode.LOCAL_CONTAINER:
                 self.mode = self.pysdk_model.mode = Mode.LOCAL_CONTAINER
@@ -160,6 +167,13 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                         dependencies=self.dependencies,
                         model_data=self.pysdk_model.model_data,
                     )
+                elif not hasattr(self, "prepared_for_mms"):
+                    self.js_model_config, self.prepared_for_mms = prepare_mms_js_resources(
+                        model_path=self.model_path,
+                        js_id=self.model,
+                        dependencies=self.dependencies,
+                        model_data=self.pysdk_model.model_data,
+                    )
 
                 self._prepare_for_mode()
                 env = {}
@@ -179,6 +193,10 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                 predictor = TgiLocalModePredictor(
                     self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
                 )
+            elif self.model_server == ModelServer.MMS:
+                predictor = TransformersLocalModePredictor(
+                    self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
+                )
 
             ram_usage_before = _get_ram_usage_mb()
             self.modes[str(Mode.LOCAL_CONTAINER)].create_server(
@@ -254,6 +272,24 @@ def _build_for_tgi_jumpstart(self):
 
         self.pysdk_model.env.update(env)
 
+    def _build_for_mms_jumpstart(self):
+        """Placeholder docstring"""
+
+        env = {}
+        if self.mode == Mode.LOCAL_CONTAINER:
+            if not hasattr(self, "prepared_for_mms"):
+                self.js_model_config, self.prepared_for_mms = prepare_mms_js_resources(
+                    model_path=self.model_path,
+                    js_id=self.model,
+                    dependencies=self.dependencies,
+                    model_data=self.pysdk_model.model_data,
+                )
+            self._prepare_for_mode()
+        elif self.mode == Mode.SAGEMAKER_ENDPOINT and hasattr(self, "prepared_for_mms"):
+            self.pysdk_model.model_data, env = self._prepare_for_mode()
+
+        self.pysdk_model.env.update(env)
+
     def _tune_for_js(self, sharded_supported: bool, max_tuning_duration: int = 1800):
         """Tune for Jumpstart Models in Local Mode.
 
@@ -264,11 +300,6 @@ def _tune_for_js(self, sharded_supported: bool, max_tuning_duration: int = 1800)
         returns:
             Tuned Model.
         """
-        if self.mode != Mode.LOCAL_CONTAINER:
-            logger.warning(
-                "Tuning is only a %s capability. Returning original model.", Mode.LOCAL_CONTAINER
-            )
-            return self.pysdk_model
 
         num_shard_env_var_name = "SM_NUM_GPUS"
         if "OPTION_TENSOR_PARALLEL_DEGREE" in self.pysdk_model.env.keys():
@@ -437,42 +468,58 @@ def _build_for_jumpstart(self):
         self.secret_key = None
         self.jumpstart = True
 
-        pysdk_model = self._create_pre_trained_js_model()
+        self.pysdk_model = self._create_pre_trained_js_model()
+        self.pysdk_model.tune = lambda *args, **kwargs: self._default_tune()
 
-        image_uri = pysdk_model.image_uri
+        logger.info(
+            "JumpStart ID %s is packaged with Image URI: %s", self.model, self.pysdk_model.image_uri
+        )
 
-        logger.info("JumpStart ID %s is packaged with Image URI: %s", self.model, image_uri)
+        if self.mode != Mode.SAGEMAKER_ENDPOINT:
+            if self._is_gated_model(self.pysdk_model):
+                raise ValueError(
+                    "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode."
+                )
 
-        if self._is_gated_model(pysdk_model) and self.mode != Mode.SAGEMAKER_ENDPOINT:
-            raise ValueError(
-                "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode."
-            )
+            if "djl-inference" in self.pysdk_model.image_uri:
+                logger.info("Building for DJL JumpStart Model ID...")
+                self.model_server = ModelServer.DJL_SERVING
+                self.image_uri = self.pysdk_model.image_uri
 
-        if "djl-inference" in image_uri:
-            logger.info("Building for DJL JumpStart Model ID...")
-            self.model_server = ModelServer.DJL_SERVING
+                self._build_for_djl_jumpstart()
 
-            self.pysdk_model = pysdk_model
-            self.image_uri = self.pysdk_model.image_uri
+                self.pysdk_model.tune = self.tune_for_djl_jumpstart
+            elif "tgi-inference" in self.pysdk_model.image_uri:
+                logger.info("Building for TGI JumpStart Model ID...")
+                self.model_server = ModelServer.TGI
+                self.image_uri = self.pysdk_model.image_uri
 
-            self._build_for_djl_jumpstart()
+                self._build_for_tgi_jumpstart()
 
-            self.pysdk_model.tune = self.tune_for_djl_jumpstart
-        elif "tgi-inference" in image_uri:
-            logger.info("Building for TGI JumpStart Model ID...")
-            self.model_server = ModelServer.TGI
+                self.pysdk_model.tune = self.tune_for_tgi_jumpstart
+            elif "huggingface-pytorch-inference:" in self.pysdk_model.image_uri:
+                logger.info("Building for MMS JumpStart Model ID...")
+                self.model_server = ModelServer.MMS
+                self.image_uri = self.pysdk_model.image_uri
 
-            self.pysdk_model = pysdk_model
-            self.image_uri = self.pysdk_model.image_uri
+                self._build_for_mms_jumpstart()
+            else:
+                raise ValueError(
+                    "JumpStart Model ID was not packaged "
+                    "with djl-inference, tgi-inference, or mms-inference container."
+                )
 
-            self._build_for_tgi_jumpstart()
+        return self.pysdk_model
 
-            self.pysdk_model.tune = self.tune_for_tgi_jumpstart
-        else:
-            raise ValueError(
-                "JumpStart Model ID was not packaged with djl-inference or tgi-inference container."
-            )
+    def _default_tune(self):
+        """Logs a warning message if tune is invoked on endpoint mode.
 
+        Returns:
+            Jumpstart Model: ``This`` model
+        """
+        logger.warning(
+            "Tuning is only a %s capability. Returning original model.", Mode.LOCAL_CONTAINER
+        )
         return self.pysdk_model
 
     def _is_gated_model(self, model) -> bool:
diff --git a/src/sagemaker/serve/model_server/multi_model_server/prepare.py b/src/sagemaker/serve/model_server/multi_model_server/prepare.py
@@ -15,7 +15,9 @@
 from __future__ import absolute_import
 import logging
 from pathlib import Path
+from typing import List
 
+from sagemaker.serve.model_server.tgi.prepare import _copy_jumpstart_artifacts
 from sagemaker.serve.utils.local_hardware import _check_disk_space, _check_docker_disk_usage
 
 logger = logging.getLogger(__name__)
@@ -36,3 +38,28 @@ def _create_dir_structure(model_path: str) -> tuple:
     _check_docker_disk_usage()
 
     return model_path, code_dir
+
+
+def prepare_mms_js_resources(
+    model_path: str,
+    js_id: str,
+    shared_libs: List[str] = None,
+    dependencies: str = None,
+    model_data: str = None,
+) -> tuple:
+    """Prepare serving when a JumpStart model id is given
+
+    Args:
+        model_path (str) : Argument
+        js_id (str): Argument
+        shared_libs (List[]) : Argument
+        dependencies (str) : Argument
+        model_data (str) : Argument
+
+    Returns:
+        ( str ) :
+
+    """
+    model_path, code_dir = _create_dir_structure(model_path)
+
+    return _copy_jumpstart_artifacts(model_data, js_id, code_dir)
diff --git a/tests/unit/sagemaker/serve/builder/test_js_builder.py b/tests/unit/sagemaker/serve/builder/test_js_builder.py
@@ -63,6 +63,10 @@
     "123456789712.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi"
     "-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu20.04"
 )
+mock_invalid_image_uri = (
+    "123456789712.dkr.ecr.us-west-2.amazonaws.com/invalid"
+    "-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu20.04"
+)
 mock_djl_image_uri = (
     "123456789712.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.24.0-neuronx-sdk2.14.1"
 )
@@ -82,6 +86,88 @@
 
 
 class TestJumpStartBuilder(unittest.TestCase):
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
+        return_value=True,
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._create_pre_trained_js_model",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.prepare_tgi_js_resources",
+        return_value=({"model_type": "t5", "n_head": 71}, True),
+    )
+    @patch("sagemaker.serve.builder.jumpstart_builder._get_ram_usage_mb", return_value=1024)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder._get_nb_instance", return_value="ml.g5.24xlarge"
+    )
+    def test__build_for_jumpstart_value_error(
+        self,
+        mock_get_nb_instance,
+        mock_get_ram_usage_mb,
+        mock_prepare_for_tgi,
+        mock_pre_trained_model,
+        mock_is_jumpstart_model,
+        mock_telemetry,
+    ):
+        builder = ModelBuilder(
+            model="facebook/invalid",
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+        )
+
+        mock_pre_trained_model.return_value.image_uri = mock_invalid_image_uri
+
+        self.assertRaises(
+            ValueError,
+            lambda: builder.build(),
+        )
+
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
+        return_value=True,
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._create_pre_trained_js_model",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.prepare_mms_js_resources",
+        return_value=({"model_type": "t5", "n_head": 71}, True),
+    )
+    @patch("sagemaker.serve.builder.jumpstart_builder._get_ram_usage_mb", return_value=1024)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder._get_nb_instance", return_value="ml.g5.24xlarge"
+    )
+    def test__build_for_mms_jumpstart(
+        self,
+        mock_get_nb_instance,
+        mock_get_ram_usage_mb,
+        mock_prepare_for_mms,
+        mock_pre_trained_model,
+        mock_is_jumpstart_model,
+        mock_telemetry,
+    ):
+        builder = ModelBuilder(
+            model="facebook/galactica-mock-model-id",
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+        )
+
+        mock_pre_trained_model.return_value.image_uri = (
+            "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface"
+            "-pytorch-inference:2.1.0-transformers4.37.0-gpu-py310-cu118"
+            "-ubuntu20.04"
+        )
+
+        builder.build()
+        builder.serve_settings.telemetry_opt_out = True
+
+        mock_prepare_for_mms.assert_called()
+
     @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
     @patch(
         "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
