Skip to content

Commit d37486e

Browse files
chiajunglienchiajunglien
committed
feat(tpu): add JobSet healthiness validation DAG and enhance utils
1 parent a529792 commit d37486e

3 files changed

Lines changed: 339 additions & 0 deletions

File tree

‎dags/common/scheduling_helper/scheduling_helper.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ class DayOfWeek(enum.Enum):
6363
"jobset_ttr_kill_process": dt.timedelta(minutes=90),
6464
"jobset_uptime_validation": dt.timedelta(minutes=90),
6565
"jobset_ttr_drain_restart": DefaultTimeout,
66+
"jobset_healthiness_validation": dt.timedelta(minutes=90),
6667
"tpu_info_metrics_verification": DefaultTimeout,
6768
"jobset_ttr_node_reboot": dt.timedelta(minutes=90),
6869
},

‎dags/tpu_observability/jobset_healthiness_validation.py‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
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 Suspended Healthiness" metric."""
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+
from airflow.models.baseoperator import chain
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 Workload, ReplicatedJobStatus
28+
from dags.tpu_observability.configs.common import (
29+
MachineConfigMap,
30+
GCS_CONFIG_PATH,
31+
GCS_JOBSET_CONFIG_PATH,
32+
)
33+
from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
34+
35+
36+
DAG_ID = "jobset_healthiness_validation"
37+
DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
38+
SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
39+
40+
FAIL_WORKLOAD = "python3 -c 'import logging; import sys; logging.error(\"Simulating Failure\"); sys.exit(1)'"
41+
42+
# Keyword arguments are generated dynamically at runtime (pylint does not
43+
# know this signature).
44+
with models.DAG( # pylint: disable=unexpected-keyword-arg
45+
dag_id=DAG_ID,
46+
start_date=datetime.datetime(2025, 8, 10),
47+
schedule=SCHEDULE if composer_env.is_prod_env() else None,
48+
dagrun_timeout=DAGRUN_TIMEOUT,
49+
catchup=False,
50+
tags=[
51+
"cloud-ml-auto-solutions",
52+
"jobset",
53+
"healthiness",
54+
"tpu-obervability",
55+
"TPU",
56+
"v6e-16",
57+
],
58+
description=(
59+
"This DAG tests the 'Suspended' status of jobset healthiness by "
60+
"comparing the number of 'suspended' replicas before and after "
61+
"a jobset is running."
62+
),
63+
doc_md="""
64+
# JobSet Healthiness Test For the "Suspended" Status
65+
### Description
66+
This DAG automates the process of creating node-pools, ensuring the
67+
correct number of "Suspended" replicas appear, then launching a jobset on
68+
multiple replicas to ensure the correct number begin running.
69+
### Prerequisites
70+
This test requires an existing cluster to run.
71+
### Procedures
72+
First a node-pool is created. The validation test is then run to
73+
check if the number of "Suspended" replicas is 0. Once the jobset is
74+
running the jobs should quickly enter the "Ready" state. Then using
75+
command to suspend entire jobset. The number of found replicas is
76+
tested against the number of replicas which should be "Suspended".
77+
If they match the DAG is a success.
78+
""",
79+
) as dag:
80+
for machine in MachineConfigMap:
81+
config = machine.value
82+
83+
# Keyword arguments are generated dynamically at runtime (pylint does not
84+
# know this signature).
85+
with TaskGroup( # pylint: disable=unexpected-keyword-arg
86+
group_id=f"v{config.tpu_version.value}"
87+
):
88+
selector = jobset.generate_node_pool_selector(
89+
"jobset-healthiness-validation"
90+
)
91+
92+
jobset_config = jobset.build_jobset_from_gcs_yaml(
93+
gcs_path=GCS_JOBSET_CONFIG_PATH,
94+
dag_name=DAG_ID,
95+
node_pool_selector=selector,
96+
)
97+
98+
cluster_info = node_pool.build_node_pool_info_from_gcs_yaml.override(
99+
task_id="build_node_pool_info_from_gcs_yaml"
100+
)(
101+
gcs_path=GCS_CONFIG_PATH,
102+
dag_name=DAG_ID,
103+
is_prod=composer_env.is_prod_env(),
104+
machine_type=config.machine_version.value,
105+
tpu_topology=config.tpu_topology,
106+
node_pool_selector=selector,
107+
)
108+
109+
create_node_pool = node_pool.create.override(task_id="create_node_pool")(
110+
node_pool=cluster_info,
111+
)
112+
113+
startup = jobset.create_jobset_startup_group(
114+
node_pool=cluster_info,
115+
jobset_config=jobset_config,
116+
workload_type=Workload.JAX_TPU_BENCHMARK,
117+
)
118+
119+
with TaskGroup(group_id="validate_running_metrics") as validate_running:
120+
running_metrics = [
121+
(ReplicatedJobStatus.SPECIFIED, "USE_CONFIG_REPLICAS"),
122+
(ReplicatedJobStatus.ACTIVE, "USE_CONFIG_REPLICAS"),
123+
(ReplicatedJobStatus.READY, "USE_CONFIG_REPLICAS"),
124+
(ReplicatedJobStatus.FAILED, 0),
125+
(ReplicatedJobStatus.SUCCEEDED, 0),
126+
(ReplicatedJobStatus.SUSPENDED, 0),
127+
]
128+
for status, expected in running_metrics:
129+
jobset.wait_for_jobset_metrics.override(
130+
task_id=f"wait_{status.value}"
131+
)(
132+
metric_name=status,
133+
expected_value=expected,
134+
node_pool=cluster_info,
135+
jobset_config=jobset_config,
136+
)
137+
138+
suspend_action = jobset.suspended_jobset.override(
139+
task_id="suspend_jobset"
140+
)(
141+
node_pool=cluster_info,
142+
jobset_config=jobset_config,
143+
)
144+
145+
with TaskGroup(
146+
group_id="validate_suspended_metrics"
147+
) as validate_suspended:
148+
suspended_metrics = [
149+
(ReplicatedJobStatus.ACTIVE, 0),
150+
(ReplicatedJobStatus.SUSPENDED, "USE_CONFIG_REPLICAS"),
151+
]
152+
for status, expected in suspended_metrics:
153+
jobset.wait_for_jobset_metrics.override(
154+
task_id=f"wait_after_suspend_{status.value}"
155+
)(
156+
metric_name=status,
157+
expected_value=expected,
158+
node_pool=cluster_info,
159+
jobset_config=jobset_config,
160+
)
161+
162+
resume_action = jobset.resume_jobset.override(task_id="resume_jobset")(
163+
node_pool=cluster_info,
164+
jobset_config=jobset_config,
165+
)
166+
167+
with TaskGroup(group_id="inject_and_validate_failure") as failure_test:
168+
cleanup_for_failure = jobset.end_workload.override(
169+
task_id="cleanup_before_failure_injection"
170+
)(
171+
node_pool=cluster_info,
172+
jobset_config=jobset_config,
173+
)
174+
175+
start_fail_job = jobset.run_workload.override(task_id="start_fail_job")(
176+
node_pool=cluster_info,
177+
jobset_config=jobset_config,
178+
workload_type=FAIL_WORKLOAD,
179+
)
180+
181+
validate_failed_metric = jobset.wait_for_jobset_metrics.override(
182+
task_id="wait_for_failed_count"
183+
)(
184+
metric_name=ReplicatedJobStatus.FAILED,
185+
expected_value="USE_CONFIG_REPLICAS",
186+
node_pool=cluster_info,
187+
jobset_config=jobset_config,
188+
)
189+
190+
chain(cleanup_for_failure, start_fail_job, validate_failed_metric)
191+
192+
cleanup_workload = jobset.end_workload.override(
193+
task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
194+
)(
195+
node_pool=cluster_info,
196+
jobset_config=jobset_config,
197+
).as_teardown(
198+
setups=startup.jobset_start_time
199+
)
200+
201+
cleanup_node_pool = node_pool.delete.override(
202+
task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE
203+
)(node_pool=cluster_info).as_teardown(
204+
setups=create_node_pool,
205+
)
206+
207+
chain(
208+
selector,
209+
jobset_config,
210+
cluster_info,
211+
create_node_pool,
212+
startup.task_group,
213+
validate_running,
214+
suspend_action,
215+
validate_suspended,
216+
resume_action,
217+
failure_test,
218+
cleanup_workload,
219+
cleanup_node_pool,
220+
)

