Fix metric column name

Author: Jonathan Makunga

src/sagemaker/jumpstart/utils.py

@@ -1220,6 +1220,8 @@ def get_metrics_from_deployment_configs(
     if not deployment_configs:
         return {}
 
+    print("deployment_configs: {}".format(deployment_configs))
+
     data = {"Instance Type": [], "Config Name": [], "Concurrent Users": []}
     instance_rate_data = {}
     for index, deployment_config in enumerate(deployment_configs):
@@ -1256,26 +1258,27 @@ def get_metrics_from_deployment_configs(
                     instance_rate_data[instance_rate_column_name].append(instance_type_rate.value)
 
                 for metric in metrics:
-                    column_name = _normalize_benchmark_metric_column_name(metric.name)
+                    column_name = _normalize_benchmark_metric_column_name(metric.name, metric.unit)
                     data[column_name] = data.get(column_name, [])
                     data[column_name].append(metric.value)
 
     data = {**data, **instance_rate_data}
     return data
 
 
-def _normalize_benchmark_metric_column_name(name: str) -> str:
+def _normalize_benchmark_metric_column_name(name: str, unit: str) -> str:
     """Normalizes benchmark metric column name.
 
     Args:
         name (str): Name of the metric.
+        unit (str): Unit of the metric.
     Returns:
         str: Normalized metric column name.
     """
     if "latency" in name.lower():
-        name = "Latency for each user (TTFT in ms)"
+        name = f"Latency, TTFT (P50 in {unit.lower()})"
     elif "throughput" in name.lower():
-        name = "Throughput per user (token/seconds)"
+        name = f"Throughput (P50 in {unit.lower()}/user)"
     return name
 
 

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -1979,14 +1979,14 @@ def test__normalize_benchmark_metrics():
 
 
 @pytest.mark.parametrize(
-    "name, expected",
+    "name, unit, expected",
     [
-        ("latency", "Latency for each user (TTFT in ms)"),
-        ("throughput", "Throughput per user (token/seconds)"),
+        ("latency", "sec", "Latency, TTFT (P50 in sec)"),
+        ("throughput", "tokens/sec", "Throughput (P50 in tokens/sec/user)"),
     ],
 )
-def test__normalize_benchmark_metric_column_name(name, expected):
-    out = utils._normalize_benchmark_metric_column_name(name)
+def test_normalize_benchmark_metric_column_name(name, unit, expected):
+    out = utils._normalize_benchmark_metric_column_name(name, unit)
 
     assert out == expected