Testing with Notebook

Author: Jonathan Makunga

src/sagemaker/jumpstart/model.py

@@ -443,10 +443,10 @@ def set_deployment_config(self, config_name: Optional[str]) -> None:
 
     @property
     def deployment_config(self) -> Optional[Dict[str, Any]]:
-        """The deployment config to apply to the model.
+        """The deployment config that will be applied to the model.
 
         Returns:
-            Union[Dict[str, Any], None]: Deployment config to apply to this model.
+            Union[Dict[str, Any], None]: Deployment config that will be applied to the model.
         """
         return self._retrieve_selected_deployment_config(self.config_name)
 
@@ -861,7 +861,7 @@ def register_deploy_wrapper(*args, **kwargs):
 
     @lru_cache
     def _get_benchmarks_data(self, config_name: str) -> Dict[str, List[str]]:
-        """Constructs deployment configs benchmark data.
+        """Deployment configs benchmark metrics.
 
         Args:
             config_name (str): The name of the selected deployment config.
@@ -878,9 +878,9 @@ def _retrieve_selected_deployment_config(self, config_name: str) -> Optional[Dic
         """Retrieve the deployment config to apply to the model.
 
         Args:
-            config_name (str): The name of the selected deployment config.
+            config_name (str): The name of the deployment config to retrieve.
         Returns:
-            Union[Dict[str, Any], None]: The deployment config to apply to the model.
+            Optional[Dict[str, Any]]: The retrieved deployment config.
         """
         if config_name is None:
             return None

src/sagemaker/jumpstart/types.py

@@ -2286,7 +2286,6 @@ class DeploymentConfigMetadata(BaseDeploymentConfigDataHolder):
     """Dataclass representing a Deployment Config Metadata"""
 
     __slots__ = [
-        "version",
         "deployment_config_name",
         "deployment_args",
         "acceleration_configs",
@@ -2301,7 +2300,6 @@ def __init__(
         deploy_kwargs: JumpStartModelDeployKwargs,
     ):
         """Instantiates DeploymentConfigMetadata object."""
-        self.version = "1.0.0"
         self.deployment_config_name = config_name
         self.deployment_args = DeploymentArgs(init_kwargs, deploy_kwargs)
         self.acceleration_configs = None

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -439,6 +439,9 @@ def set_deployment_config(self, config_name: Optional[str]) -> None:
                 The name of the deployment config. Set to None to unset
                 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")
+
         self.pysdk_model.set_deployment_config(config_name)
 
     def get_deployment_config(self) -> Optional[Dict[str, Any]]:
@@ -447,56 +450,58 @@ def get_deployment_config(self) -> Optional[Dict[str, Any]]:
         Returns:
             Union[Dict[str, Any], None]: Deployment config to apply to this model.
         """
-        return getattr(self, "pysdk_model", self._create_pre_trained_js_model()).deployment_config
+        if not hasattr(self, "pysdk_model") or self.pysdk_model is None:
+            self._build_for_jumpstart()
+
+        return self.pysdk_model.deployment_config
 
     def display_benchmark_metrics(self):
         """Display Markdown Benchmark Metrics for deployment configs."""
-        getattr(
-            self, "pysdk_model", self._create_pre_trained_js_model()
-        ).display_benchmark_metrics()
+        if not hasattr(self, "pysdk_model") or self.pysdk_model is None:
+            self._build_for_jumpstart()
+
+        self.pysdk_model.display_benchmark_metrics()
 
     def list_deployment_configs(self) -> List[Dict[str, Any]]:
         """List deployment configs for ``This`` model in the current region.
 
         Returns:
             List[Dict[str, Any]]: A list of deployment configs.
         """
-        return getattr(
-            self, "pysdk_model", self._create_pre_trained_js_model()
-        ).list_deployment_configs()
+        if not hasattr(self, "pysdk_model") or self.pysdk_model is None:
+            self._build_for_jumpstart()
+
+        return self.pysdk_model.list_deployment_configs()
 
     def _build_for_jumpstart(self):
         """Placeholder docstring"""
         # we do not pickle for jumpstart. set to none
         self.secret_key = None
         self.jumpstart = True
 
-        pysdk_model = getattr(self, "pysdk_model", self._create_pre_trained_js_model())
-
-        image_uri = pysdk_model.image_uri
+        if not hasattr(self, "pysdk_model") or self.pysdk_model is None:
+            self.pysdk_model = self._create_pre_trained_js_model()
 
