move startup to test group

Author: yuna-tzeng

dags/tpu_observability/jobset_ttr_drain_restart.py

@@ -160,6 +160,10 @@ def check_nodes_number(
             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,
@@ -168,10 +172,6 @@ def check_nodes_number(
             workload_type=Workload.JAX_TPU_BENCHMARK,
         )
 
-      with TaskGroupWithTimeout(
-          group_id="test",
-          timeout=TEST_TIMEOUT,
-      ) as test:
         select_node = node_pool.draw_random_node.override(
             task_id="select_node"
         )(node_pool=cluster_info)
@@ -207,6 +207,7 @@ def check_nodes_number(
         )
 
         chain(
+            *startup.tasks,
             select_node,
             drained_node,
             check_nodes_number_task,

dags/tpu_observability/jobset_ttr_node_pool_resize.py

@@ -125,6 +125,10 @@
             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,
@@ -133,10 +137,6 @@
             workload_type=Workload.JAX_TPU_BENCHMARK,
         )
 
-      with TaskGroupWithTimeout(
-          group_id="test",
-          timeout=TEST_TIMEOUT,
-      ) as test:
         node_pool_resize_start_time = node_pool.update.override(
             task_id="node_pool_resize"
         )(
@@ -172,6 +172,7 @@
         )
 
         chain(
+            *startup.tasks,
             node_pool_resize_start_time,
             wait_for_recovery,
             verify_duration,

dags/tpu_observability/jobset_ttr_pod_delete.py

@@ -128,6 +128,10 @@
             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,
@@ -136,10 +140,6 @@
             workload_type=Workload.JAX_TPU_BENCHMARK,
         )
 
-      with TaskGroupWithTimeout(
-          group_id="test",
-          timeout=TEST_TIMEOUT,
-      ) as test:
         deletion_start_time = jobset.delete_one_random_pod.override(
             task_id="delete_random_pod"
         )(
@@ -174,6 +174,7 @@
         )
 
         chain(
+            *startup.tasks,
             deletion_start_time,
             wait_for_recovery,
             verify_duration,

dags/tpu_observability/jobset_ttr_rollback.py

@@ -129,6 +129,10 @@
             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,
@@ -137,10 +141,6 @@
             workload_type=Workload.JAX_TPU_BENCHMARK,
         )
 
-      with TaskGroupWithTimeout(
-          group_id="test",
-          timeout=TEST_TIMEOUT,
-      ) as test:
         rollback_node_pool = node_pool.rollback.override(
             task_id="rollback_node_pool"
         )(node_pool=cluster_info)
@@ -171,6 +171,7 @@
         )
 
         chain(
+            *startup.tasks,
             rollback_node_pool,
             wait_for_recovery,
             verify_duration,

dags/tpu_observability/jobset_uptime_validation.py

@@ -135,6 +135,10 @@ def get_current_time() -> TimeUtil:
             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,
@@ -143,10 +147,6 @@ def get_current_time() -> TimeUtil:
             workload_type=Workload.JAX_TPU_BENCHMARK,
         )
 
-      with TaskGroupWithTimeout(
-          group_id="test",
-          timeout=TEST_TIMEOUT,
-      ) as test:
         wait_for_jobset_uptime_data = (
             jobset.wait_for_jobset_uptime_data.override(
                 task_id="wait_for_jobset_uptime_data"
@@ -184,6 +184,7 @@ def get_current_time() -> TimeUtil:
         )
 
         chain(
+            *startup.tasks,
             wait_for_jobset_uptime_data,
             clean_up_workload,
             jobset_clear_time,