‎dags/tpu_observability/utils/jobset_util.py‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -326,6 +326,17 @@ class JobSetStartupOutput:
326326
jobset_start_time: XComArg
327327

328328

329+
class ReplicatedJobStatus(enum.Enum):
330+
"""Defines status of a replicated job."""
331+
332+
READY = "ready"
333+
SUSPENDED = "suspended"
334+
ACTIVE = "active"
335+
FAILED = "failed"
336+
SUCCEEDED = "succeeded"
337+
SPECIFIED = "specified"
338+
339+
329340
class Command:
330341
"""
331342
A collection of static methods to generate Kubernetes and gcloud commands.
@@ -411,6 +422,26 @@ def k8s_get_pod_name_command(
411422
"""Alias for getting just the names, maintaining existing API."""
412423
return Command.k8s_get_pods(jobset_name, namespace, output)
413424

425+
@staticmethod
426+
def k8s_suspend_jobset_command(
427+
kubeconfig: str, jobset_name: str, namespace: str
428+
) -> str:
429+
return (
430+
f"kubectl --kubeconfig={kubeconfig} "
431+
f"patch jobset {jobset_name} -n {namespace} "
432+
'--type=merge -p \'{"spec": {"suspend": true}}\''
433+
)
434+
435+
@staticmethod
436+
def k8s_resume_jobset_command(
437+
kubeconfig: str, jobset_name: str, namespace: str
438+
) -> str:
439+
return (
440+
f"kubectl --kubeconfig={kubeconfig} "
441+
f"patch jobset {jobset_name} -n {namespace} "
442+
'--type=merge -p \'{"spec": {"suspend": false}}\''
443+
)
444+
414445

415446
def get_replica_num(
416447
replica_type: str, job_name: str, node_pool: node_pool_info
@@ -866,6 +897,93 @@ def wait_for_jobset_started(
866897
return all(p > threshold_value for p in last_n_data_points)
867898

868899

900+
@task.sensor(poke_interval=60, timeout=3600, mode="poke")
901+
def wait_for_jobset_metrics(
902+
metric_name: ReplicatedJobStatus,
903+
expected_value: any,
904+
node_pool: node_pool_info,
905+
jobset_config: JobSet,
906+
start_time: TimeUtil = None,
907+
) -> bool:
908+
"""Polls Cloud Monitoring for a specific JobSet replicated job metric.
909+
910+
A sensor task which polls various JobSet metrics (e.g., ready, active,
911+
failed) every 60 seconds for 60 minutes. This sensor handles dynamic
912+
scaling expectations and ensures compatibility with different data
913+
types returned by the Prometheus exporter.
914+
915+
Args:
916+
metric_name (ReplicatedJobStatus): The status of the replicated job to
917+
monitor.
918+
expected_value (any): The target value to wait for. If set to
919+
"USE_CONFIG_REPLICAS", the sensor will dynamically retrieve the
920+
expected count from the jobset_config.
921+
node_pool (Info): An instance of the Info class containing GKE cluster
922+
metadata.
923+
jobset_config (JobSet): An instance of the JobSet class representing
924+
the target jobset configuration.
925+
start_time (TimeUtil, optional): The UTC timestamp to start polling
926+
from. If not provided, defaults to 60 minutes before the current time.
927+
928+
Returns:
929+
bool: True if the current metric value matches the expected value,
930+
False otherwise.
931+
"""
932+
final_expected = expected_value
933+
if expected_value == "USE_CONFIG_REPLICAS":
934+
final_expected = jobset_config.replicas
935+
936+
metric_mapping = {
937+
"ready": "prometheus.googleapis.com/kube_jobset_ready_replicas/gauge",
938+
"active": "prometheus.googleapis.com/kube_jobset_active_replicas/gauge",
939+
"suspended": "prometheus.googleapis.com/kube_jobset_suspended_replicas/gauge",
940+
"succeeded": "prometheus.googleapis.com/kube_jobset_succeeded_replicas/gauge",
941+
"failed": "prometheus.googleapis.com/kube_jobset_failed_replicas/gauge",
942+
"specified": "prometheus.googleapis.com/kube_jobset_specified_replicas/gauge",
943+
}
944+
945+
name_str = (
946+
metric_name.value
947+
if hasattr(metric_name, "value")
948+
else str(metric_name).lower()
949+
)
950+
metric_type = metric_mapping.get(name_str)
951+
query_start = (
952+
start_time if start_time else TimeUtil.now() - timedelta(minutes=60)
953+
)
954+
955+
time_series = list_time_series(
956+
project_id=node_pool.project_id,
957+
filter_str=(
958+
f'metric.type="{metric_type}" '
959+
f'resource.type="prometheus_target" '
960+
f'resource.labels.cluster="{node_pool.cluster_name}" '
961+
f'metric.labels.jobset_name="{jobset_config.jobset_name}"'
962+
),
963+
start_time=query_start,
964+
end_time=TimeUtil.now(),
965+
)
966+
967+
if not time_series or len(time_series) == 0 or not time_series[0].points:
968+
return False
969+
970+
point_value = time_series[0].points[0].value
971+
if (
972+
hasattr(point_value, "double_value")
973+
and point_value.double_value is not None
974+
):
975+
latest_value = point_value.double_value
976+
else:
977+
latest_value = float(point_value.int64_value)
978+
979+
logging.info(
980+
f"Metric {name_str} for JobSet {jobset_config.jobset_name}: "
981+
f"current={latest_value}, expected={final_expected}"
982+
)
983+
984+
return float(latest_value) == float(final_expected)
985+
986+
869987
@task.sensor(poke_interval=60, timeout=3600, mode="poke")
870988
def wait_for_jobset_ttr_to_be_found(
871989
node_pool: node_pool_info,

0 commit comments

Comments
 (0)