Skip to content

Commit 99acfc4

Browse files
chiajunglienchiajunglien
committed
refactor: encapsulate JobSet metric mapping into JobSetHealthiness enum
1 parent 234847f commit 99acfc4

2 files changed

Lines changed: 116 additions & 29 deletions

File tree

‎dags/tpu_observability/jobset_healthiness_validation.py‎

Lines changed: 29 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -88,42 +88,40 @@
8888
with TaskGroup( # pylint: disable=unexpected-keyword-arg
8989
group_id=f"v{config.tpu_version.value}"
9090
):
91-
selector = jobset.generate_node_pool_selector(
92-
"jobset-healthiness-validation"
91+
cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(
92+
gcs_path=GCS_CONFIG_PATH,
93+
dag_name=DAG_ID,
94+
is_prod=composer_env.is_prod_env(),
95+
machine_type=config.machine_version.value,
96+
tpu_topology=config.tpu_topology,
9397
)
9498

9599
jobset_config = jobset.build_jobset_from_gcs_yaml(
96100
gcs_path=GCS_JOBSET_CONFIG_PATH,
97101
dag_name=DAG_ID,
98-
node_pool_selector=selector,
99102
)
100103

101-
cluster_info = node_pool.build_node_pool_info_from_gcs_yaml.override(
102-
task_id="build_node_pool_info_from_gcs_yaml"
103-
)(
104-
gcs_path=GCS_CONFIG_PATH,
105-
dag_name=DAG_ID,
106-
is_prod=composer_env.is_prod_env(),
107-
machine_type=config.machine_version.value,
108-
tpu_topology=config.tpu_topology,
109-
node_pool_selector=selector,
110-
)
104+
selector = jobset.generate_node_pool_selector(DAG_ID)
105+
jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix)
111106

112107
create_node_pool = node_pool.create.override(task_id="create_node_pool")(
113108
node_pool=cluster_info,
109+
node_pool_selector=selector,
114110
)
115111

116112
startup = jobset.create_jobset_startup_tasks(
117113
node_pool=cluster_info,
118114
jobset_config=jobset_config,
115+
jobset_name=jobset_name,
116+
node_pool_selector=selector,
119117
workload_type=Workload.JAX_TPU_BENCHMARK,
120118
)
121119

122120
with TaskGroup(group_id="validate_running_metrics") as validate_running:
123121
running_metrics = [
124-
(JobSetHealthiness.SPECIFIED, "USE_CONFIG_REPLICAS"),
125-
(JobSetHealthiness.ACTIVE, "USE_CONFIG_REPLICAS"),
126-
(JobSetHealthiness.READY, "USE_CONFIG_REPLICAS"),
122+
(JobSetHealthiness.SPECIFIED, jobset_config.replicas),
123+
(JobSetHealthiness.ACTIVE, jobset_config.replicas),
124+
(JobSetHealthiness.READY, jobset_config.replicas),
127125
(JobSetHealthiness.FAILED, 0),
128126
(JobSetHealthiness.SUCCEEDED, 0),
129127
(JobSetHealthiness.SUSPENDED, 0),
@@ -135,22 +133,23 @@
135133
metric_name=status,
136134
expected_value=expected,
137135
node_pool=cluster_info,
138-
jobset_config=jobset_config,
136+
jobset_name=jobset_name,
139137
)
140138

141139
suspend_action = jobset.suspended_jobset.override(
142140
task_id="suspend_jobset"
143141
)(
144142
node_pool=cluster_info,
145143
jobset_config=jobset_config,
144+
jobset_name=jobset_name,
146145
)
147146

148147
with TaskGroup(
149148
group_id="validate_suspended_metrics"
150149
) as validate_suspended:
151150
suspended_metrics = [
152151
(JobSetHealthiness.ACTIVE, 0),
153-
(JobSetHealthiness.SUSPENDED, "USE_CONFIG_REPLICAS"),
152+
(JobSetHealthiness.SUSPENDED, jobset_config.replicas),
154153
]
155154
for status, expected in suspended_metrics:
156155
jobset.wait_for_jobset_metrics.override(
@@ -159,12 +158,13 @@
159158
metric_name=status,
160159
expected_value=expected,
161160
node_pool=cluster_info,
162-
jobset_config=jobset_config,
161+
jobset_name=jobset_name,
163162
)
164163

