fix: Model server override logic

Author: samruds

src/sagemaker/serve/builder/model_builder.py

@@ -88,11 +88,15 @@
 
 logger = logging.getLogger(__name__)
 
+# Any new server type should be added here
 supported_model_server = {
     ModelServer.TORCHSERVE,
     ModelServer.TRITON,
     ModelServer.DJL_SERVING,
     ModelServer.TENSORFLOW_SERVING,
+    ModelServer.MMS,
+    ModelServer.TGI,
+    ModelServer.TEI,
 }
 
 
@@ -281,31 +285,6 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing,
         },
     )
 
-    def _build_validations(self):
-        """Placeholder docstring"""
-        # TODO: Beta validations - remove after the launch
-        if self.mode == Mode.IN_PROCESS:
-            raise ValueError("IN_PROCESS mode is not supported yet!")
-
-        if self.inference_spec and self.model:
-            raise ValueError("Cannot have both the Model and Inference spec in the builder")
-
-        if self.image_uri and not is_1p_image_uri(self.image_uri) and self.model_server is None:
-            raise ValueError(
-                "Model_server must be set when non-first-party image_uri is set. "
-                + "Supported model servers: %s" % supported_model_server
-            )
-
-        # Set TorchServe as default model server
-        if not self.model_server:
-            self.model_server = ModelServer.TORCHSERVE
-
-        if self.model_server not in supported_model_server:
-            raise ValueError(
-                "%s is not supported yet! Supported model servers: %s"
-                % (self.model_server, supported_model_server)
-            )
-
     def _save_model_inference_spec(self):
         """Placeholder docstring"""
         # check if path exists and create if not
