Skip to content

Commit 7da7f18

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 -Created "validate_jobset_replica_number" funtion in jobset_util.py -Changed workload name to "IDLE_READY_TPU_20M" -Changed schedule time and added statement to only run on PROD environment Changed to new nodepool creation format Remove unused DAG flow
1 parent 0535667 commit 7da7f18

2 files changed

Lines changed: 224 additions & 10 deletions

File tree

‎dags/tpu_observability/jobset_healthiness_ready.py‎

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,202 @@
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 import models
20+
from airflow.decorators import task
21+
from airflow.utils.trigger_rule import TriggerRule
22+
from airflow.utils.task_group import TaskGroup
23+
24+
from dags import composer_env
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, GCS_CONFIG_PATH
29+
30+
31+
# Keyword arguments are generated dynamically at runtime (pylint does not
32+
# know this signature).
33+
with models.DAG( # pylint: disable=unexpected-keyword-arg
34+
dag_id="jobset_healthiness_ready",
35+
start_date=datetime.datetime(2025, 8, 10),
36+
schedule="30 19 * * *" if composer_env.is_prod_env() else None,
37+
catchup=False,
38+
tags=[
39+
"cloud-ml-auto-solutions",
40+
"jobset",
41+
"healthiness",
42+
"tpu-obervability",
43+
"TPU",
44+
"v6e-16",
45+
],
46+
description=(
47+
"This DAG tests the 'Ready' status of jobset healthiness by "
48+
"comparing the number of 'Ready' replicas before and after "
49+
"a jobset is running."
50+
),
51+
doc_md="""
52+
# JobSet Healthiness Test For the "Ready" Status
53+
54+
### Description
55+
This DAG automates the process of creating node-pools, ensuring the
56+
correct number of "Ready" replicas appear, then launching a jobset on
57+
multiple replicas to ensure the correct number begin running.
58+
59+
### Prerequisites
60+
This test requires an existing cluster to run.
61+
62+
### Procedures
63+
First two node-pools are created. The validation test is then run to
64+
check if the number of "Ready" replicas is 0. A jobset is then launched
65+
which uses 2 replicas. Once the jobset is running the jobs should
66+
quickly enter the "Ready" state. The number of found replicas is
67+
tested against the number of replicas which should be "Ready". If they
68+
match the DAG is a success.
69+
""",
70+
) as dag:
71+
for machine in MachineConfigMap:
72+
config = machine.value
73+
74+
@task
75+
def generate_second_node_pool_name(
76+
node_pool_info: node_pool.Info,
77+
) -> str:
78+
"""Generates a second node pool name."""
79+
return f"{node_pool_info.node_pool_name}-2"
80+
81+
jobset_config = JobSet(
82+
jobset_name="jobset-healthiness-ready",
83+
namespace="default",
84+
max_restarts=0,
85+
replicated_job_name="tpu-job-slice",
86+
replicas=2,
87+
backoff_limit=0,
88+
completions=4,
89+
parallelism=4,
90+
tpu_accelerator_type="tpu-v6e-slice",
91+
tpu_topology="4x4",
92+
container_name="jax-tpu-worker",
93+
image="python:3.11",
94+
tpu_cores_per_pod=4,
95+
)
96+
97+
# Keyword arguments are generated dynamically at runtime (pylint does not
98+
# know this signature).
99+
with TaskGroup( # pylint: disable=unexpected-keyword-arg
100+
group_id=f"v{config.tpu_version.value}"
101+
):
102+
cluster_info = node_pool.build_node_pool_info_from_gcs_yaml.override(
103+
task_id="build_node_pool_info_from_gcs_yaml"
104+
)(
105+
gcs_path=GCS_CONFIG_PATH,
106+
dag_name="jobset_healthiness_ready",
107+
is_prod=composer_env.is_prod_env(),
108+
machine_type=config.machine_version.value,
109+
tpu_topology=config.tpu_topology,
110+
)
111+
112+
cluster_info_2 = node_pool.copy_node_pool_info_with_override(
113+
info=cluster_info,
114+
node_pool_name=generate_second_node_pool_name(cluster_info),
115+
)
116+
117+
# Keyword arguments are generated dynamically at runtime (pylint does not
118+
# know this signature).
119+
with TaskGroup( # pylint: disable=unexpected-keyword-arg
120+
group_id="create_node_pool"
121+
) as create_node_pool:
122+
create_first_node_pool = node_pool.create.override(
123+
task_id="node_pool_1",
124+
retries=2,
125+
)(
126+
node_pool=cluster_info,
127+
)
128+
129+
create_second_node_pool = node_pool.create.override(
130+
task_id="node_pool_2",
131+
retries=2,
132+
)(
133+
node_pool=cluster_info_2,
134+
)
135+
136+
validate_zero_replicas = jobset.validate_jobset_replica_number(
137+
node_pool=cluster_info,
138+
jobset_config=jobset_config,
139+
replica_type="ready",
140+
correct_replica_num=0,
141+
)
142+
143+
start_workload = jobset.run_workload(
144+
node_pool=cluster_info,
145+
yaml_config=jobset_config.generate_yaml(
146+
workload_script=Workload.IDLE_READY_TPU_20M
147+
),
148+
namespace=jobset_config.namespace,
149+
)
150+
151+
validate_ready_replicas = jobset.validate_jobset_replica_number(
152+
node_pool=cluster_info,
153+
jobset_config=jobset_config,
154+
replica_type="ready",
155+
correct_replica_num=jobset_config.replicas,
156+
)
157+
158+
cleanup_workload = jobset.end_workload.override(
159+
task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
160+
)(
161+
node_pool=cluster_info,
162+
jobset_name=jobset_config.jobset_name,
163+
namespace=jobset_config.namespace,
164+
).as_teardown(
165+
setups=start_workload
166+
)
167+
168+
# Keyword arguments are generated dynamically at runtime (pylint does not
169+
# know this signature).
170+
with TaskGroup( # pylint: disable=unexpected-keyword-arg
171+
group_id="cleanup_node_pool"
172+
) as cleanup_node_pool:
173+
cleanup_first_node_pool = node_pool.delete.override(
174+
task_id="cleanup_node_pool_1",
175+
trigger_rule=TriggerRule.ALL_DONE,
176+
retries=2,
177+
)(node_pool=cluster_info).as_teardown(
178+
setups=create_node_pool,
179+
)
180+
181+
cleanup_second_node_pool = node_pool.delete.override(
182+
task_id="cleanup_node_pool_2",
183+
trigger_rule=TriggerRule.ALL_DONE,
184+
retries=2,
185+
)(node_pool=cluster_info_2).as_teardown(
186+
setups=create_node_pool,
187+
)
188+
189+
# Airflow uses >> for task chaining, which is pointless for pylint.
190+
# pylint: disable=pointless-statement
191+
(
192+
cluster_info
193+
>> cluster_info_2
194+
>> create_node_pool
195+
>> validate_zero_replicas
196+
>> start_workload
197+
>> validate_ready_replicas
198+
>> cleanup_workload
199+
>> cleanup_node_pool
200+
201+
)
202+
# pylint: enable=pointless-statement