165164
resume_action = jobset.resume_jobset.override(task_id="resume_jobset")(
166165
node_pool=cluster_info,
167166
jobset_config=jobset_config,
167+
jobset_name=jobset_name,
168168
)
169169

170170
with TaskGroup(group_id="inject_and_validate_success") as success_test:
@@ -173,23 +173,25 @@
173173
)(
174174
node_pool=cluster_info,
175175
jobset_config=jobset_config,
176+
jobset_name=jobset_name,
176177
)
177178

178179
start_success_job = jobset.run_workload.override(
179180
task_id="start_success_job"
180181
)(
181182
node_pool=cluster_info,
182183
jobset_config=jobset_config,
184+
jobset_name=jobset_name,
183185
workload_type=SUCCESS_WORKLOAD,
184186
)
185187

186188
validate_succeeded_metric = jobset.wait_for_jobset_metrics.override(
187189
task_id="wait_for_succeeded_count"
188190
)(
189191
metric_name=JobSetHealthiness.SUCCEEDED,
190-
expected_value="USE_CONFIG_REPLICAS",
192+
expected_value=jobset_config.replicas,
191193
node_pool=cluster_info,
192-
jobset_config=jobset_config,
194+
jobset_name=jobset_name,
193195
)
194196

195197
chain(cleanup_for_success, start_success_job, validate_succeeded_metric)
@@ -200,21 +202,23 @@
200202
)(
201203
node_pool=cluster_info,
202204
jobset_config=jobset_config,
205+
jobset_name=jobset_name,
203206
)
204207

205208
start_fail_job = jobset.run_workload.override(task_id="start_fail_job")(
206209
node_pool=cluster_info,
207210
jobset_config=jobset_config,
211+
jobset_name=jobset_name,
208212
workload_type=FAIL_WORKLOAD,
209213
)
210214

211215
validate_failed_metric = jobset.wait_for_jobset_metrics.override(
212216
task_id="wait_for_failed_count"
213217
)(
214218
metric_name=JobSetHealthiness.FAILED,
215-
expected_value="USE_CONFIG_REPLICAS",
219+
expected_value=jobset_config.replicas,
216220
node_pool=cluster_info,
217-
jobset_config=jobset_config,
221+
jobset_name=jobset_name,
218222
)
219223

220224
chain(cleanup_for_failure, start_fail_job, validate_failed_metric)
@@ -224,6 +228,7 @@
224228
)(
225229
node_pool=cluster_info,
226230
jobset_config=jobset_config,
231+
jobset_name=jobset_name,
227232
).as_teardown(
228233
setups=startup.jobset_start_time
229234
)
@@ -236,8 +241,7 @@
236241

237242
chain(
238243
selector,
239-
jobset_config,
240-
cluster_info,
244+
jobset_name,
241245
create_node_pool,
242246
*startup.tasks,
243247
validate_running,

‎dags/tpu_observability/utils/jobset_util.py‎

Lines changed: 87 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -380,6 +380,18 @@ class JobSetHealthiness(enum.Enum):
380380
SUCCEEDED = "succeeded"
381381
SPECIFIED = "specified"
382382

383+
@property
384+
def healthiness_metric_type(self) -> str:
385+
mapping = {
386+
"ready": "prometheus.googleapis.com/kube_jobset_ready_replicas/gauge",
387+
"active": "prometheus.googleapis.com/kube_jobset_active_replicas/gauge",
388+
"suspended": "prometheus.googleapis.com/kube_jobset_suspended_replicas/gauge",
389+
"succeeded": "prometheus.googleapis.com/kube_jobset_succeeded_replicas/gauge",
390+
"failed": "prometheus.googleapis.com/kube_jobset_failed_replicas/gauge",
391+
"specified": "prometheus.googleapis.com/kube_jobset_specified_replicas/gauge",
392+
}
393+
return mapping.get(self.value)
394+
383395

384396
class Command:
385397
"""
@@ -914,7 +926,9 @@ def operate_pod(
914926

915927

916928
@task
917-
def suspended_jobset(node_pool: node_pool_info, jobset_config: JobSet):
929+
def suspended_jobset(
930+
node_pool: node_pool_info, jobset_config: JobSet, jobset_name: str
931+
):
918932
"""
919933
Suspend a jobset from the GKE cluster.
920934
@@ -935,7 +949,7 @@ def suspended_jobset(node_pool: node_pool_info, jobset_config: JobSet):
935949
Command.get_credentials_command(node_pool),
936950
Command.k8s_suspend_jobset_command(
937951
temp_config_file.name,
938-
jobset_config.jobset_name,
952+
jobset_name,
939953
jobset_config.namespace,
940954
),
941955
])
@@ -944,7 +958,9 @@ def suspended_jobset(node_pool: node_pool_info, jobset_config: JobSet):
944958

