Change the "validate_jobset_replica_number" function to work using the

new jobset storage method. Since "jobset_confifg.replicas" cannot be
accessed in DAG itself.

Author: QuinnMMcGarry

dags/tpu_observability/jobset_healthiness_ready.py

@@ -139,7 +139,7 @@ def generate_second_node_pool_name(
           node_pool=cluster_info,
           jobset_config=jobset_config,
           replica_type="ready",
-          correct_replica_num=0,
+          pre_startup=True,
       )
 
       start_workload = jobset.run_workload(
@@ -152,7 +152,7 @@ def generate_second_node_pool_name(
           node_pool=cluster_info,
           jobset_config=jobset_config,
           replica_type="ready",
-          correct_replica_num=jobset_config.replicas,
+          pre_startup=False,
       )
 
       cleanup_workload = jobset.end_workload.override(

dags/tpu_observability/utils/jobset_util.py

@@ -889,10 +889,12 @@ def validate_jobset_replica_number(
     node_pool: node_pool_info,
     jobset_config: JobSet,
     replica_type: str,
-    correct_replica_num: int,
+    pre_startup: bool,
 ):
   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
+  if pre_startup:
+    return found_replica_number == 0
+  return found_replica_number == jobset_config.replicas