-        logger.info("JumpStart ID %s is packaged with Image URI: %s", self.model, image_uri)
+        logger.info(
+            "JumpStart ID %s is packaged with Image URI: %s", self.model, self.pysdk_model.image_uri
+        )
 
-        if self._is_gated_model(pysdk_model) and self.mode != Mode.SAGEMAKER_ENDPOINT:
+        if self._is_gated_model() and self.mode != Mode.SAGEMAKER_ENDPOINT:
             raise ValueError(
                 "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode."
             )
 
-        if "djl-inference" in image_uri:
+        if "djl-inference" in self.pysdk_model.image_uri:
             logger.info("Building for DJL JumpStart Model ID...")
             self.model_server = ModelServer.DJL_SERVING
-
-            self.pysdk_model = pysdk_model
             self.image_uri = self.pysdk_model.image_uri
 
             self._build_for_djl_jumpstart()
 
             self.pysdk_model.tune = self.tune_for_djl_jumpstart
-        elif "tgi-inference" in image_uri:
+        elif "tgi-inference" in self.pysdk_model.image_uri:
             logger.info("Building for TGI JumpStart Model ID...")
             self.model_server = ModelServer.TGI
-
-            self.pysdk_model = pysdk_model
             self.image_uri = self.pysdk_model.image_uri
 
             self._build_for_tgi_jumpstart()
@@ -509,15 +514,13 @@ def _build_for_jumpstart(self):
 
         return self.pysdk_model
 
-    def _is_gated_model(self, model) -> bool:
+    def _is_gated_model(self) -> bool:
         """Determine if ``this`` Model is Gated
 
-        Args:
-            model (Model): Jumpstart Model
         Returns:
             bool: ``True`` if ``this`` Model is Gated
         """
-        s3_uri = model.model_data
+        s3_uri = self.pysdk_model.model_data
         if isinstance(s3_uri, dict):
             s3_uri = s3_uri.get("S3DataSource").get("S3Uri")
 

tests/unit/sagemaker/jumpstart/constants.py

@@ -7911,7 +7911,6 @@
 
 DEPLOYMENT_CONFIGS = [
     {
-        "Version": "1.0.0",
         "DeploymentConfigName": "neuron-inference",
         "DeploymentArgs": {
             "ImageUri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4"
@@ -7944,7 +7943,6 @@
         "BenchmarkMetrics": [{"name": "Instance Rate", "value": "0.0083000000", "unit": "USD/Hrs"}],
     },
     {
-        "Version": "1.0.0",
         "DeploymentConfigName": "neuron-inference-budget",
         "DeploymentArgs": {
             "ImageUri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4"
@@ -7977,7 +7975,6 @@
         "BenchmarkMetrics": [{"name": "Instance Rate", "value": "0.0083000000", "unit": "USD/Hrs"}],
     },
     {
-        "Version": "1.0.0",
         "DeploymentConfigName": "gpu-inference-budget",
         "DeploymentArgs": {
             "ImageUri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4"
@@ -8010,7 +8007,6 @@
         "BenchmarkMetrics": [{"name": "Instance Rate", "value": "0.0083000000", "unit": "USD/Hrs"}],
     },
     {
-        "Version": "1.0.0",
         "DeploymentConfigName": "gpu-inference",
         "DeploymentArgs": {
             "ImageUri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4"

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

@@ -735,7 +735,7 @@ def test_get_deployment_config(
     @patch(
         "sagemaker.serve.builder.jumpstart_builder._get_nb_instance", return_value="ml.g5.24xlarge"
     )
-    def test_set_deployment_config(
+    def test_display_benchmark_metrics(
         self,
         mock_get_nb_instance,
         mock_get_ram_usage_mb,
@@ -750,43 +750,20 @@ def test_set_deployment_config(
         )
 
         mock_pre_trained_model.return_value.image_uri = mock_tgi_image_uri
-
-        builder.build()
-        builder.serve_settings.telemetry_opt_out = True
-
-        builder.set_deployment_config("config_name")
-
-        mock_pre_trained_model.return_value.set_deployment_config.assert_called_once_with(
-            "config_name"
+        mock_pre_trained_model.return_value.list_deployment_configs.side_effect = (
+            lambda: DEPLOYMENT_CONFIGS
         )
 
-    @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")
-    @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_display_benchmark_metrics(
-        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,
-        )
+        builder.list_deployment_configs()
+
         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(),
+        )