Skip to content

Commit b93a74c

Browse files
QuinnMMcGarryQuinnMMcGarry
committed
Adds a new DAG to test "Jobset Ready Healthiness" metric
- Adds an entire new DAG to test this metric automatically - Adds the "validate_replica_number" function - Adds a "READY_TPU" Jobset in jobset_util.py -Fix automatic cluster naming script
1 parent 0535667 commit b93a74c

2 files changed

Lines changed: 234 additions & 0 deletions

File tree

‎dags/tpu_observability/jobset_healthiness_ready.py‎

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
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 Ready Healthiness" metric."""
16+
17+
import datetime
18+
19+
from airflow.decorators import task
20+
from airflow import models
21+
from airflow.utils.trigger_rule import TriggerRule
22+
from airflow.utils.task_group import TaskGroup
23+
24+
from dags.common.vm_resource import Region, Zone
25+
from dags.tpu_observability.utils import jobset_util as jobset
26+
from dags.tpu_observability.utils import node_pool_util as node_pool
27+
from dags.tpu_observability.utils.jobset_util import JobSet, Workload
28+
from dags.tpu_observability.configs.common import MachineConfigMap
29+
from dags import composer_env
30+
31+
# Can be moved to jobset_util.py
32+
@task.sensor(poke_interval=30, timeout=900, mode="reschedule")
33+
def validate_replica_number(
34+
node_pool: node_pool.Info,
35+
jobset_config: JobSet,
36+
replica_type: str,
37+
replica_num: int
38+
):
39+
found_replia_num = jobset.get_replica_num(
40+
replica_type=replica_type,
41+
job_name=jobset_config.replicated_job_name,
42+
node_pool=node_pool
43+
)
44+
45+
return (found_replia_num == replica_num)
46+
47+
48+
# Keyword arguments are generated dynamically at runtime (pylint does not
49+
# know this signature).
50+
with models.DAG( # pylint: disable=unexpected-keyword-arg
51+
dag_id="jobset_healthiness_ready",
52+
start_date=datetime.datetime(2025, 8, 10),
53+
schedule="00 03 * * *",
54+
catchup=False,
55+
tags=[
56+
"cloud-ml-auto-solutions",
57+
"jobset",
58+
"healthiness",
59+
"tpu-obervability",
60+
"TPU",
61+
"v6e-16",
62+
],
63+
description=(
64+
"This DAG tests the 'Ready' status of jobset healthiness by "
65+
"comparing the number of 'Ready' replicas before and after "
66+
"a jobset is running."
67+
),
68+
doc_md="""
69+
# JobSet Healthiness Test For the "Ready" Status
70+
71+
### Description
72+
This DAG automates the process of creating node-pools, ensuring the
73+
correct number of "Ready" replicas appear, then launching a jobset on
74+
multiple replicas to ensure the correct number begin running.
75+
76+
### Prerequisites
77+
This test requires an existing cluster to run.
78+
79+
### Procedures
80+
First two node-pools are created. The validation test is then run to
81+
check if the number of "Ready" replicas is 0. A jobset is then launched
82+
which uses 2 replicas. Once the jobset is running the jobs should
83+
quickly enter the "Ready" state. The number of found replicas is
84+
tested against the number of replicas which should be "Ready". If they
85+
match the DAG is a success.
86+
""",
87+
) as dag:
88+
cluster_name = "tpu-observability-automation"
89+
cluster_name += "-prod" if composer_env.is_prod_env() else "-dev"
90+
91+
for machine in MachineConfigMap:
92+
config = machine.value
93+
cluster_info = node_pool.Info(
94+
project_id=models.Variable.get("PROJECT_ID", default_var="cienet-cmcs"),
95+
cluster_name=models.Variable.get(
96+
"CLUSTER_NAME", default_var=cluster_name
97+
),
98+
node_pool_name=models.Variable.get(
99+
"NODE_POOL_NAME", default_var="jobset-healthiness-ready-v6e"
100+
),
101+
region=models.Variable.get(
102+
"REGION", default_var=Region.US_CENTRAL1.value
103+
),
104+
location=models.Variable.get(
105+
"LOCATION", default_var=Region.US_CENTRAL1.value
106+
),
107+
node_locations=models.Variable.get(
108+
"LOCATIONS", default_var=Zone.US_CENTRAL1_B.value
109+
),
110+
num_nodes=models.Variable.get("NUM_NODES", default_var=4),
111+
machine_type=config.machine_version.value,
112+
tpu_topology=config.tpu_topology,
113+
)
114+
115+
cluster_info_2 = node_pool.Info(
116+
project_id=models.Variable.get("PROJECT_ID", default_var="cienet-cmcs"),
117+
cluster_name=models.Variable.get(
118+
"CLUSTER_NAME", default_var=cluster_name
119+
),
120+
node_pool_name=models.Variable.get(
121+
"NODE_POOL_NAME", default_var="jobset-healthiness-ready-v6e-2"
122+
),
123+
region=models.Variable.get(
124+
"REGION", default_var=Region.US_CENTRAL1.value
125+
),
126+
location=models.Variable.get(
127+
"LOCATION", default_var=Region.US_CENTRAL1.value
128+
),
129+
node_locations=models.Variable.get(
130+
"LOCATIONS", default_var=Zone.US_CENTRAL1_B.value
131+
),
132+
num_nodes=models.Variable.get("NUM_NODES", default_var=4),
133+
machine_type=config.machine_version.value,
134+
tpu_topology=config.tpu_topology,
135+
)
136+
137+
jobset_config = JobSet(
138+
jobset_name="jobset-healthiness-ready",
139+
namespace="default",
140+
max_restarts=0,
141+
replicated_job_name="tpu-job-slice",
142+
replicas=2,
143+
backoff_limit=0,
144+
completions=4,
145+
parallelism=4,
146+
tpu_accelerator_type="tpu-v6e-slice",
147+
tpu_topology="4x4",
148+
container_name="jax-tpu-worker",
149+
image="python:3.11",
150+
tpu_cores_per_pod=4,
151+
)
152+
153+
# Keyword arguments are generated dynamically at runtime (pylint does not
154+
# know this signature).
155+
with TaskGroup( # pylint: disable=unexpected-keyword-arg
156+
group_id=f"v{config.tpu_version.value}"
157+
):
158+
create_node_pool = node_pool.create(
159+
node_pool=cluster_info,
160+
reservation="cloudtpu-20251107233000-1246578561",
161+
)
162+
163+
create_node_pool_2 = node_pool.create(
164+
node_pool=cluster_info_2,
165+
reservation="cloudtpu-20251107233000-1246578561",
166+
)
167+
168+
validate_zero_replicas = validate_replica_number(
169+
node_pool=cluster_info,
170+
jobset_config=jobset_config,
171+
replica_type="ready",
172+
replica_num=0
173+
)
174+
175+
start_workload = jobset.run_workload(
176+
node_pool=cluster_info,
177+
yaml_config=jobset_config.generate_yaml(
178+
workload_script=Workload.READY_TPU
179+
),
180+
namespace=jobset_config.namespace,
181+
)
182+
183+
validate_ready_replicas = validate_replica_number(
184+
node_pool=cluster_info,
185+
jobset_config=jobset_config,
186+
replica_type="ready",
187+
replica_num=jobset_config.replicas
188+
)
189+
190+
cleanup_workload = jobset.end_workload.override(
191+
task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
192+
)(
193+
node_pool=cluster_info,
194+
jobset_name=jobset_config.jobset_name,
195+
namespace=jobset_config.namespace,
196+
).as_teardown(
197+
setups=start_workload
198+
)
199+
200+
cleanup_node_pool = node_pool.delete.override(
201+
task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE
202+
)(node_pool=cluster_info).as_teardown(
203+
setups=create_node_pool,
204+
)
205+
206+
cleanup_node_pool_2 = node_pool.delete.override(
207+
task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE
208+
)(node_pool=cluster_info_2).as_teardown(
209+
setups=create_node_pool_2,
210+
)
211+
212+
# Airflow uses >> for task chaining, which is pointless for pylint.
213+
# pylint: disable=pointless-statement
214+
(
215+
create_node_pool
216+
>> create_node_pool_2
217+
>> validate_zero_replicas
218+
>> start_workload
219+
>> validate_ready_replicas
220+
>> cleanup_workload
221+
>> cleanup_node_pool
222+
>> cleanup_node_pool_2
223+
224+
)
225+
# pylint: enable=pointless-statement

‎dags/tpu_observability/utils/jobset_util.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,15 @@ class Workload:
4343
Each workload is a JSON-escaped string, ready to be used as a shell argument.
4444
"""
4545

46+
READY_TPU = json.dumps(
47+
textwrap.dedent(
48+
"""
49+
sleep 1000
50+
"""
51+
),
52+
ensure_ascii=False
53+
)
54+
4655
JAX_TPU_BENCHMARK = json.dumps(
4756
textwrap.dedent(
4857
"""

0 commit comments

Comments
 (0)