fix: Docstring

Author: chiajunglien

dags/tpu_observability/jobset_healthiness_validation.py

@@ -24,7 +24,7 @@
 from dags import composer_env
 from dags.tpu_observability.utils import jobset_util as jobset
 from dags.tpu_observability.utils import node_pool_util as node_pool
-from dags.tpu_observability.utils.jobset_util import Workload, ReplicatedJobStatus
+from dags.tpu_observability.utils.jobset_util import Workload, JobSetHealthiness
 from dags.tpu_observability.configs.common import (
     MachineConfigMap,
     GCS_CONFIG_PATH,
@@ -64,18 +64,20 @@
     doc_md="""
       # JobSet Healthiness Test For the "Suspended" Status
       ### Description
-      This DAG automates the process of creating node-pools, ensuring the
-      correct number of "Suspended" replicas appear, then launching a jobset on
-      multiple replicas to ensure the correct number begin running.
+      This DAG automates node-pool creation and validates JobSet healthiness
+      by examining replica-based metrics: Specified, Active, Ready,
+      Suspended, Succeeded, and Failed. It ensures the JobSet controller
+      accurately reports these states during startup, maintenance,
+      and failure scenarios.
       ### Prerequisites
       This test requires an existing cluster to run.
       ### Procedures
-      First a node-pool is created. The validation test is then run to
-      check if the number of "Suspended" replicas is 0. Once the jobset is
-      running the jobs should quickly enter the "Ready" state. Then using
-      command to suspend entire jobset. The number of found replicas is
-      tested against the number of replicas which should be "Suspended".
-      If they match the DAG is a success.
+      First a node-pool is created. This test uses a State-Trigger-Observe pattern:
+      it triggers lifecycle transitions (e.g., suspension, failure or succeeded)
+      and verifies that GKE telemetry reflects these shifts. Using sensors,
+      the DAG polls for eventual consistency to account for ingestion latency,
+      dynamically matching runtime JobSet configurations against normalized monitoring
+      data types to ensure accurate state validation.
       """,
 ) as dag:
   for machine in MachineConfigMap:
@@ -119,12 +121,12 @@
 
       with TaskGroup(group_id="validate_running_metrics") as validate_running:
         running_metrics = [
-            (ReplicatedJobStatus.SPECIFIED, "USE_CONFIG_REPLICAS"),
-            (ReplicatedJobStatus.ACTIVE, "USE_CONFIG_REPLICAS"),
-            (ReplicatedJobStatus.READY, "USE_CONFIG_REPLICAS"),
-            (ReplicatedJobStatus.FAILED, 0),
-            (ReplicatedJobStatus.SUCCEEDED, 0),
-            (ReplicatedJobStatus.SUSPENDED, 0),
+            (JobSetHealthiness.SPECIFIED, "USE_CONFIG_REPLICAS"),
+            (JobSetHealthiness.ACTIVE, "USE_CONFIG_REPLICAS"),
+            (JobSetHealthiness.READY, "USE_CONFIG_REPLICAS"),
+            (JobSetHealthiness.FAILED, 0),
+            (JobSetHealthiness.SUCCEEDED, 0),
+            (JobSetHealthiness.SUSPENDED, 0),
         ]
         for status, expected in running_metrics:
           jobset.wait_for_jobset_metrics.override(
@@ -147,8 +149,8 @@
           group_id="validate_suspended_metrics"
       ) as validate_suspended:
         suspended_metrics = [
-            (ReplicatedJobStatus.ACTIVE, 0),
-            (ReplicatedJobStatus.SUSPENDED, "USE_CONFIG_REPLICAS"),
+            (JobSetHealthiness.ACTIVE, 0),
+            (JobSetHealthiness.SUSPENDED, "USE_CONFIG_REPLICAS"),
         ]
         for status, expected in suspended_metrics:
           jobset.wait_for_jobset_metrics.override(
@@ -184,7 +186,7 @@
         validate_succeeded_metric = jobset.wait_for_jobset_metrics.override(
             task_id="wait_for_succeeded_count"
         )(
-            metric_name=ReplicatedJobStatus.SUCCEEDED,
+            metric_name=JobSetHealthiness.SUCCEEDED,
             expected_value="USE_CONFIG_REPLICAS",
             node_pool=cluster_info,
             jobset_config=jobset_config,
@@ -209,7 +211,7 @@
         validate_failed_metric = jobset.wait_for_jobset_metrics.override(
             task_id="wait_for_failed_count"
         )(
-            metric_name=ReplicatedJobStatus.FAILED,
+            metric_name=JobSetHealthiness.FAILED,
             expected_value="USE_CONFIG_REPLICAS",
             node_pool=cluster_info,
             jobset_config=jobset_config,

dags/tpu_observability/utils/jobset_util.py

@@ -326,8 +326,11 @@ class JobSetStartupOutput:
   jobset_start_time: XComArg
 
 
-class ReplicatedJobStatus(enum.Enum):
-  """Defines status of a replicated job."""
+class JobSetHealthiness(enum.Enum):
+  """Enumeration of JobSet healthiness metrics based on replica states.
+  The values represent the metric name suffix used in Prometheus/GKE
+  observability to track the count of replicas in each lifecycle state.
+  """
 
   READY = "ready"
   SUSPENDED = "suspended"
@@ -899,7 +902,7 @@ def wait_for_jobset_started(
 
 @task.sensor(poke_interval=60, timeout=3600, mode="poke")
 def wait_for_jobset_metrics(
-    metric_name: ReplicatedJobStatus,
+    metric_name: JobSetHealthiness,
     expected_value: any,
     node_pool: node_pool_info,
     jobset_config: JobSet,