feat(tpu_observability): use custom labels for distribution metrics in tpu_info verification

Author: yuna-tzeng

dags/tpu_observability/tpu_info_metric.py

@@ -81,6 +81,16 @@ def parse_from_tpu_info(
     """
     pass
 
+  def get_labels(
+      self, tpu_info_metric_output: list[tpu_info.Table]
+  ) -> list[str] | None:
+    """Gets labels for the metric values.
+
+    Default returns None, which implies using default device labels.
+    """
+    return None
+
+
 
 class _BaseSimplePointStrategy(BaseMetricStrategy):
   """Base strategy for parsing single numeric data points from Cloud Monitoring.
@@ -325,6 +335,30 @@ def parse_from_tpu_info(
 
     return tpu_info_data_values
 
+  def get_labels(
+      self, tpu_info_metric_output: list[tpu_info.Table]
+  ) -> list[str] | None:
+    """Parses labels for percentiles from `tpu-info` output tables."""
+    parsed_values_by_group: dict[str, dict[str, float]] = {}
+    table_name = self._tpu_info_table_name
+    group_key = self._tpu_info_group_by_key
+
+    for metric_table in tpu_info_metric_output:
+      if metric_table.name == table_name:
+        for row_dict in metric_table.body:
+          group_value = row_dict.get(group_key, "summary")
+
+          if group_value not in parsed_values_by_group:
+            parsed_values_by_group[group_value] = {}
+
+    labels = []
+    for group_value in sorted(parsed_values_by_group.keys()):
+      for p in self.percentiles_to_check:
+        labels.append(f"{group_value} {p.name}")
+
+    return labels
+
+
 
 class MemoryUsedStrategy(_BaseSimplePointStrategy):
   """Strategy for verifying Used HBM Memory."""

dags/tpu_observability/tpu_info_metrics_dag.py

@@ -61,6 +61,7 @@ def compare_metric_values(
     pod_name: str,
     metric_display_name: str,
     tolerance_percent: float,
+    labels: list[str] | None = None,
 ):
   """Compares two lists of metric values and checks if they are within a tolerance range."""
   if len(cmd_values) != len(monitoring_values):
@@ -76,15 +77,15 @@ def compare_metric_values(
       metric_display_name,
   )
   logging.info(
-      "%-12s%-15s%-17s%-12s%-15s%-10s",
-      "Device",
+      "%-20s%-15s%-17s%-12s%-15s%-10s",
+      "Label" if labels else "Device",
       "TPU-Info Val",
       "Monitoring Val",
       "Difference",
       "Allowed Diff",
       "Result",
   )
-  logging.info("-" * 85)
+  logging.info("-" * 95)
 
   all_passed = True
   for i, (log_val, mon_val) in enumerate(zip(cmd_values, monitoring_values)):
@@ -93,9 +94,10 @@ def compare_metric_values(
     passed = diff <= allowed_diff
     if not passed:
       all_passed = False
+    label = labels[i] if labels else f"Device {i}"
     logging.info(
-        "%-12s%-15.2f%-17.2f%-12.2f%-15.2f%-10s",
-        f"Device {i}",
+        "%-20s%-15.2f%-17.2f%-12.2f%-15.2f%-10s",
+        label,
         log_val,
         mon_val,
         diff,
@@ -114,6 +116,7 @@ def compare_metric_values(
   )
 
 
+
 @task
 def get_tpu_info_metric_from_pod(
     node_pool: node_pool.Info,
@@ -172,14 +175,17 @@ def run_metric_verification(
       metric_strategy.dag_id_suffix,
   )
 
+  labels = metric_strategy.get_labels(tpu_info_output)
   compare_metric_values(
       cmd_values,
       monitoring_values,
       pod_name,
       metric_display_name=metric_strategy.dag_id_suffix,
       tolerance_percent=tolerance_for_metric,
+      labels=labels,
   )
 
+
   return True