Benchmark feature v2 (#4618)

* Add funtionalities to get and set deployment config

* Resolve PR comments

* ModelBuilder-JS

* Add Unit tests

* Refactoring

* Testing with Notebook

* Test backward compatibility

* Remove Accelerated column if all not enabled

* Fix docstring

* Resolved PR Review comments

* Docstring

* increase code coverage

* Testing fix with Notebook

* Only fetch instance rate metrics if not present

* Increase code coverage

---------

Co-authored-by: Jonathan Makunga <[email protected]>

Author: makungaj1

src/sagemaker/jumpstart/model.py

@@ -48,7 +48,7 @@
     validate_model_id_and_get_type,
     verify_model_region_and_return_specs,
     get_jumpstart_configs,
-    extract_metrics_from_deployment_configs,
+    get_metrics_from_deployment_configs,
 )
 from sagemaker.jumpstart.constants import JUMPSTART_LOGGER
 from sagemaker.jumpstart.enums import JumpStartModelType
@@ -868,7 +868,7 @@ def _get_benchmarks_data(self, config_name: str) -> Dict[str, List[str]]:
         Returns:
             Dict[str, List[str]]: Deployment config benchmark data.
         """
-        return extract_metrics_from_deployment_configs(
+        return get_metrics_from_deployment_configs(
             self._deployment_configs,
             config_name,
         )
@@ -905,20 +905,29 @@ def _convert_to_deployment_config_metadata(
             "default_inference_instance_type"
         )
 
-        instance_rate = get_instance_rate_per_hour(
-            instance_type=default_inference_instance_type, region=self.region
-        )
-
         benchmark_metrics = (
             metadata_config.benchmark_metrics.get(default_inference_instance_type)
             if metadata_config.benchmark_metrics is not None
             else None
         )
-        if instance_rate is not None:
-            if benchmark_metrics is not None:
-                benchmark_metrics.append(JumpStartBenchmarkStat(instance_rate))
+
+        should_fetch_instance_rate_metric = True
+        if benchmark_metrics is not None:
+            for benchmark_metric in benchmark_metrics:
+                if benchmark_metric.name.lower() == "instance rate":
+                    should_fetch_instance_rate_metric = False
+                    break
+
+        if should_fetch_instance_rate_metric:
+            instance_rate = get_instance_rate_per_hour(
+                instance_type=default_inference_instance_type, region=self.region
+            )
+            instance_rate_metric = JumpStartBenchmarkStat(instance_rate)
+
+            if benchmark_metrics is None:
+                benchmark_metrics = [instance_rate_metric]
             else:
-                benchmark_metrics = [JumpStartBenchmarkStat(instance_rate)]
+                benchmark_metrics.append(instance_rate_metric)
 
         init_kwargs = get_init_kwargs(
             model_id=self.model_id,

src/sagemaker/jumpstart/utils.py

@@ -1030,7 +1030,7 @@ def get_jumpstart_configs(
     )
 
 
-def extract_metrics_from_deployment_configs(
+def get_metrics_from_deployment_configs(
     deployment_configs: List[Dict[str, Any]], config_name: str
 ) -> Dict[str, List[str]]:
     """Extracts metrics from deployment configs.

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -51,6 +51,7 @@
     get_base_spec_with_prototype_configs,
     get_mock_init_kwargs,
     get_base_deployment_configs,
+    get_base_spec_with_prototype_configs_with_missing_benchmarks,
 )
 import boto3
 
@@ -1790,7 +1791,7 @@ def test_model_retrieve_deployment_config(
 
         mock_get_init_kwargs.side_effect = lambda *args, **kwargs: get_mock_init_kwargs(model_id)
         mock_verify_model_region_and_return_specs.side_effect = (
-            lambda *args, **kwargs: get_base_spec_with_prototype_configs()
+            lambda *args, **kwargs: get_base_spec_with_prototype_configs_with_missing_benchmarks()
         )
         mock_get_instance_rate_per_hour.side_effect = lambda *args, **kwargs: {
             "name": "Instance Rate",
@@ -1838,7 +1839,7 @@ def test_model_display_benchmark_metrics(
 
         mock_get_init_kwargs.side_effect = lambda *args, **kwargs: get_mock_init_kwargs(model_id)
         mock_verify_model_region_and_return_specs.side_effect = (
-            lambda *args, **kwargs: get_base_spec_with_prototype_configs()
+            lambda *args, **kwargs: get_base_spec_with_prototype_configs_with_missing_benchmarks()
         )
         mock_get_instance_rate_per_hour.side_effect = lambda *args, **kwargs: {
             "name": "Instance Rate",

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -1810,6 +1810,6 @@ def test_get_jumpstart_benchmark_stats_training(
     ],
 )
 def test_extract_metrics_from_deployment_configs(config_name, configs, expected):
-    data = utils.extract_metrics_from_deployment_configs(configs, config_name)
+    data = utils.get_metrics_from_deployment_configs(configs, config_name)
 
     assert data == expected

tests/unit/sagemaker/jumpstart/utils.py

@@ -226,6 +226,26 @@ def get_base_spec_with_prototype_configs(
     return JumpStartModelSpecs(spec)
 
 
+def get_base_spec_with_prototype_configs_with_missing_benchmarks(
+    region: str = None,
+    model_id: str = None,
+    version: str = None,
+    s3_client: boto3.client = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+) -> JumpStartModelSpecs:
+    spec = copy.deepcopy(BASE_SPEC)
+    copy_inference_configs = copy.deepcopy(INFERENCE_CONFIGS)
+    copy_inference_configs["inference_configs"]["neuron-inference"]["benchmark_metrics"] = None
+
+    inference_configs = {**INFERENCE_CONFIGS, **INFERENCE_CONFIG_RANKINGS}
+    training_configs = {**TRAINING_CONFIGS, **TRAINING_CONFIG_RANKINGS}
+
+    spec.update(inference_configs)
+    spec.update(training_configs)
+
+    return JumpStartModelSpecs(spec)
+
+
 def get_prototype_spec_with_configs(
     region: str = None,
     model_id: str = None,