‎dags/tpu_observability/utils/jobset_util.py‎

Lines changed: 22 additions & 10 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+
IDLE_READY_TPU_20M = json.dumps(
47+
textwrap.dedent(
48+
"""
49+
sleep 1200
50+
"""
51+
),
52+
ensure_ascii=False,
53+
)
54+
4655
JAX_TPU_BENCHMARK = json.dumps(
4756
textwrap.dedent(
4857
"""
@@ -551,26 +560,29 @@ def wait_for_jobset_ttr_to_be_found(node_pool: node_pool.Info) -> bool:
551560
return len(time_series) > 0
552561

553562

554-
@task.sensor(poke_interval=30, timeout=600, mode="reschedule")
555-
def wait_for_jobset_status_occurrence(
556-
replica_type: str, job_name: str, node_pool: node_pool.Info
563+
@task.sensor(poke_interval=30, timeout=900, mode="reschedule")
564+
def validate_jobset_replica_number(
565+
node_pool: node_pool.Info,
566+
jobset_config: JobSet,
567+
replica_type: str,
568+
correct_replica_num: int,
557569
):
558570
"""
559-
A sensor which checks if are any jobset replicas in a status type.
571+
A sensor which checks if the correct number jobset replicas in a status type.
560572
561573
Args:
562-
replica_type(str): The type of status being checked for.
563-
job_name(str): The name of the job replica which is run from the jobset.
564-
node_pool(Info): The Info object containing the cluster information needed
565-
for the kubernetes API to connect to it.
574+
node_pool: Configuration object with cluster details.
575+
jobset_config: Configuration of JobSet which is being run.
576+
replica_type(str): The name of the type of status being checked for.
577+
correct_replica_num(int): The expected number of replicas to be found.
566578
"""
567579
logging.info("Checking for number of replicas of type: %s", replica_type)
568580
ready_replicas = get_replica_num(
569581
replica_type=replica_type,
570-
job_name=job_name,
582+
job_name=jobset_config.replicated_job_name,
571583
node_pool=node_pool,
572584
)
573-
return ready_replicas > 0
585+
return ready_replicas == correct_replica_num
574586

575587

576588
@task.sensor(poke_interval=30, timeout=600, mode="reschedule")

0 commit comments

Comments
 (0)