Remove Accelerated column if all not enabled

Author: Jonathan Makunga

src/sagemaker/jumpstart/utils.py

@@ -1084,4 +1084,6 @@ def extract_metrics_from_deployment_configs(
                 if column_name in data.keys():
                     data[column_name].append(benchmark_metric.get("value"))
 
+    if "Yes" not in data["Accelerated"]:
+        del data["Accelerated"]
     return data

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -50,6 +50,7 @@
     get_special_model_spec,
     get_prototype_manifest,
     get_base_deployment_configs,
+    get_base_deployment_configs_with_acceleration_configs,
 )
 from mock import MagicMock
 
@@ -1763,10 +1764,11 @@ def test_get_jumpstart_benchmark_stats_training(
 
 
 @pytest.mark.parametrize(
-    "config_name, expected",
+    "config_name, configs, expected",
     [
         (
             None,
+            get_base_deployment_configs(),
             {
                 "Config Name": [
                     "neuron-inference",
@@ -1776,7 +1778,6 @@ def test_get_jumpstart_benchmark_stats_training(
                 ],
                 "Instance Type": ["ml.p2.xlarge", "ml.p2.xlarge", "ml.p2.xlarge", "ml.p2.xlarge"],
                 "Selected": ["No", "No", "No", "No"],
-                "Accelerated": ["Yes", "No", "No", "No"],
                 "Instance Rate (USD/Hrs)": [
                     "0.0083000000",
                     "0.0083000000",
@@ -1787,6 +1788,7 @@ def test_get_jumpstart_benchmark_stats_training(
         ),
         (
             "neuron-inference",
+            get_base_deployment_configs_with_acceleration_configs(),
             {
                 "Config Name": [
                     "neuron-inference",
@@ -1807,12 +1809,7 @@ def test_get_jumpstart_benchmark_stats_training(
         ),
     ],
 )
-def test_extract_metrics_from_deployment_configs(config_name, expected):
-    configs = get_base_deployment_configs()
-    configs[0]["AccelerationConfigs"] = [
-        {"Type": "Speculative-Decoding", "Enabled": True, "Spec": {"Version": "0.1"}}
-    ]
-
+def test_extract_metrics_from_deployment_configs(config_name, configs, expected):
     data = utils.extract_metrics_from_deployment_configs(configs, config_name)
 
     assert data == expected

tests/unit/sagemaker/jumpstart/utils.py

@@ -307,6 +307,14 @@ def get_base_deployment_configs() -> List[Dict[str, Any]]:
     return DEPLOYMENT_CONFIGS
 
 
+def get_base_deployment_configs_with_acceleration_configs() -> List[Dict[str, Any]]:
+    configs = copy.deepcopy(DEPLOYMENT_CONFIGS)
+    configs[0]["AccelerationConfigs"] = [
+        {"Type": "Speculative-Decoding", "Enabled": True, "Spec": {"Version": "0.1"}}
+    ]
+    return configs
+
+
 def get_mock_init_kwargs(model_id) -> JumpStartModelInitKwargs:
     return JumpStartModelInitKwargs(
         model_id=model_id,