fix

Author: yuna-tzeng

dags/tpu_observability/multi_host_nodepool_rollback_dag.py

@@ -156,7 +156,6 @@
         )(node_pool=node_pool_info)
 
       chain(
-          node_pool_info,
           pre_test,
           test,
           post_test,

dags/tpu_observability/node_pool_status.py

@@ -117,6 +117,20 @@
             node_pool=node_pool_info,
         )
 
+        # Intentionally create a node pool with problematic configurations
+        # to validate that it enters the ERROR state.
+        task_id = "create_problematic_node_pool_info"
+        create_problematic_node_pool_info = node_pool.create.override(
+            task_id=task_id,
+            owner=test_owner.YUNA_T,
+        )(
+            node_pool=problematic_node_pool_info,
+            # The failure is intentionally ignored because we want to validate
+            # that the status of the node pool (which fails to be created) is
+            # "ERROR".
+            ignore_failure=True,
+        )
+
         task_id = "wait_for_provisioning"
         wait_for_provisioning = node_pool.wait_for_status.override(
             task_id=task_id
@@ -165,30 +179,12 @@
             node_pool=node_pool_info, status=node_pool.Status.STOPPING
         )
 
-        # Intentionally create a node pool with problematic configurations
-        # to validate that it enters the ERROR state.
-        task_id = "create_problematic_node_pool_info"
-        create_problematic_node_pool_info = node_pool.create.override(
-            task_id=task_id,
-            owner=test_owner.YUNA_T,
-        )(
-            node_pool=problematic_node_pool_info,
-            # The failure is intentionally ignored because we want to validate
-            # that the status of the node pool (which fails to be created) is
-            # "ERROR".
-            ignore_failure=True,
-        )
-
         task_id = "wait_for_error"
-        wait_for_error = node_pool.wait_for_status.override(task_id=task_id)(
-            node_pool=problematic_node_pool_info, status=node_pool.Status.ERROR
-        )
-
-        task_id = "cleanup_wrong_node_pool"
-        cleanup_wrong_node_pool = node_pool.delete.override(
-            task_id=task_id, trigger_rule=TriggerRule.ALL_DONE
-        )(node_pool=problematic_node_pool_info).as_teardown(
-            setups=create_problematic_node_pool_info,
+        validate_problematic_node_pool_enter_error_state = (
+            node_pool.wait_for_status.override(task_id=task_id)(
+                node_pool=problematic_node_pool_info,
+                status=node_pool.Status.ERROR,
+            )
         )
 
         chain(
@@ -202,7 +198,7 @@
 
         chain(
             create_problematic_node_pool_info,
-            wait_for_error,
+            validate_problematic_node_pool_enter_error_state,
         )
 
       with TaskGroupWithTimeout(

dags/tpu_observability/update_node_pool_label.py

@@ -145,7 +145,6 @@
         )(node_pool=node_pool_info)
 
       chain(
-          node_pool_info,
           pre_test,
           test,
           post_test,

dags/tpu_observability/utils/node_pool_util.py

@@ -553,7 +553,7 @@ def _query_status_metric(node_pool: Info) -> Status:
   return Status.from_str(latest_status)
 
 
-@task.sensor(poke_interval=60, timeout=600, mode="poke")
+@task.sensor(poke_interval=60, timeout=600, mode="poke", retries=0)
 def wait_for_status(
     node_pool: Info,
     status: Status,
@@ -608,7 +608,7 @@ def rollback(node_pool: Info) -> None:
   return TimeUtil.from_datetime(current_time_utc)
 
 
-@task.sensor(poke_interval=30, timeout=1200, mode="poke")
+@task.sensor(poke_interval=30, timeout=1200, mode="poke", retries=0)
 def wait_for_availability(
     node_pool: Info,
     availability: bool,
@@ -684,7 +684,7 @@ def wait_for_availability(
   return availability == state
 
 
-@task.sensor(poke_interval=30, timeout=3600, mode="poke")
+@task.sensor(poke_interval=30, timeout=3600, mode="poke", retries=0)
 def wait_for_ttr(
     node_pool: Info,
     operation_start_time: TimeUtil,