Fix format/linting

Linting fix

Delete dags/dags/tpu_observability directory

Delete dags/airflow_monitoring.py

Author: QuinnMMcGarry

dags/airflow_monitoring.py

@@ -197,6 +197,5 @@ def generate_second_node_pool_name(
           >> validate_ready_replicas
           >> cleanup_workload
           >> cleanup_node_pool
-
       )
       # pylint: enable=pointless-statement

…bservability/jobset_healthiness_ready.py …bservability/jobset_healthiness_ready.pydags/dags/tpu_observability/jobset_healthiness_ready.py renamed to dags/tpu_observability/jobset_healthiness_ready.py dags/dags/tpu_observability/jobset_healthiness_ready.py renamed to dags/tpu_observability/jobset_healthiness_ready.py

@@ -563,21 +563,21 @@ def validate_jobset_replica_number(
     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")