Skip to content

Commit b3e6f9c

Browse files
chiajunglienchiajunglien
authored
fix: Standardize node pool and jobset TTR validation flows in TPU-Observability related DAGs (GoogleCloudPlatform#1232)
The primary goal of this change is to unify the architectural design and implementation logic specifically across all TPU TTR (Time To Recovery) related DAGs. Previously, different TTR DAGs used inconsistent methods to verify workload status—some called `wait_for_jobset_started`, some used `wait_for_jobset_ttr_to_be_found`, and some used both. This PR refactors these into a **standardized "Procedure of TTR"**, ensuring that all DAGs triggering workloads follow a consistent, unified function to verify that JobSets and have entered the proper state.
1 parent 9fa1c0d commit b3e6f9c

9 files changed

Lines changed: 121 additions & 136 deletions

‎dags/tpu_observability/jobset_ttr_kill_process.py‎

Lines changed: 4 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -150,31 +150,16 @@ def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
150150
node_pool=cluster_info,
151151
)
152152

153-
apply_time = jobset.run_workload.override(task_id="run_workload")(
153+
startup = jobset.create_jobset_startup_group(
154154
node_pool=cluster_info,
155155
jobset_config=jobset_config,
156156
workload_type=Workload.JAX_TPU_BENCHMARK,
157157
)
158158

159-
running_pods = jobset.wait_for_all_pods_running.override(
160-
task_id="ensure_all_pods_running"
161-
)(
162-
node_pool=cluster_info,
163-
jobset_config=jobset_config,
164-
)
165-
166-
wait_for_job_start = jobset.wait_for_jobset_started.override(
167-
task_id="wait_for_job_start"
168-
)(
169-
cluster_info,
170-
pod_name_list=running_pods,
171-
job_apply_time=apply_time,
172-
)
173-
174159
kill_tasks = (
175160
kill_tpu_pod_workload.override(task_id="kill_tpu_pod_workload")
176161
.partial(info=cluster_info)
177-
.expand(pod_name=running_pods)
162+
.expand(pod_name=startup.running_pods)
178163
)
179164

180165
wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override(
@@ -190,7 +175,7 @@ def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
190175
node_pool=cluster_info,
191176
jobset_config=jobset_config,
192177
).as_teardown(
193-
setups=apply_time
178+
setups=startup.jobset_start_time
194179
)
195180

196181
cleanup_node_pool = node_pool.delete.override(
@@ -204,9 +189,7 @@ def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
204189
jobset_config,
205190
cluster_info,
206191
create_node_pool,
207-
apply_time,
208-
running_pods,
209-
wait_for_job_start,
192+
startup.task_group,
210193
kill_tasks,
211194
wait_for_metric_upload,
212195
cleanup_workload,

‎dags/tpu_observability/jobset_ttr_node_pool_resize.py‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -114,19 +114,12 @@
114114
node_pool=cluster_info,
115115
)
116116

117-
start_workload = jobset.run_workload.override(task_id="start_workload")(
117+
startup = jobset.create_jobset_startup_group(
118118
node_pool=cluster_info,
119119
jobset_config=jobset_config,
120120
workload_type=Workload.JAX_TPU_BENCHMARK,
121121
)
122122

123-
ensure_all_pods_running = jobset.wait_for_all_pods_running.override(
124-
task_id="ensure_all_pods_running"
125-
)(
126-
node_pool=cluster_info,
127-
jobset_config=jobset_config,
128-
)
129-
130123
node_pool_resize = node_pool.update.override(task_id="node_pool_resize")(
131124
node_pool=cluster_info,
132125
spec=node_pool.NodePoolUpdateSpec.DiskSize(
@@ -147,7 +140,7 @@
147140
node_pool=cluster_info,
148141
jobset_config=jobset_config,
149142
).as_teardown(
150-
setups=start_workload
143+
setups=startup.jobset_start_time
151144
)
152145

153146
cleanup_node_pool = node_pool.delete.override(
@@ -161,8 +154,7 @@
161154
jobset_config,
162155
cluster_info,
163156
create_node_pool,
164-
start_workload,
165-
ensure_all_pods_running,
157+
startup.task_group,
166158
node_pool_resize,
167159
wait_for_metric_upload,
168160
cleanup_workload,

‎dags/tpu_observability/jobset_ttr_pod_delete.py‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -110,19 +110,12 @@
110110
node_pool=cluster_info,
111111
)
112112

113-
start_workload = jobset.run_workload.override(task_id="start_workload")(
113+
startup = jobset.create_jobset_startup_group(
114114
node_pool=cluster_info,
115115
jobset_config=jobset_config,
116116
workload_type=Workload.JAX_TPU_BENCHMARK,
117117
)
118118

119-
ensure_all_pods_running = jobset.wait_for_all_pods_running.override(
120-
task_id="ensure_all_pods_running"
121-
)(
122-
node_pool=cluster_info,
123-
jobset_config=jobset_config,
124-
)
125-
126119
delete_random_pod = jobset.delete_one_random_pod.override(
127120
task_id="delete_random_pod"
128121
)(
@@ -140,7 +133,7 @@
140133
cleanup_workload = jobset.end_workload.override(
141134
task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
142135
)(node_pool=cluster_info, jobset_config=jobset_config).as_teardown(
143-
setups=start_workload
136+
setups=startup.jobset_start_time
144137
)
145138

146139
cleanup_node_pool = node_pool.delete.override(
@@ -154,8 +147,7 @@
154147
jobset_config,
155148
cluster_info,
156149
create_node_pool,
157-
start_workload,
158-
ensure_all_pods_running,
150+
startup.task_group,
159151
delete_random_pod,
160152
wait_for_metric_upload,
161153
cleanup_workload,

‎dags/tpu_observability/jobset_ttr_rollback.py‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -111,19 +111,12 @@
111111
node_pool=cluster_info,
112112
)
113113

114-
start_workload = jobset.run_workload.override(task_id="start_workload")(
114+
startup = jobset.create_jobset_startup_group(
115115
node_pool=cluster_info,
116116
jobset_config=jobset_config,
117117
workload_type=Workload.JAX_TPU_BENCHMARK,
118118
)
119119

120-
ensure_all_pods_running = jobset.wait_for_all_pods_running.override(
121-
task_id="ensure_all_pods_running"
122-
)(
123-
node_pool=cluster_info,
124-
jobset_config=jobset_config,
125-
)
126-
127120
rollback_node_pool = node_pool.rollback.override(
128121
task_id="rollback_node_pool"
129122
)(node_pool=cluster_info)
@@ -141,7 +134,7 @@
141134
node_pool=cluster_info,
142135
jobset_config=jobset_config,
143136
).as_teardown(
144-
setups=start_workload
137+
setups=startup.jobset_start_time
145138
)
146139

147140
cleanup_node_pool = node_pool.delete.override(
@@ -155,8 +148,7 @@
155148
jobset_config,
156149
cluster_info,
157150
create_node_pool,
158-
start_workload,
159-
ensure_all_pods_running,
151+
startup.task_group,
160152
rollback_node_pool,
161153
wait_for_metric_upload,
162154
cleanup_workload,

‎dags/tpu_observability/jobset_uptime_validation.py‎

Lines changed: 6 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -122,27 +122,18 @@ def get_current_time() -> TimeUtil:
122122
node_pool=cluster_info,
123123
)
124124

125-
apply_time = jobset.run_workload.override(task_id="run_workload")(
125+
startup = jobset.create_jobset_startup_group(
126126
node_pool=cluster_info,
127127
jobset_config=jobset_config,
128128
workload_type=Workload.JAX_TPU_BENCHMARK,
129129
)
130130

131-
pod_names = jobset.list_pod_names.override(task_id="list_pod_names")(
132-
node_pool=cluster_info,
133-
jobset_config=jobset_config,
134-
)
135-
136-
wait_for_job_start = jobset.wait_for_jobset_started.override(
137-
task_id="wait_for_job_start"
138-
)(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time)
139-
140131
wait_for_jobset_uptime_data = jobset.wait_for_jobset_uptime_data.override(
141132
task_id="wait_for_jobset_uptime_data"
142133
)(
143134
node_pool=cluster_info,
144135
jobset_config=jobset_config,
145-
jobset_apply_time=apply_time,
136+
jobset_apply_time=startup.jobset_start_time,
146137
)
147138

148139
clean_up_workload = jobset.end_workload.override(
@@ -151,7 +142,7 @@ def get_current_time() -> TimeUtil:
151142
node_pool=cluster_info,
152143
jobset_config=jobset_config,
153144
).as_teardown(
154-
setups=apply_time
145+
setups=startup.jobset_start_time
155146
)
156147

157148
jobset_clear_time = get_current_time.override(
@@ -175,11 +166,11 @@ def get_current_time() -> TimeUtil:
175166
)
176167

177168
chain(
169+
selector,
170+
jobset_config,
178171
cluster_info,
179172
create_node_pool,
180-
apply_time,
181-
pod_names,
182-
wait_for_job_start,
173+
startup.task_group,
183174
wait_for_jobset_uptime_data,
184175
clean_up_workload,
185176
jobset_clear_time,

‎dags/tpu_observability/tpu_info_format_validation_dags.py‎

Lines changed: 8 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -297,7 +297,10 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]):
297297
with models.DAG( # pylint: disable=unexpected-keyword-arg
298298
dag_id=DAG_ID,
299299
start_date=datetime.datetime(2025, 8, 15),
300-
default_args={"retries": 0},
300+
default_args={
301+
"retries": 0,
302+
"owner": test_owner.YUNA_T,
303+
},
301304
schedule=SCHEDULE if composer_env.is_prod_env() else None,
302305
dagrun_timeout=DAGRUN_TIMEOUT,
303306
catchup=False,
@@ -398,33 +401,16 @@ def generate_second_node_pool_name(
398401
node_pool=cluster_info_2,
399402
)
400403

401-
apply_time = jobset.run_workload.override(
402-
owner=test_owner.YUNA_T, task_id="run_workload"
403-
)(
404+
startup = jobset.create_jobset_startup_group(
404405
node_pool=cluster_info,
405406
jobset_config=jobset_config,
406407
workload_type=Workload.JAX_TPU_BENCHMARK,
407408
)
408409

409-
running_pods = jobset.wait_for_all_pods_running.override(
410-
task_id="ensure_all_pods_running"
411-
)(
412-
node_pool=cluster_info,
413-
jobset_config=jobset_config,
414-
)
415-
416-
wait_for_job_start = jobset.wait_for_jobset_started.override(
417-
task_id="wait_for_job_start"
418-
)(
419-
cluster_info,
420-
pod_name_list=running_pods,
421-
job_apply_time=apply_time,
422-
)
423-
424410
outputs_of_tpu_info = (
425411
get_tpu_info_from_pod.override(task_id="get_tpu_info")
426412
.partial(info=cluster_info)
427-
.expand(pod_name=running_pods)
413+
.expand(pod_name=startup.running_pods)
428414
)
429415

430416
output_of_tpu_info = (
@@ -478,7 +464,7 @@ def generate_second_node_pool_name(
478464
node_pool=cluster_info,
479465
jobset_config=jobset_config,
480466
).as_teardown(
481-
setups=apply_time
467+
setups=startup.jobset_start_time
482468
)
483469

484470
# Keyword arguments are generated dynamically at runtime (pylint does not
@@ -522,9 +508,7 @@ def generate_second_node_pool_name(
522508
cluster_info,
523509
cluster_info_2,
524510
create_node_pool,
525-
apply_time,
526-
running_pods,
527-
wait_for_job_start,
511+
startup.task_group,
528512
outputs_of_tpu_info,
529513
output_of_tpu_info,
530514
verification_group,

‎dags/tpu_observability/tpu_info_metrics_dag.py‎

Lines changed: 11 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -327,25 +327,12 @@ def generate_second_node_pool_name(
327327

328328
_ = [create_first_node_pool, create_second_node_pool]
329329

330-
apply_time = jobset.run_workload(
330+
startup = jobset.create_jobset_startup_group(
331331
node_pool=cluster_info,
332332
jobset_config=jobset_config,
333333
workload_type=Workload.JAX_TPU_BENCHMARK,
334334
)
335335

336-
pod_names = jobset.list_pod_names.override(
337-
task_id="list_pod_names",
338-
retries=5,
339-
retry_delay=datetime.timedelta(seconds=10),
340-
)(
341-
node_pool=cluster_info,
342-
jobset_config=jobset_config,
343-
)
344-
345-
wait_for_job_start = jobset.wait_for_jobset_started.override(
346-
task_id="wait_for_job_start"
347-
)(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time)
348-
349336
verification_results = {}
350337
all_verification_groups = []
351338

@@ -362,7 +349,7 @@ def generate_second_node_pool_name(
362349
jobset_config=jobset_config,
363350
metric_name=strategy.tpu_info_metric_name,
364351
)
365-
.expand(pod_name=pod_names)
352+
.expand(pod_name=startup.running_pods)
366353
)
367354

368355
tpu_info_metric_output = (
@@ -377,10 +364,14 @@ def generate_second_node_pool_name(
377364
run_metric_verification.override(task_id="run_verification")
378365
.partial(
379366
node_pool=cluster_info,
380-
job_apply_time=apply_time,
367+
job_apply_time=startup.jobset_start_time,
381368
metric_strategy=strategy,
382369
)
383-
.expand(comparison_data=pod_names.zip(tpu_info_metric_output))
370+
.expand(
371+
comparison_data=startup.running_pods.zip(
372+
tpu_info_metric_output
373+
)
374+
)
384375
)
385376

386377
all_verification_groups.append(verification_group)
@@ -391,7 +382,7 @@ def generate_second_node_pool_name(
391382
task_id="summarize_results", trigger_rule=TriggerRule.ALL_DONE
392383
)(
393384
verification_results_dict=verification_results,
394-
active_pods=pod_names,
385+
active_pods=startup.running_pods,
395386
)
396387

397388
clean_up_workload = jobset.end_workload.override(
@@ -400,7 +391,7 @@ def generate_second_node_pool_name(
400391
node_pool=cluster_info,
401392
jobset_config=jobset_config,
402393
).as_teardown(
403-
setups=apply_time
394+
setups=startup.jobset_start_time
404395
)
405396

406397
with TaskGroup(group_id="cleanup_node_pool") as cleanup_node_pool:
@@ -426,9 +417,7 @@ def generate_second_node_pool_name(
426417
cluster_info,
427418
cluster_info_2,
428419
create_node_pool,
429-
apply_time,
430-
pod_names,
431-
wait_for_job_start,
420+
startup.task_group,
432421
all_verification_groups,
433422
summary,
434423
clean_up_workload,

0 commit comments

Comments
 (0)