Add Unit Tests

Author: Jonathan Makunga

src/sagemaker/jumpstart/model.py

@@ -14,7 +14,7 @@
 
 from __future__ import absolute_import
 
-from functools import lru_cache
+from functools import cached_property
 from typing import Dict, List, Optional, Union, Any
 import pandas as pd
 from botocore.exceptions import ClientError
@@ -293,17 +293,6 @@ def __init__(
             ValueError: If the model ID is not recognized by JumpStart.
         """
 
-        metadata_configs = get_jumpstart_configs(
-            region=region,
-            model_id=model_id,
-            model_version=model_version,
-            sagemaker_session=sagemaker_session,
-        )
-        self._deployment_configs = self._deployment_configs = [
-            self._convert_to_deployment_config_metadata(config_name, config)
-            for config_name, config in metadata_configs.items()
-        ]
-
         def _validate_model_id_and_type():
             return validate_model_id_and_get_type(
                 model_id=model_id,
@@ -372,6 +361,17 @@ def _validate_model_id_and_type():
         self.model_package_arn = model_init_kwargs.model_package_arn
         self.init_kwargs = model_init_kwargs.to_kwargs_dict(False)
 
+        metadata_configs = get_jumpstart_configs(
+            region=region,
+            model_id=model_id,
+            model_version=model_version,
+            sagemaker_session=sagemaker_session,
+        )
+        self._deployment_configs = [
+            self._convert_to_deployment_config_metadata(config_name, config)
+            for config_name, config in metadata_configs.items()
+        ]
+
     def log_subscription_warning(self) -> None:
         """Log message prompting the customer to subscribe to the proprietary model."""
         subscription_link = verify_model_region_and_return_specs(
@@ -828,8 +828,7 @@ def register_deploy_wrapper(*args, **kwargs):
 
         return model_package
 
-    @lru_cache
-    @property
+    @cached_property
     def benchmark_metrics(self) -> pd.DataFrame:
         """Pandas DataFrame object of Benchmark Metrics for deployment configs"""
         data = extract_metrics_from_deployment_configs(
@@ -844,16 +843,17 @@ def display_benchmark_metrics(self, *args, **kwargs):
 
     def list_deployment_configs(self) -> List[Dict[str, Any]]:
         """List deployment configs for ``This`` model in the current region.
+
         Returns:
-            A list of deployment configs (List[Dict[str, Any]]).
+            List[Dict[str, Any]]: A list of deployment configs.
         """
         return self._deployment_configs
 
-    @lru_cache
     def _convert_to_deployment_config_metadata(
         self, config_name: str, metadata_config: JumpStartMetadataConfig
     ) -> Dict[str, Any]:
         """Retrieve deployment config for config name.
+
         Args:
             config_name (str): Name of deployment config.
             metadata_config (JumpStartMetadataConfig): Metadata config for deployment config.

src/sagemaker/jumpstart/utils.py

@@ -1005,6 +1005,7 @@ def extract_metrics_from_deployment_configs(
     deployment_configs: list[dict[str, Any]], config_name: str
 ) -> Dict[str, List[str]]:
     """Extracts metrics from deployment configs.
+
     Args:
         deployment_configs (list[dict[str, Any]]): List of deployment configs.
         config_name (str): The name of the deployment config use by the model.

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -431,14 +431,15 @@ def tune_for_tgi_jumpstart(self, max_tuning_duration: int = 1800):
             sharded_supported=sharded_supported, max_tuning_duration=max_tuning_duration
         )
 
-    def display_benchmark_metrics(self):
+    def display_benchmark_metrics(self, *args, **kwargs):
         """Display Markdown Benchmark Metrics for deployment configs."""
-        self.pysdk_model.display_benchmark_metrics()
+        self.pysdk_model.display_benchmark_metrics(args, kwargs)
 
     def list_deployment_configs(self) -> list[dict[str, Any]]:
         """List deployment configs for ``This`` model in the current region.
+
         Returns:
-            A list of deployment configs (List[Dict[str, Any]]).
+            List[Dict[str, Any]]: A list of deployment configs.
         """
         return self.pysdk_model.list_deployment_configs()
 

src/sagemaker/utils.py

@@ -1666,6 +1666,7 @@ def get_instance_rate_per_hour(
     region: str,
 ) -> Union[Dict[str, str], None]:
     """Gets instance rate per hour for the given instance type.
+
     Args:
         instance_type (str): The instance type.
         region (str): The region.