Benchmark feature fixes

src/sagemaker/jumpstart/model.py

@@ -14,7 +14,6 @@
 
 from __future__ import absolute_import
 
-from functools import lru_cache
 from typing import Dict, List, Optional, Any, Union
 import pandas as pd
 from botocore.exceptions import ClientError
@@ -48,6 +47,8 @@
     get_jumpstart_configs,
     get_metrics_from_deployment_configs,
     add_instance_rate_stats_to_benchmark_metrics,
+    deployment_config_response_data,
+    _deployment_config_lru_cache,
 )
 from sagemaker.jumpstart.constants import JUMPSTART_LOGGER
 from sagemaker.jumpstart.enums import JumpStartModelType
@@ -449,10 +450,12 @@ def deployment_config(self) -> Optional[Dict[str, Any]]:
         Returns:
             Optional[Dict[str, Any]]: Deployment config.
         """
-        deployment_config = self._retrieve_selected_deployment_config(
-            self.config_name, self.instance_type
-        )
-        return deployment_config.to_json() if deployment_config is not None else None
+        if self.config_name is None:
+            return None
+        for config in self.list_deployment_configs():
+            if config.get("DeploymentConfigName") == self.config_name:
+                return config
+        return None
 
     @property
     def benchmark_metrics(self) -> pd.DataFrame:
@@ -461,29 +464,24 @@ def benchmark_metrics(self) -> pd.DataFrame:
         Returns:
             Benchmark Metrics: Pandas DataFrame object.
         """
-        benchmark_metrics_data = self._get_deployment_configs_benchmarks_data(
-            self.config_name, self.instance_type
-        )
-        keys = list(benchmark_metrics_data.keys())
-        df = pd.DataFrame(benchmark_metrics_data).sort_values(by=[keys[0], keys[1]])
-        return df
+        df = pd.DataFrame(self._get_deployment_configs_benchmarks_data())
+        default_mask = df.apply(lambda row: any("Default" in str(val) for val in row), axis=1)
+        sorted_df = pd.concat([df[default_mask], df[~default_mask]])
+        return sorted_df
 
-    def display_benchmark_metrics(self) -> None:
+    def display_benchmark_metrics(self, *args, **kwargs) -> None:
         """Display deployment configs benchmark metrics."""
-        print(self.benchmark_metrics.to_markdown(index=False))
+        print(self.benchmark_metrics.to_markdown(index=False), *args, **kwargs)
 
     def list_deployment_configs(self) -> List[Dict[str, Any]]:
         """List deployment configs for ``This`` model.
 
         Returns:
             List[Dict[str, Any]]: A list of deployment configs.
         """
-        return [
-            deployment_config.to_json()
-            for deployment_config in self._get_deployment_configs(
-                self.config_name, self.instance_type
-            )
-        ]
+        return deployment_config_response_data(
+            self._get_deployment_configs(self.config_name, self.instance_type)
+        )
 
     def _create_sagemaker_model(
         self,
@@ -873,71 +871,46 @@ def register_deploy_wrapper(*args, **kwargs):
 
         return model_package
 
-    @lru_cache
-    def _get_deployment_configs_benchmarks_data(
-        self, config_name: str, instance_type: str
-    ) -> Dict[str, Any]:
+    @_deployment_config_lru_cache
+    def _get_deployment_configs_benchmarks_data(self) -> Dict[str, Any]:
         """Deployment configs benchmark metrics.
 
-        Args:
-            config_name (str): Name of selected deployment config.
-            instance_type (str): The selected Instance type.
         Returns:
             Dict[str, List[str]]: Deployment config benchmark data.
         """
         return get_metrics_from_deployment_configs(
-            self._get_deployment_configs(config_name, instance_type)
+            self._get_deployment_configs(None, None),
         )
 
-    @lru_cache
-    def _retrieve_selected_deployment_config(
-        self, config_name: str, instance_type: str
-    ) -> Optional[DeploymentConfigMetadata]:
-        """Retrieve the deployment config to apply to `This` model.
-
-        Args:
-            config_name (str): The name of the deployment config to retrieve.
-            instance_type (str): The instance type of the deployment config to retrieve.
-        Returns:
-            Optional[Dict[str, Any]]: The retrieved deployment config.
-        """
-        if config_name is None:
-            return None
-
-        for deployment_config in self._get_deployment_configs(config_name, instance_type):
-            if deployment_config.deployment_config_name == config_name:
-                return deployment_config
-        return None
-
-    @lru_cache
+    @_deployment_config_lru_cache
     def _get_deployment_configs(
-        self, selected_config_name: str, selected_instance_type: str
+        self, selected_config_name: Optional[str], selected_instance_type: Optional[str]
     ) -> List[DeploymentConfigMetadata]:
         """Retrieve deployment configs metadata.
 
         Args:
-            selected_config_name (str): The name of the selected deployment config.
-            selected_instance_type (str): The selected instance type.
+            selected_config_name (Optional[str]): The name of the selected deployment config.
+            selected_instance_type (Optional[str]): The selected instance type.
         """
         deployment_configs = []
-        if self._metadata_configs is None:
+        if not self._metadata_configs:
             return deployment_configs
 
         err = None
         for config_name, metadata_config in self._metadata_configs.items():
-            if err is None or "is not authorized to perform: pricing:GetProducts" not in err:
-                err, metadata_config.benchmark_metrics = (
-                    add_instance_rate_stats_to_benchmark_metrics(
-                        self.region, metadata_config.benchmark_metrics
-                    )
-                )
-
             resolved_config = metadata_config.resolved_config
             if selected_config_name == config_name:
                 instance_type_to_use = selected_instance_type
             else:
                 instance_type_to_use = resolved_config.get("default_inference_instance_type")
 
+            if metadata_config.benchmark_metrics:
+                err, metadata_config.benchmark_metrics = (
+                    add_instance_rate_stats_to_benchmark_metrics(
+                        self.region, metadata_config.benchmark_metrics
+                    )
+                )
+
             init_kwargs = get_init_kwargs(
                 model_id=self.model_id,
                 instance_type=instance_type_to_use,
@@ -957,9 +930,9 @@ def _get_deployment_configs(
             )
             deployment_configs.append(deployment_config_metadata)
 
-        if err is not None and "is not authorized to perform: pricing:GetProducts" in err:
+        if err and err["Code"] == "AccessDeniedException":
             error_message = "Instance rate metrics will be omitted. Reason: %s"
-            JUMPSTART_LOGGER.warning(error_message, err)
+            JUMPSTART_LOGGER.warning(error_message, err["Message"])
 
         return deployment_configs
 

src/sagemaker/jumpstart/types.py

@@ -2255,6 +2255,8 @@ def _val_to_json(self, val: Any) -> Any:
             Any: The converted json value.
         """
         if issubclass(type(val), JumpStartDataHolderType):
+            if isinstance(val, JumpStartBenchmarkStat):
+                val.name = val.name.replace("_", " ").title()
             return val.to_json()
         if isinstance(val, list):
             list_obj = []