feat: Apply TaskGroupWithTimeout for TPU Observability DAGs - part 1 (GoogleCloudPlatform#1299)

Integrates `TaskGroupWithTimeout` directly into all TPU Observability
DAGs that do not depend on Kubernetes JobSets. This provides strict,
shared deadline timeout enforcement across distinct stages of GKE node
pool validation.

Specifically:
1. Configures stage-specific timeouts for all modified DAGs:
   - `PRE_TEST_TIMEOUT`: 10 minutes (for GKE provisioning and availability setup)
   - `POST_TEST_TIMEOUT`: 10 minutes (for GKE teardown and cleanup)
   - `TEST_TIMEOUT`: Evaluates dynamically as `DAGRUN_TIMEOUT - 20 mins`

2. Restructures the following 5 DAGs into three separate `TaskGroupWithTimeout` blocks:
   - `dags/tpu_observability/node_pool_status.py`
   - `dags/tpu_observability/node_pool_ttr_disk_size.py`
   - `dags/tpu_observability/node_pool_ttr_update_label.py`
   - `dags/tpu_observability/multi_host_nodepool_rollback_dag.py`
   - `dags/tpu_observability/update_node_pool_label.py`

3. Refactors GKE node pool operations into:
   - `pre_test`: Handles node pool creation and waiting for provisioning/running.
   - `test`: Executes the core test mutation and verification flow.
   - `post_test`: Manages cleanup/teardown of GKE node pools (using `is_teardown=True`
     to guarantee execution even if upstream stages timeout or fail).

* refactor(tpu-observability): wrap TPU observability DAGs with TaskGroupWithTimeout

Refactors five TPU observability DAGs to use TaskGroupWithTimeout for proper timeout
management during pre-test setup, test execution, and post-test teardown. This aligns
their implementation pattern with node_pool_status.py.

Specifically, this change:
- Imports TaskGroupWithTimeout in each DAG.
- Defines PRE_TEST_TIMEOUT, TEST_TIMEOUT, and POST_TEST_TIMEOUT at the module level.
- Groups tasks into pre_test, test, and post_test TaskGroupWithTimeout blocks.
- Chains tasks sequentially within each group and updates the overall DAG chain.

Modified DAGs:
- dags/tpu_observability/jobset_ttr_drain_restart.py
- dags/tpu_observability/jobset_ttr_node_pool_resize.py
- dags/tpu_observability/jobset_ttr_pod_delete.py
- dags/tpu_observability/jobset_ttr_rollback.py
- dags/tpu_observability/jobset_uptime_validation.py

Author: yuna-tzeng

dags/tpu_observability/jobset_ttr_drain_restart.py

@@ -21,6 +21,11 @@
 from airflow.models.baseoperator import chain
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
+from airflow.utils.task_group import TaskGroup
+from dags.common.task_group_with_timeout import TaskGroupWithTimeout
+
+
+from airflow.decorators import task
 
 from dags import composer_env
 from dags.common.scheduling_helper.scheduling_helper import (
@@ -42,6 +47,10 @@
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
 SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
 
+PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10)
+POST_TEST_TIMEOUT = datetime.timedelta(minutes=10)
+TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT
+
 
 @task
 def check_nodes_number(
@@ -140,75 +149,102 @@ def check_nodes_number(
       selector = jobset.generate_node_pool_selector(DAG_ID)
       jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix)
 
-      create_node_pool = node_pool.create.override(task_id="create_node_pool")(
-          node_pool=cluster_info,
-          node_pool_selector=selector,
-      )
-
-      startup = jobset.create_jobset_startup_tasks(
-          node_pool=cluster_info,
-          jobset_config=jobset_config,
-          jobset_name=jobset_name,
-          node_pool_selector=selector,
-          workload_type=Workload.JAX_TPU_BENCHMARK,
-      )
-
-      select_node = node_pool.draw_random_node.override(task_id="select_node")(
-          node_pool=cluster_info
-      )
-
-      drained_node = node_pool.operate_node.override(task_id="drained_node")(
-          node_pool=cluster_info,
-          operation=NodeOperationSpec.Drain(),
-          node_name=select_node,
-      )
-
-      check_nodes_number = check_nodes_number.override(
-          task_id="check_nodes_number"
-      )(
-          pool=cluster_info,
-          drained_node_number=1,
-      )
-
-      uncordon_node = node_pool.operate_node.override(task_id="uncordon_node")(
-          node_pool=cluster_info,
-          operation=NodeOperationSpec.Uncordon(),
-          node_name=select_node,
-      )
-
-      wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override(
-          task_id="wait_for_metric_upload"
-      )(
-          node_pool=cluster_info,
-          jobset_name=jobset_name,
-      )
-
-      cleanup_workload = jobset.end_workload.override(
-          task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
-      )(
-          node_pool=cluster_info,
-          jobset_config=jobset_config,
-          jobset_name=jobset_name,
-      ).as_teardown(
-          setups=startup.jobset_start_time
-      )
-
-      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,
-      )
+      with TaskGroupWithTimeout(
+          group_id="pre_test",
+          timeout=PRE_TEST_TIMEOUT,
+      ) as pre_test:
+        create_node_pool = node_pool.create.override(
+            task_id="create_node_pool"
+        )(
+            node_pool=cluster_info,
+            node_pool_selector=selector,
+        )
+
+      with TaskGroupWithTimeout(
+          group_id="test",
+          timeout=TEST_TIMEOUT,
+      ) as test:
+        startup = jobset.create_jobset_startup_tasks(
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+            jobset_name=jobset_name,
+            node_pool_selector=selector,
+            workload_type=Workload.JAX_TPU_BENCHMARK,
+        )
+
+        select_node = node_pool.draw_random_node.override(
+            task_id="select_node"
+        )(node_pool=cluster_info)
+
+        drained_node = node_pool.operate_node.override(task_id="drained_node")(
+            node_pool=cluster_info,
+            operation=NodeOperationSpec.Drain(),
+            node_name=select_node,
+        )
+
+        check_nodes_number_task = check_nodes_number.override(
+            task_id="check_nodes_number"
+        )(
+            pool=cluster_info,
+            drained_node_number=1,
+        )
+
+        uncordon_node = node_pool.operate_node.override(
+            task_id="uncordon_node"
+        )(
+            node_pool=cluster_info,
+            operation=NodeOperationSpec.Uncordon(),
+            node_name=select_node,
+        )
+
+        wait_for_metric_upload = (
+            jobset.wait_for_jobset_ttr_to_be_found.override(
+                task_id="wait_for_metric_upload"
+            )(
+                node_pool=cluster_info,
+                jobset_name=jobset_name,
+            )
+        )
+
+        chain(
+            *startup.tasks,
+            select_node,
+            drained_node,
+            check_nodes_number_task,
+            uncordon_node,
+            wait_for_metric_upload,
+        )
+
+      with TaskGroupWithTimeout(
+          group_id="post_test",
+          timeout=POST_TEST_TIMEOUT,
+          is_teardown=True,
+      ) as post_test:
+        cleanup_workload = jobset.end_workload.override(
+            task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
+        )(
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+            jobset_name=jobset_name,
+        ).as_teardown(
+            setups=startup.jobset_start_time
+        )
+
+        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,
+        )
+
+        chain(
+            cleanup_workload,
+            cleanup_node_pool,
+        )
 
       chain(
           selector,
           jobset_name,
-          create_node_pool,
-          *startup.tasks,
-          select_node,
-          drained_node,
-          check_nodes_number,
-          uncordon_node,
-          wait_for_metric_upload,
-          cleanup_workload,
-          cleanup_node_pool,
+          pre_test,
+          test,
+          post_test,
       )