945959

946960
@task
947-
def resume_jobset(node_pool: node_pool_info, jobset_config: JobSet):
961+
def resume_jobset(
962+
node_pool: node_pool_info, jobset_config: JobSet, jobset_name: str
963+
):
948964
"""
949965
Resume a jobset from the GKE cluster.
950966
@@ -965,7 +981,7 @@ def resume_jobset(node_pool: node_pool_info, jobset_config: JobSet):
965981
Command.get_credentials_command(node_pool),
966982
Command.k8s_resume_jobset_command(
967983
temp_config_file.name,
968-
jobset_config.jobset_name,
984+
jobset_name,
969985
jobset_config.namespace,
970986
),
971987
])
@@ -1055,6 +1071,73 @@ def verify_recovery_duration(start_time: TimeUtil, end_time: TimeUtil):
10551071
"The 'jobset_time_to_recover' metric requires > 60s to be recorded. "
10561072
"Failing fast to avoid waiting for a missing metric."
10571073
)
1074+
@task.sensor(poke_interval=60, timeout=3600, mode="poke")
1075+
def wait_for_jobset_metrics(
1076+
metric_name: JobSetHealthiness,
1077+
expected_value: int,
1078+
node_pool: node_pool_info,
1079+
jobset_name: str,
1080+
start_time: TimeUtil = None,
1081+
) -> bool:
1082+
"""Polls Cloud Monitoring for a specific JobSet replicated job metric.
1083+
1084+
A sensor task which polls various JobSet metrics (e.g., ready, active,
1085+
failed) every 60 seconds for 60 minutes. This sensor handles dynamic
1086+
scaling expectations and ensures compatibility with different data
1087+
types returned by the Prometheus exporter.
1088+
1089+
Args:
1090+
metric_name (JobSetHealthiness): The status of the replicated job to
1091+
monitor.
1092+
expected_value (int): The target value to wait for.
1093+
node_pool (Info): An instance of the Info class containing GKE cluster
1094+
metadata.
1095+
jobset_name: The name of the JobSet.
1096+
start_time (TimeUtil, optional): The UTC timestamp to start polling
1097+
from. If not provided, defaults to 60 minutes before the current time.
1098+
1099+
Returns:
1100+
bool: True if the current metric value matches the expected value,
1101+
False otherwise.
1102+
"""
1103+
1104+
metric_type = metric_name.healthiness_metric_type
1105+
name_str = metric_name.value
1106+
1107+
query_start = (
1108+
start_time if start_time else TimeUtil.now() - timedelta(minutes=60)
1109+
)
1110+
1111+
time_series = list_time_series(
1112+
project_id=node_pool.project_id,
1113+
filter_str=(
1114+
f'metric.type="{metric_type}" '
1115+
f'resource.type="prometheus_target" '
1116+
f'resource.labels.cluster="{node_pool.cluster_name}" '
1117+
f'metric.labels.jobset_name="{jobset_name}"'
1118+
),
1119+
start_time=query_start,
1120+
end_time=TimeUtil.now(),
1121+
)
1122+
1123+
if not time_series or len(time_series) == 0 or not time_series[0].points:
1124+
return False
1125+
1126+
point_value = time_series[0].points[0].value
1127+
if (
1128+
hasattr(point_value, "double_value")
1129+
and point_value.double_value is not None
1130+
):
1131+
latest_value = point_value.double_value
1132+
else:
1133+
latest_value = float(point_value.int64_value)
1134+
1135+
logging.info(
1136+
f"Metric {name_str} for JobSet {jobset_name}: "
1137+
f"current={latest_value}, expected={expected_value}"
1138+
)
1139+
1140+
return float(latest_value) == float(expected_value)
10581141

10591142

10601143
@task.sensor(poke_interval=60, timeout=3600, mode="poke", retries=0)

0 commit comments

Comments
 (0)