-Created "validate_jobset_replica_number" funtion in jobset_util.py

QuinnMMcGarry · QuinnMMcGarry · commit aea57e1cfc81 · 2025-12-28T23:21:54.000Z
-Changed workload name to "IDLE_READY_TPU_20M"
-Changed schedule time and added statement to only run on PROD environment

Changed to new nodepool creation format

Remove unused DAG flow
diff --git a/dags/tpu_observability/jobset_healthiness_ready.py b/dags/tpu_observability/jobset_healthiness_ready.py
@@ -16,41 +16,24 @@
 
 import datetime
 
-from airflow.decorators import task
 from airflow import models
+from airflow.decorators import task
 from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
 
-from dags.common.vm_resource import Region, Zone
+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 JobSet, Workload
-from dags.tpu_observability.configs.common import MachineConfigMap
-from dags import composer_env
-
-# Can be moved to jobset_util.py
-@task.sensor(poke_interval=30, timeout=900, mode="reschedule")
-def validate_replica_number(
-    node_pool: node_pool.Info,
-    jobset_config: JobSet,
-    replica_type: str,
-    replica_num: int
-):
-  found_replia_num = jobset.get_replica_num(
-      replica_type=replica_type,
-      job_name=jobset_config.replicated_job_name,
-      node_pool=node_pool
-  )
-
-  return (found_replia_num == replica_num)
+from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH
 
 
 # Keyword arguments are generated dynamically at runtime (pylint does not
 # know this signature).
 with models.DAG(  # pylint: disable=unexpected-keyword-arg
     dag_id="jobset_healthiness_ready",
     start_date=datetime.datetime(2025, 8, 10),
-    schedule="00 03 * * *",
+    schedule="30 19 * * *" if composer_env.is_prod_env() else None,
     catchup=False,
     tags=[
         "cloud-ml-auto-solutions",
@@ -85,54 +68,15 @@ def validate_replica_number(
       match the DAG is a success.
       """,
 ) as dag:
-  cluster_name = "tpu-observability-automation"
-  cluster_name += "-prod" if composer_env.is_prod_env() else "-dev"
-
   for machine in MachineConfigMap:
     config = machine.value
-    cluster_info = node_pool.Info(
-        project_id=models.Variable.get("PROJECT_ID", default_var="cienet-cmcs"),
-        cluster_name=models.Variable.get(
-            "CLUSTER_NAME", default_var=cluster_name
-        ),
-        node_pool_name=models.Variable.get(
-            "NODE_POOL_NAME", default_var="jobset-healthiness-ready-v6e"
-        ),
-        region=models.Variable.get(
-            "REGION", default_var=Region.US_CENTRAL1.value
-        ),
-        location=models.Variable.get(
-            "LOCATION", default_var=Region.US_CENTRAL1.value
-        ),
-        node_locations=models.Variable.get(
-            "LOCATIONS", default_var=Zone.US_CENTRAL1_B.value
-        ),
-        num_nodes=models.Variable.get("NUM_NODES", default_var=4),
-        machine_type=config.machine_version.value,
-        tpu_topology=config.tpu_topology,
-    )
 
-    cluster_info_2 = node_pool.Info(
-        project_id=models.Variable.get("PROJECT_ID", default_var="cienet-cmcs"),
-        cluster_name=models.Variable.get(
-            "CLUSTER_NAME", default_var=cluster_name
-        ),
-        node_pool_name=models.Variable.get(
-            "NODE_POOL_NAME", default_var="jobset-healthiness-ready-v6e-2"
-        ),
-        region=models.Variable.get(
-            "REGION", default_var=Region.US_CENTRAL1.value
-        ),
-        location=models.Variable.get(
-            "LOCATION", default_var=Region.US_CENTRAL1.value
-        ),
-        node_locations=models.Variable.get(
-            "LOCATIONS", default_var=Zone.US_CENTRAL1_B.value
-        ),
-        num_nodes=models.Variable.get("NUM_NODES", default_var=4),
-        machine_type=config.machine_version.value,
-        tpu_topology=config.tpu_topology,
-    )
+    @task
+    def generate_second_node_pool_name(
+        node_pool_info: node_pool.Info,
+    ) -> str:
+      """Generates a second node pool name."""
+      return f"{node_pool_info.node_pool_name}-2"
 
     jobset_config = JobSet(
         jobset_name="jobset-healthiness-ready",
@@ -155,36 +99,60 @@ def validate_replica_number(
     with TaskGroup(  # pylint: disable=unexpected-keyword-arg
         group_id=f"v{config.tpu_version.value}"
     ):
-      create_node_pool = node_pool.create(
-          node_pool=cluster_info,
-          reservation="cloudtpu-20251107233000-1246578561",
+      cluster_info = node_pool.build_node_pool_info_from_gcs_yaml.override(
+          task_id="build_node_pool_info_from_gcs_yaml"
+      )(
+          gcs_path=GCS_CONFIG_PATH,
+          dag_name="jobset_healthiness_ready",
+          is_prod=composer_env.is_prod_env(),
+          machine_type=config.machine_version.value,
+          tpu_topology=config.tpu_topology,
       )
 
-      create_node_pool_2 = node_pool.create(
-          node_pool=cluster_info_2,
-          reservation="cloudtpu-20251107233000-1246578561",
+      cluster_info_2 = node_pool.copy_node_pool_info_with_override(
+          info=cluster_info,
+          node_pool_name=generate_second_node_pool_name(cluster_info),
       )
 
-      validate_zero_replicas = validate_replica_number(
+      # Keyword arguments are generated dynamically at runtime (pylint does not
+      # know this signature).
+      with TaskGroup(  # pylint: disable=unexpected-keyword-arg
+          group_id="create_node_pool"
+      ) as create_node_pool:
+        create_first_node_pool = node_pool.create.override(
+            task_id="node_pool_1",
+            retries=2,
+        )(
+            node_pool=cluster_info,
+        )
+
+        create_second_node_pool = node_pool.create.override(
+            task_id="node_pool_2",
+            retries=2,
+        )(
+            node_pool=cluster_info_2,
+        )
+
+      validate_zero_replicas = jobset.validate_jobset_replica_number(
           node_pool=cluster_info,
           jobset_config=jobset_config,
           replica_type="ready",
-          replica_num=0
+          correct_replica_num=0,
       )
 
       start_workload = jobset.run_workload(
           node_pool=cluster_info,
           yaml_config=jobset_config.generate_yaml(
-              workload_script=Workload.READY_TPU
+              workload_script=Workload.IDLE_READY_TPU_20M
           ),
           namespace=jobset_config.namespace,
       )
 
-      validate_ready_replicas = validate_replica_number(
+      validate_ready_replicas = jobset.validate_jobset_replica_number(
           node_pool=cluster_info,
           jobset_config=jobset_config,
           replica_type="ready",
-          replica_num=jobset_config.replicas
+          correct_replica_num=jobset_config.replicas,
       )
 
       cleanup_workload = jobset.end_workload.override(
@@ -197,29 +165,38 @@ def validate_replica_number(
           setups=start_workload
       )
 
-      cleanup_node_pool = node_pool.delete.override(
-          task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE
-      )(node_pool=cluster_info).as_teardown(
-          setups=create_node_pool,
-      )
-
-      cleanup_node_pool_2 = node_pool.delete.override(
-          task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE
-      )(node_pool=cluster_info_2).as_teardown(
-          setups=create_node_pool_2,
-      )
+      # Keyword arguments are generated dynamically at runtime (pylint does not
+      # know this signature).
+      with TaskGroup(  # pylint: disable=unexpected-keyword-arg
+          group_id="cleanup_node_pool"
+      ) as cleanup_node_pool:
+        cleanup_first_node_pool = node_pool.delete.override(
+            task_id="cleanup_node_pool_1",
+            trigger_rule=TriggerRule.ALL_DONE,
+            retries=2,
+        )(node_pool=cluster_info).as_teardown(
+            setups=create_node_pool,
+        )
+
+        cleanup_second_node_pool = node_pool.delete.override(
+            task_id="cleanup_node_pool_2",
+            trigger_rule=TriggerRule.ALL_DONE,
+            retries=2,
+        )(node_pool=cluster_info_2).as_teardown(
+            setups=create_node_pool,
+        )
 
       # Airflow uses >> for task chaining, which is pointless for pylint.
       # pylint: disable=pointless-statement
       (
-          create_node_pool
-          >> create_node_pool_2
+          cluster_info
+          >> cluster_info_2
+          >> create_node_pool
           >> validate_zero_replicas
           >> start_workload
           >> validate_ready_replicas
           >> cleanup_workload
           >> cleanup_node_pool
-          >> cleanup_node_pool_2
 
       )
       # pylint: enable=pointless-statement
diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py
@@ -43,13 +43,13 @@ class Workload:
   Each workload is a JSON-escaped string, ready to be used as a shell argument.
   """
 
-  READY_TPU = json.dumps(
+  IDLE_READY_TPU_20M = json.dumps(
       textwrap.dedent(
           """
-          sleep 1000
+          sleep 1200
           """
       ),
-      ensure_ascii=False
+      ensure_ascii=False,
   )
 
   JAX_TPU_BENCHMARK = json.dumps(
@@ -560,26 +560,29 @@ def wait_for_jobset_ttr_to_be_found(node_pool: node_pool.Info) -> bool:
   return len(time_series) > 0
 
 
-@task.sensor(poke_interval=30, timeout=600, mode="reschedule")
-def wait_for_jobset_status_occurrence(
-    replica_type: str, job_name: str, node_pool: node_pool.Info
+@task.sensor(poke_interval=30, timeout=900, mode="reschedule")
+def validate_jobset_replica_number(
+    node_pool: node_pool.Info,
+    jobset_config: JobSet,
+    replica_type: str,
+    correct_replica_num: int,
 ):
   """
-  A sensor which checks if are any jobset replicas in a status type.
+  A sensor which checks if the correct number jobset replicas in a status type.
 
   Args:
-    replica_type(str): The type of status being checked for.
-    job_name(str): The name of the job replica which is run from the jobset.
-    node_pool(Info): The Info object containing the cluster information needed
-    for the kubernetes API to connect to it.
+    node_pool: Configuration object with cluster details.
+    jobset_config: Configuration of JobSet which is being run.
+    replica_type(str): The name of the type of status being checked for.
+    correct_replica_num(int): The expected number of replicas to be found.
   """
   logging.info("Checking for number of replicas of type: %s", replica_type)
   ready_replicas = get_replica_num(
       replica_type=replica_type,
-      job_name=job_name,
+      job_name=jobset_config.replicated_job_name,
       node_pool=node_pool,
   )
-  return ready_replicas > 0
+  return ready_replicas == correct_replica_num
 
 
 @task.sensor(poke_interval=30, timeout=600, mode="reschedule")