dags/tpu_observability/jobset_ttr_node_pool_resize.py

@@ -34,12 +34,18 @@
 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 Workload
+from dags.common.task_group_with_timeout import TaskGroupWithTimeout
+
 
 DAG_ID = "jobset_ttr_node_pool_resize"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
 SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
 _DISK_SIZE_INCREMENT = 100
 
+PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10)
+POST_TEST_TIMEOUT = datetime.timedelta(minutes=10)
+TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT
+
 # Keyword arguments are generated dynamically at runtime (pylint does not
 # know this signature).
 with models.DAG(  # pylint: disable=unexpected-keyword-arg
@@ -108,76 +114,101 @@
 
       jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix)
 
-      create_node_pool = node_pool.create.override(task_id="create_node_pool")(
-          node_pool=cluster_info,
-          node_pool_selector=selector,
-      )
-
-      startup = jobset.create_jobset_startup_tasks(
-          node_pool=cluster_info,
-          jobset_config=jobset_config,
-          jobset_name=jobset_name,
-          node_pool_selector=selector,
-          workload_type=Workload.JAX_TPU_BENCHMARK,
-      )
-
-      node_pool_resize_start_time = node_pool.update.override(
-          task_id="node_pool_resize"
-      )(
-          node_pool=cluster_info,
-          spec=node_pool.NodePoolUpdateSpec.DiskSize(
-              delta=_DISK_SIZE_INCREMENT
-          ),
-      )
-
-      wait_for_recovery = jobset.wait_for_jobset_recovered.override(
-          task_id="wait_for_recovery"
-      )(
-          node_pool=cluster_info,
-          jobset_config=jobset_config,
-          jobset_name=jobset_name,
-      )
-
-      verify_duration = jobset.verify_recovery_duration.override(
-          task_id="verify_recovery_duration"
-      )(
-          start_time=node_pool_resize_start_time,
-          end_time=wait_for_recovery,
-      )
-
-      wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override(
-          task_id="wait_for_jobset_ttr_to_be_found",
-      )(
-          node_pool=cluster_info,
-          jobset_name=jobset_name,
-          start_time=node_pool_resize_start_time,
-      )
-
-      cleanup_workload = jobset.end_workload.override(
-          task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
-      )(
-          node_pool=cluster_info,
-          jobset_config=jobset_config,
-          jobset_name=jobset_name,
-      ).as_teardown(
-          setups=startup.jobset_start_time
-      )
-
-      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,
-      )
+      with TaskGroupWithTimeout(
+          group_id="pre_test",
+          timeout=PRE_TEST_TIMEOUT,
+      ) as pre_test:
+        create_node_pool = node_pool.create.override(
+            task_id="create_node_pool"
+        )(
+            node_pool=cluster_info,
+            node_pool_selector=selector,
+        )
+
+      with TaskGroupWithTimeout(
+          group_id="test",
+          timeout=TEST_TIMEOUT,
+      ) as test:
+        startup = jobset.create_jobset_startup_tasks(
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+            jobset_name=jobset_name,
+            node_pool_selector=selector,
+            workload_type=Workload.JAX_TPU_BENCHMARK,
+        )
+
+        node_pool_resize_start_time = node_pool.update.override(
+            task_id="node_pool_resize"
+        )(
+            node_pool=cluster_info,
+            spec=node_pool.NodePoolUpdateSpec.DiskSize(
+                delta=_DISK_SIZE_INCREMENT
+            ),
+        )
+
+        wait_for_recovery = jobset.wait_for_jobset_recovered.override(
+            task_id="wait_for_recovery"
+        )(
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+            jobset_name=jobset_name,
+        )
+
+        verify_duration = jobset.verify_recovery_duration.override(
+            task_id="verify_recovery_duration"
+        )(
+            start_time=node_pool_resize_start_time,
+            end_time=wait_for_recovery,
+        )
+
+        wait_for_metric_upload = (
+            jobset.wait_for_jobset_ttr_to_be_found.override(
+                task_id="wait_for_jobset_ttr_to_be_found",
+            )(
+                node_pool=cluster_info,
+                jobset_name=jobset_name,
+                start_time=node_pool_resize_start_time,
+            )
+        )
+
+        chain(
+            *startup.tasks,
+            node_pool_resize_start_time,
+            wait_for_recovery,
+            verify_duration,
+            wait_for_metric_upload,
+        )
+
+      with TaskGroupWithTimeout(
+          group_id="post_test",
+          timeout=POST_TEST_TIMEOUT,
+          is_teardown=True,
+      ) as post_test:
+        cleanup_workload = jobset.end_workload.override(
+            task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
+        )(
+            node_pool=cluster_info,
+            jobset_config=jobset_config,
+            jobset_name=jobset_name,
+        ).as_teardown(
+            setups=startup.jobset_start_time
+        )
+
+        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,
+        )
+
+        chain(
+            cleanup_workload,
+            cleanup_node_pool,
+        )
 
       chain(
           selector,
           jobset_name,
-          create_node_pool,
-          *startup.tasks,
-          node_pool_resize_start_time,
-          wait_for_recovery,
-          verify_duration,
-          wait_for_metric_upload,
-          cleanup_workload,
-          cleanup_node_pool,
+          pre_test,
+          test,
+          post_test,
       )