Reintroduce the "validate_jobset_replica_number" function in the jobset_util.py file

Refactor for new DAG formatting

Author: QuinnMMcGarry

dags/tpu_observability/jobset_healthiness_ready.py

@@ -144,10 +144,8 @@ def generate_second_node_pool_name(
 
       start_workload = jobset.run_workload(
           node_pool=cluster_info,
-          yaml_config=jobset_config.generate_yaml(
-              workload_script=Workload.IDLE_READY_TPU_20M
-          ),
-          namespace=jobset_config.namespace,
+          jobset_config=jobset_config,
+          workload_type=Workload.IDLE_READY_TPU_20M,
       )
 
       validate_ready_replicas = jobset.validate_jobset_replica_number(
@@ -161,8 +159,7 @@ def generate_second_node_pool_name(
           task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
       )(
           node_pool=cluster_info,
-          jobset_name=jobset_config.jobset_name,
-          namespace=jobset_config.namespace,
+          jobset_config=jobset_config,
       ).as_teardown(
           setups=start_workload
       )

dags/tpu_observability/utils/jobset_util.py

@@ -882,3 +882,17 @@ def ensure_no_jobset_uptime_data(
     logging.info("Stability period passed with no data detected.")
     return True
   return False
+
+
+@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,
+):
+  name = jobset_config.replicated_job_name
+  found_replica_number = get_replica_num(
+      replica_type=replica_type, job_name=name, node_pool=node_pool
+  )
+  return found_replica_number == correct_replica_num