diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py index f939bc303b..d98b9e7dd6 100644 --- a/src/sagemaker/jumpstart/model.py +++ b/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, diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py index 3fce6dd105..357bdb6eb7 100644 --- a/src/sagemaker/jumpstart/utils.py +++ b/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. diff --git a/tests/unit/sagemaker/jumpstart/model/test_model.py b/tests/unit/sagemaker/jumpstart/model/test_model.py index 5d8d048501..3f07df8bfa 100644 --- a/tests/unit/sagemaker/jumpstart/model/test_model.py +++ b/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", diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py index 210bd8e074..f576e36185 100644 --- a/tests/unit/sagemaker/jumpstart/test_utils.py +++ b/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 diff --git a/tests/unit/sagemaker/jumpstart/utils.py b/tests/unit/sagemaker/jumpstart/utils.py index 96662837b4..77913ea73e 100644 --- a/tests/unit/sagemaker/jumpstart/utils.py +++ b/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,