Skip to content

Commit 7d71e4d

Browse files
QuinnMMcGarryQuinnMMcGarry
authored
Add a new DAG to test jobset time-to-recover from a node pool rollback (GoogleCloudPlatform#1077)
This change adds a new DAG which automates the process of checking if the jobset time-to-recover metric is properly updated from a node pool rollback
1 parent 894f3d9 commit 7d71e4d

2 files changed

Lines changed: 358 additions & 10 deletions

File tree

‎dags/tpu_observability/jobset_ttr_rollback.py‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""A DAG to test jobset time-to-recover metric using a node-pool rollback."""
16+
17+
import datetime
18+
19+
from airflow import models
20+
from airflow.utils.trigger_rule import TriggerRule
21+
from airflow.utils.task_group import TaskGroup
22+
23+
from dags.common.vm_resource import Region, Zone
24+
from dags.tpu_observability.utils import jobset_util as jobset
25+
from dags.tpu_observability.utils import node_pool_util as node_pool
26+
from dags.tpu_observability.utils.jobset_util import JobSet, Workload
27+
from dags.tpu_observability.configs.common import MachineConfigMap
28+
29+
30+
# Keyword arguments are generated dynamically at runtime (pylint does not
31+
# know this signature).
32+
with models.DAG( # pylint: disable=unexpected-keyword-arg
33+
dag_id="jobset_rollback_ttr",
34+
start_date=datetime.datetime(2025, 8, 10),
35+
schedule="00 02 * * *",
36+
catchup=False,
37+
tags=[
38+
"cloud-ml-auto-solutions",
39+
"jobset",
40+
"time-to-recover",
41+
"tpu-obervability",
42+
"rollback",
43+
"TPU",
44+
"v6e-16",
45+
],
46+
description=(
47+
"This DAG tests the use of a node-pool rollback to interrupt a "
48+
"jobset, then polls the jobset time-to-recover metric to check "
49+
"if it is updated."
50+
),
51+
doc_md="""
52+
# JobSet Time-To-Recover (TTR) Test Using Node-Pool Rollback
53+
54+
### Description
55+
This DAG automates the process of creating a node-pool, launching a jobset
56+
then using a node-pool rollback to interrupt the node-pool, and afterwards
57+
monitors if the jobset TTR metric gets updated. Finally the DAG cleans up
58+
the jobset and node-pool which were created.
59+
60+
### Prerequisites
61+
This test requires an existing cluster to run.
62+
63+
### Procedures
64+
First the node-pool is created, a jobset yaml is then launched on the
65+
cluster and given a short period of time to initialize. After this a
66+
rollback is run on the previously created node-pool to interrupt it.
67+
A sensor is finally run which will either detect that the jobset
68+
time-to-recover metric has been updated, resulting in a success, or
69+
timeout, and fail.
70+
""",
71+
) as dag:
72+
for machine in MachineConfigMap:
73+
config = machine.value
74+
cluster_info = node_pool.Info(
75+
project_id=models.Variable.get("PROJECT_ID", default_var="cienet-cmcs"),
76+
cluster_name=models.Variable.get(
77+
"CLUSTER_NAME", default_var="tpu-observability-automation"
78+
),
79+
node_pool_name=models.Variable.get(
80+
"NODE_POOL_NAME", default_var="jobset-ttr-rollback-v6e"
81+
),
82+
region=models.Variable.get(
83+
"REGION", default_var=Region.US_CENTRAL1.value
84+
),
85+
location=models.Variable.get(
86+
"LOCATION", default_var=Region.US_CENTRAL1.value
87+
),
88+
node_locations=models.Variable.get(
89+
"LOCATIONS", default_var=Zone.US_CENTRAL1_B.value
90+
),
91+
num_nodes=models.Variable.get("NUM_NODES", default_var=4),
92+
machine_type=config.machine_version.value,
93+
tpu_topology=config.tpu_topology,
94+
)
95+
96+
jobset_config = JobSet(
97+
jobset_name="ttr-rollback-v6e-workload",
98+
namespace="default",
99+
max_restarts=5,
100+
replicated_job_name="tpu-job-slice",
101+
replicas=1,
102+
backoff_limit=0,
103+
completions=4,
104+
parallelism=4,
105+
tpu_accelerator_type="tpu-v6e-slice",
106+
tpu_topology="4x4",
107+
container_name="jax-tpu-worker",
108+
image="python:3.11",
109+
tpu_cores_per_pod=4,
110+
)
111+
112+
# Keyword arguments are generated dynamically at runtime (pylint does not
113+
# know this signature).
114+
with TaskGroup( # pylint: disable=unexpected-keyword-arg
115+
group_id=f"v{config.tpu_version.value}"
116+
):
117+
create_node_pool = node_pool.create(
118+
node_pool=cluster_info,
119+
reservation="cloudtpu-20251107233000-1246578561",
120+
)
121+
122+
start_workload = jobset.run_workload(
123+
node_pool=cluster_info,
124+
yaml_config=jobset_config.generate_yaml(
125+
workload_script=Workload.JAX_TPU_BENCHMARK
126+
),
127+
namespace=jobset_config.namespace,
128+
)
129+
130+
ensure_all_pods_running = jobset.wait_for_all_pods_running(
131+
num_pods=(jobset_config.replicas * jobset_config.parallelism),
132+
node_pool=cluster_info,
133+
)
134+
135+
rollback_node_pool = node_pool.rollback(node_pool=cluster_info)
136+
137+
wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found(
138+
node_pool=cluster_info
139+
)
140+
141+
cleanup_workload = jobset.end_workload.override(
142+
task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
143+
)(
144+
node_pool=cluster_info,
145+
jobset_name=jobset_config.jobset_name,
146+
namespace=jobset_config.namespace,
147+
).as_teardown(
148+
setups=start_workload
149+
)
150+
151+
cleanup_node_pool = node_pool.delete.override(
152+
task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE
153+
)(node_pool=cluster_info).as_teardown(
154+
setups=create_node_pool,
155+
)
156+
157+
# Airflow uses >> for task chaining, which is pointless for pylint.
158+
# pylint: disable=pointless-statement
159+
(
160+
create_node_pool
161+
>> start_workload
162+
>> ensure_all_pods_running
163+
>> rollback_node_pool
164+
>> wait_for_metric_upload
165+
>> cleanup_workload
166+
>> cleanup_node_pool
167+
)
168+
# pylint: enable=pointless-statement

0 commit comments

Comments
 (0)