@@ -748,6 +727,11 @@ def build(  # pylint: disable=R0911
             self._initialize_for_mlflow()
             _validate_input_for_mlflow(self.model_server, self.env_vars.get("MLFLOW_MODEL_FLAVOR"))
 
+        self._build_validations()
+
+        if self.model_server:
+            return self._build_for_model_server()
+
         if isinstance(self.model, str):
             model_task = None
             if self.model_metadata:
@@ -779,7 +763,37 @@ def build(  # pylint: disable=R0911
                 else:
                     return self._build_for_transformers()
 
-        self._build_validations()
+        # Set TorchServe as default model server
+        if not self.model_server:
+            self.model_server = ModelServer.TORCHSERVE
+            return self._build_for_torchserve()
+
+        raise ValueError("%s model server is not supported" % self.model_server)
+
+    def _build_validations(self):
+        """Validations needed for model server overrides, or auto-detectection or fallback"""
+        if self.mode == Mode.IN_PROCESS:
+            raise ValueError("IN_PROCESS mode is not supported yet!")
+
+        if self.inference_spec and self.model:
+            raise ValueError("Cannot have both the Model and Inference spec in the builder")
+
+        if self.image_uri and not is_1p_image_uri(self.image_uri) and self.model_server is None:
+            raise ValueError(
+                "Model_server must be set when non-first-party image_uri is set. "
+                + "Supported model servers: %s" % supported_model_server
+            )
+
+    def _build_for_model_server(self):  # pylint: disable=R0911, R1710
+        """Model server overrides"""
+        if self.model_server not in supported_model_server:
+            raise ValueError(
+                "%s is not supported yet! Supported model servers: %s"
+                % (self.model_server, supported_model_server)
+            )
+
+        if not self.model and not self.model_metadata:
+            raise ValueError("Missing required parameter `model` or 'ml_flow' path")
 
         if self.model_server == ModelServer.TORCHSERVE:
             return self._build_for_torchserve()
@@ -790,7 +804,17 @@ def build(  # pylint: disable=R0911
         if self.model_server == ModelServer.TENSORFLOW_SERVING:
             return self._build_for_tensorflow_serving()
 
-        raise ValueError("%s model server is not supported" % self.model_server)
+        if self.model_server == ModelServer.DJL_SERVING:
+            return self._build_for_djl()
+
+        if self.model_server == ModelServer.TEI:
+            return self._build_for_tei()
+
+        if self.model_server == ModelServer.TGI:
+            return self._build_for_tgi()
+
+        if self.model_server == ModelServer.MMS:
+            return self._build_for_transformers()
 
     def save(
         self,

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -54,6 +54,9 @@
     ModelServer.TRITON,
     ModelServer.DJL_SERVING,
     ModelServer.TENSORFLOW_SERVING,
+    ModelServer.MMS,
+    ModelServer.TGI,
+    ModelServer.TEI,
 }
 
 mock_session = MagicMock()
@@ -124,6 +127,120 @@ def test_save_model_throw_exception_when_none_of_model_and_inference_spec_is_set
             mock_session,
         )
 
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_djl")
+    def test_model_server_override_djl_with_model(self, mock_build_for_djl, mock_serve_settings):
+        mock_setting_object = mock_serve_settings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        builder = ModelBuilder(model_server=ModelServer.DJL_SERVING, model="gpt_llm_burt")
+        builder.build(sagemaker_session=mock_session)
+
+        mock_build_for_djl.assert_called_once()
+
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    def test_model_server_override_djl_without_model_or_mlflow(self, mock_serve_settings):
+        builder = ModelBuilder(
+            model_server=ModelServer.DJL_SERVING, model=None, inference_spec=None
+        )
+        self.assertRaisesRegex(
+            Exception,
+            "Missing required parameter `model` or 'ml_flow' path",
+            builder.build,
+            Mode.SAGEMAKER_ENDPOINT,
+            mock_role_arn,
+            mock_session,
+        )
+
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_torchserve")
+    def test_model_server_override_torchserve_with_model(
+        self, mock_build_for_ts, mock_serve_settings
+    ):
+        mock_setting_object = mock_serve_settings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        builder = ModelBuilder(model_server=ModelServer.TORCHSERVE, model="gpt_llm_burt")
+        builder.build(sagemaker_session=mock_session)
+
+        mock_build_for_ts.assert_called_once()
+
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    def test_model_server_override_torchserve_without_model_or_mlflow(self, mock_serve_settings):
+        builder = ModelBuilder(model_server=ModelServer.TORCHSERVE)
+        self.assertRaisesRegex(
+            Exception,
+            "Missing required parameter `model` or 'ml_flow' path",
+            builder.build,
+            Mode.SAGEMAKER_ENDPOINT,
+            mock_role_arn,
+            mock_session,
+        )
+
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_triton")
+    def test_model_server_override_triton_with_model(self, mock_build_for_ts, mock_serve_settings):
+        mock_setting_object = mock_serve_settings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        builder = ModelBuilder(model_server=ModelServer.TRITON, model="gpt_llm_burt")
+        builder.build(sagemaker_session=mock_session)
+
+        mock_build_for_ts.assert_called_once()
+
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_tensorflow_serving")
+    def test_model_server_override_tensor_with_model(self, mock_build_for_ts, mock_serve_settings):
+        mock_setting_object = mock_serve_settings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        builder = ModelBuilder(model_server=ModelServer.TENSORFLOW_SERVING, model="gpt_llm_burt")
+        builder.build(sagemaker_session=mock_session)
+
+        mock_build_for_ts.assert_called_once()
+
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_tei")
+    def test_model_server_override_tei_with_model(self, mock_build_for_ts, mock_serve_settings):
+        mock_setting_object = mock_serve_settings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        builder = ModelBuilder(model_server=ModelServer.TEI, model="gpt_llm_burt")
+        builder.build(sagemaker_session=mock_session)
+
+        mock_build_for_ts.assert_called_once()
+
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_tgi")
+    def test_model_server_override_tgi_with_model(self, mock_build_for_ts, mock_serve_settings):
+        mock_setting_object = mock_serve_settings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        builder = ModelBuilder(model_server=ModelServer.TGI, model="gpt_llm_burt")
+        builder.build(sagemaker_session=mock_session)
+
+        mock_build_for_ts.assert_called_once()
+
+    @patch("sagemaker.serve.builder.model_builder._ServeSettings")
+    @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_transformers")
+    def test_model_server_override_transformers_with_model(
+        self, mock_build_for_ts, mock_serve_settings
+    ):
+        mock_setting_object = mock_serve_settings.return_value
+        mock_setting_object.role_arn = mock_role_arn
+        mock_setting_object.s3_model_data_url = mock_s3_model_data_url
+
+        builder = ModelBuilder(model_server=ModelServer.MMS, model="gpt_llm_burt")
+        builder.build(sagemaker_session=mock_session)
+
+        mock_build_for_ts.assert_called_once()
+
     @patch("os.makedirs", Mock())
     @patch("sagemaker.serve.builder.model_builder._detect_framework_and_version")
     @patch("sagemaker.serve.builder.model_builder.prepare_for_torchserve")