Test backward compatibility

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -440,7 +440,7 @@ def set_deployment_config(self, config_name: Optional[str]) -> None:
                 any existing config that is applied to the model.
         """
         if not hasattr(self, "pysdk_model") or self.pysdk_model is None:
-            raise Exception("Cannot set deployment config to an uninitialized model")
+            raise Exception("Cannot set deployment config to an uninitialized model.")
 
         self.pysdk_model.set_deployment_config(config_name)
 

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

@@ -718,6 +718,80 @@ def test_get_deployment_config(
 
         self.assertEqual(builder.get_deployment_config(), expected)
 
+    @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_set_deployment_config(
+        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/galactica-mock-model-id",
+            schema_builder=mock_schema_builder,
+        )
+
+        mock_pre_trained_model.return_value.image_uri = mock_tgi_image_uri
+
+        builder.build()
+        builder.set_deployment_config("config-1")
+
+        mock_pre_trained_model.return_value.set_deployment_config.assert_called_with("config-1")
+
+    @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_set_deployment_config_ex(
+        self,
+        mock_get_nb_instance,
+        mock_get_ram_usage_mb,
+        mock_prepare_for_tgi,
+        mock_pre_trained_model,
+        mock_is_jumpstart_model,
+        mock_telemetry,
+    ):
+        mock_pre_trained_model.return_value.image_uri = mock_tgi_image_uri
+
+        self.assertRaisesRegex(
+            Exception,
+            "Cannot set deployment config to an uninitialized model.",
+            lambda: ModelBuilder(
+                model="facebook/galactica-mock-model-id", schema_builder=mock_schema_builder
+            ).set_deployment_config("config-2"),
+        )
+
     @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
     @patch(
         "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
@@ -759,11 +833,3 @@ def test_display_benchmark_metrics(
         builder.display_benchmark_metrics()
 
         mock_pre_trained_model.return_value.display_benchmark_metrics.assert_called_once()
-
-    def test_display_benchmark_metrics_ex(self):
-        self.assertRaises(
-            Exception,
-            lambda: ModelBuilder(
-                model="facebook/galactica-mock-model-id", schema_builder=mock_schema_builder
-            ).display_benchmark_metrics(),
-        )