Skip to content

Commit bc9fad7

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 Fix format/linting Linting fix Delete dags/dags/tpu_observability directory Delete dags/airflow_monitoring.py
1 parent a29b9ea commit bc9fad7

2 files changed

Lines changed: 214 additions & 10 deletions

File tree

‎dags/tpu_observability/jobset_healthiness_ready.py‎

Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,201 @@
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+
# pylint: enable=pointless-statement

‎dags/tpu_observability/utils/jobset_util.py‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -559,26 +559,29 @@ def wait_for_jobset_ttr_to_be_found(node_pool: node_pool_info) -> bool:
559559
return len(time_series) > 0
560560

561561

562-
@task.sensor(poke_interval=30, timeout=600, mode="reschedule")
563-
def wait_for_jobset_status_occurrence(
564-
replica_type: str, job_name: str, node_pool: node_pool_info
562+
@task.sensor(poke_interval=30, timeout=900, mode="reschedule")
563+
def validate_jobset_replica_number(
564+
node_pool: node_pool_info,
565+
jobset_config: JobSet,
566+
replica_type: str,
567+
correct_replica_num: int,
565568
):
566569
"""
567-
A sensor which checks if are any jobset replicas in a status type.
570+
A sensor which checks if the correct number jobset replicas in a status type.
568571
569572
Args:
570-
replica_type(str): The type of status being checked for.
571-
job_name(str): The name of the job replica which is run from the jobset.
572-
node_pool(Info): The Info object containing the cluster information needed
573-
for the kubernetes API to connect to it.
573+
node_pool: Configuration object with cluster details.
574+
jobset_config: Configuration of JobSet which is being run.
575+
replica_type(str): The name of the type of status being checked for.
576+
correct_replica_num(int): The expected number of replicas to be found.
574577
"""
575578
logging.info("Checking for number of replicas of type: %s", replica_type)
576579
ready_replicas = get_replica_num(
577580
replica_type=replica_type,
578-
job_name=job_name,
581+
job_name=jobset_config.replicated_job_name,
579582
node_pool=node_pool,
580583
)
581-
return ready_replicas > 0
584+
return ready_replicas == correct_replica_num
582585

583586

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

0 commit comments

Comments
 (0)