refactor: encapsulate JobSet metric mapping into JobSetHealthiness enum

Author: chiajunglien

dags/tpu_observability/utils/jobset_util.py

@@ -339,6 +339,18 @@ class JobSetHealthiness(enum.Enum):
   SUCCEEDED = "succeeded"
   SPECIFIED = "specified"
 
+  @property
+  def healthiness_metric_type(self) -> str:
+    mapping = {
+        "ready": "prometheus.googleapis.com/kube_jobset_ready_replicas/gauge",
+        "active": "prometheus.googleapis.com/kube_jobset_active_replicas/gauge",
+        "suspended": "prometheus.googleapis.com/kube_jobset_suspended_replicas/gauge",
+        "succeeded": "prometheus.googleapis.com/kube_jobset_succeeded_replicas/gauge",
+        "failed": "prometheus.googleapis.com/kube_jobset_failed_replicas/gauge",
+        "specified": "prometheus.googleapis.com/kube_jobset_specified_replicas/gauge",
+    }
+    return mapping.get(self)
+
 
 class Command:
   """
@@ -996,21 +1008,9 @@ def wait_for_jobset_metrics(
   if expected_value == "USE_CONFIG_REPLICAS":
     final_expected = jobset_config.replicas
 
-  metric_mapping = {
-      "ready": "prometheus.googleapis.com/kube_jobset_ready_replicas/gauge",
-      "active": "prometheus.googleapis.com/kube_jobset_active_replicas/gauge",
-      "suspended": "prometheus.googleapis.com/kube_jobset_suspended_replicas/gauge",
-      "succeeded": "prometheus.googleapis.com/kube_jobset_succeeded_replicas/gauge",
-      "failed": "prometheus.googleapis.com/kube_jobset_failed_replicas/gauge",
-      "specified": "prometheus.googleapis.com/kube_jobset_specified_replicas/gauge",
-  }
+  metric_type = metric_name.healthiness_metric_type
+  name_str = metric_name.value
 
-  name_str = (
-      metric_name.value
-      if hasattr(metric_name, "value")
-      else str(metric_name).lower()
-  )
-  metric_type = metric_mapping.get(name_str)
   query_start = (
       start_time if start_time else TimeUtil.now() - timedelta(minutes=60)
   )