Skip to content

Commit c5f8abf

Browse files
chiajunglienchiajunglien
committed
fix: Apply TaskGroupWithTimeout into DAG.
1 parent 99acfc4 commit c5f8abf

2 files changed

Lines changed: 116 additions & 98 deletions

File tree

‎dags/tpu_observability/jobset_healthiness_validation.py‎

Lines changed: 99 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
"""A DAG to test "Jobset Suspended Healthiness" metric."""
1616

1717
import datetime
18+
from datetime import timedelta
1819

1920
from airflow import models
2021
from airflow.utils.trigger_rule import TriggerRule
@@ -24,6 +25,7 @@
2425
from dags import composer_env
2526
from dags.tpu_observability.utils import jobset_util as jobset
2627
from dags.tpu_observability.utils import node_pool_util as node_pool
28+
from dags.common.task_group_with_timeout import TaskGroupWithTimeout
2729
from dags.tpu_observability.utils.jobset_util import Workload, JobSetHealthiness
2830
from dags.tpu_observability.configs.common import (
2931
MachineConfigMap,
@@ -85,8 +87,9 @@
8587

8688
# Keyword arguments are generated dynamically at runtime (pylint does not
8789
# know this signature).
88-
with TaskGroup( # pylint: disable=unexpected-keyword-arg
89-
group_id=f"v{config.tpu_version.value}"
90+
with TaskGroupWithTimeout( # pylint: disable=unexpected-keyword-arg
91+
group_id=f"v{config.tpu_version.value}",
92+
timeout=timedelta(minutes=90),
9093
):
9194
cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(
9295
gcs_path=GCS_CONFIG_PATH,
@@ -117,24 +120,25 @@
117120
workload_type=Workload.JAX_TPU_BENCHMARK,
118121
)
119122

120-
with TaskGroup(group_id="validate_running_metrics") as validate_running:
121-
running_metrics = [
122-
(JobSetHealthiness.SPECIFIED, jobset_config.replicas),
123-
(JobSetHealthiness.ACTIVE, jobset_config.replicas),
124-
(JobSetHealthiness.READY, jobset_config.replicas),
125-
(JobSetHealthiness.FAILED, 0),
126-
(JobSetHealthiness.SUCCEEDED, 0),
127-
(JobSetHealthiness.SUSPENDED, 0),
128-
]
129-
for status, expected in running_metrics:
130-
jobset.wait_for_jobset_metrics.override(
131-
task_id=f"wait_{status.value}"
132-
)(
133-
metric_name=status,
134-
expected_value=expected,
135-
node_pool=cluster_info,
136-
jobset_name=jobset_name,
137-
)
123+
running_metrics = [
124+
(JobSetHealthiness.SPECIFIED, jobset_config.replicas),
125+
(JobSetHealthiness.ACTIVE, jobset_config.replicas),
126+
(JobSetHealthiness.READY, jobset_config.replicas),
127+
(JobSetHealthiness.FAILED, 0),
128+
(JobSetHealthiness.SUCCEEDED, 0),
129+
(JobSetHealthiness.SUSPENDED, 0),
130+
]
131+
validate_running_tasks = []
132+
for status, expected in running_metrics:
133+
t = jobset.wait_for_jobset_metrics.override(
134+
task_id=f"validate_running_wait_{status.value}"
135+
)(
136+
metric_name=status,
137+
expected_value=expected,
138+
node_pool=cluster_info,
139+
jobset_name=jobset_name,
140+
)
141+
validate_running_tasks.append(t)
138142

139143
suspend_action = jobset.suspended_jobset.override(
140144
task_id="suspend_jobset"
@@ -144,84 +148,77 @@
144148
jobset_name=jobset_name,
145149
)
146150

147-
with TaskGroup(
148-
group_id="validate_suspended_metrics"
149-
) as validate_suspended:
150-
suspended_metrics = [
151-
(JobSetHealthiness.ACTIVE, 0),
152-
(JobSetHealthiness.SUSPENDED, jobset_config.replicas),
153-
]
154-
for status, expected in suspended_metrics:
155-
jobset.wait_for_jobset_metrics.override(
156-
task_id=f"wait_after_suspend_{status.value}"
157-
)(
158-
metric_name=status,
159-
expected_value=expected,
160-
node_pool=cluster_info,
161-
jobset_name=jobset_name,
162-
)
151+
suspended_metrics = [
152+
(JobSetHealthiness.ACTIVE, 0),
153+
(JobSetHealthiness.SUSPENDED, jobset_config.replicas),
154+
]
155+
validate_suspended_tasks = []
156+
for status, expected in suspended_metrics:
157+
t = jobset.wait_for_jobset_metrics.override(
158+
task_id=f"validate_suspended_wait_after_suspend_{status.value}"
159+
)(
160+
metric_name=status,
161+
expected_value=expected,
162+
node_pool=cluster_info,
163+
jobset_name=jobset_name,
164+
)
165+
validate_suspended_tasks.append(t)
163166

164167
resume_action = jobset.resume_jobset.override(task_id="resume_jobset")(
165168
node_pool=cluster_info,
166169
jobset_config=jobset_config,
167170
jobset_name=jobset_name,
168171
)
169172

170-
with TaskGroup(group_id="inject_and_validate_success") as success_test:
171-
cleanup_for_success = jobset.end_workload.override(
172-
task_id="cleanup_before_success_injection"
173-
)(
174-
node_pool=cluster_info,
175-
jobset_config=jobset_config,
176-
jobset_name=jobset_name,
177-
)
178-
179-
start_success_job = jobset.run_workload.override(
180-
task_id="start_success_job"
181-
)(
182-
node_pool=cluster_info,
183-
jobset_config=jobset_config,
184-
jobset_name=jobset_name,
185-
workload_type=SUCCESS_WORKLOAD,
186-
)
187-
188-
validate_succeeded_metric = jobset.wait_for_jobset_metrics.override(
189-
task_id="wait_for_succeeded_count"
190-
)(
191-
metric_name=JobSetHealthiness.SUCCEEDED,
192-
expected_value=jobset_config.replicas,
193-
node_pool=cluster_info,
194-
jobset_name=jobset_name,
195-
)
173+
cleanup_for_success = jobset.end_workload.override(
174+
task_id="cleanup_before_success_injection"
175+
)(
176+
node_pool=cluster_info,
177+
jobset_config=jobset_config,
178+
jobset_name=jobset_name,
179+
)
196180

197-
chain(cleanup_for_success, start_success_job, validate_succeeded_metric)
181+
start_success_job = jobset.run_workload.override(
182+
task_id="start_success_job"
183+
)(
184+
node_pool=cluster_info,
185+
jobset_config=jobset_config,
186+
jobset_name=jobset_name,
187+
workload_type=SUCCESS_WORKLOAD,
188+
)
198189

199-
with TaskGroup(group_id="inject_and_validate_failure") as failure_test:
200-
cleanup_for_failure = jobset.end_workload.override(
201-
task_id="cleanup_before_failure_injection"
202-
)(
203-
node_pool=cluster_info,
204-
jobset_config=jobset_config,
205-
jobset_name=jobset_name,
206-
)
190+
validate_succeeded_metric = jobset.wait_for_jobset_metrics.override(
191+
task_id="wait_for_succeeded_count"
192+
)(
193+
metric_name=JobSetHealthiness.SUCCEEDED,
194+
expected_value=jobset_config.replicas,
195+
node_pool=cluster_info,
196+
jobset_name=jobset_name,
197+
)
207198

208-
start_fail_job = jobset.run_workload.override(task_id="start_fail_job")(
209-
node_pool=cluster_info,
210-
jobset_config=jobset_config,
211-
jobset_name=jobset_name,
212-
workload_type=FAIL_WORKLOAD,
213-
)
199+
cleanup_for_failure = jobset.end_workload.override(
200+
task_id="cleanup_before_failure_injection"
201+
)(
202+
node_pool=cluster_info,
203+
jobset_config=jobset_config,
204+
jobset_name=jobset_name,
205+
)
214206

215-
validate_failed_metric = jobset.wait_for_jobset_metrics.override(
216-
task_id="wait_for_failed_count"
217-
)(
218-
metric_name=JobSetHealthiness.FAILED,
219-
expected_value=jobset_config.replicas,
220-
node_pool=cluster_info,
221-
jobset_name=jobset_name,
222-
)
207+
start_fail_job = jobset.run_workload.override(task_id="start_fail_job")(
208+
node_pool=cluster_info,
209+
jobset_config=jobset_config,
210+
jobset_name=jobset_name,
211+
workload_type=FAIL_WORKLOAD,
212+
)
223213

224-
chain(cleanup_for_failure, start_fail_job, validate_failed_metric)
214+
validate_failed_metric = jobset.wait_for_jobset_metrics.override(
215+
task_id="wait_for_failed_count"
216+
)(
217+
metric_name=JobSetHealthiness.FAILED,
218+
expected_value=jobset_config.replicas,
219+
node_pool=cluster_info,
220+
jobset_name=jobset_name,
221+
)
225222

226223
cleanup_workload = jobset.end_workload.override(
227224
task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
@@ -244,12 +241,24 @@
244241
jobset_name,
245242
create_node_pool,
246243
*startup.tasks,
247-
validate_running,
244+
*validate_running_tasks,
248245
suspend_action,
249-
validate_suspended,
246+
*validate_suspended_tasks,
250247
resume_action,
251-
success_test,
252-
failure_test,
248+
cleanup_for_success,
249+
)
250+
251+
chain(
252+
cleanup_for_success,
253+
start_success_job,
254+
validate_succeeded_metric,
255+
cleanup_for_failure,
256+
)
257+
258+
chain(
259+
cleanup_for_failure,
260+
start_fail_job,
261+
validate_failed_metric,
253262
cleanup_workload,
254263
cleanup_node_pool,
255264
)

‎dags/tpu_observability/utils/jobset_util.py‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,9 @@
3535
from airflow.models.baseoperator import chain
3636
from airflow.models.xcom_arg import XComArg
3737
from airflow.sensors.base import PokeReturnValue
38+
from google.cloud import monitoring_v3
3839
from google.cloud.monitoring_v3 import types
40+
from google.protobuf.message import Message
3941
from websocket import WebSocketConnectionClosedException
4042

4143
from dags.tpu_observability.utils import subprocess_util as subprocess
@@ -1071,6 +1073,8 @@ def verify_recovery_duration(start_time: TimeUtil, end_time: TimeUtil):
10711073
"The 'jobset_time_to_recover' metric requires > 60s to be recorded. "
10721074
"Failing fast to avoid waiting for a missing metric."
10731075
)
1076+
1077+
10741078
@task.sensor(poke_interval=60, timeout=3600, mode="poke")
10751079
def wait_for_jobset_metrics(
10761080
metric_name: JobSetHealthiness,
@@ -1120,17 +1124,22 @@ def wait_for_jobset_metrics(
11201124
end_time=TimeUtil.now(),
11211125
)
11221126

1123-
if not time_series or len(time_series) == 0 or not time_series[0].points:
1127+
if not time_series or not time_series[0].points:
11241128
return False
11251129

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
1130+
point = time_series[0].points[0]
1131+
msg = monitoring_v3.TypedValue.pb(point.value)
1132+
1133+
if isinstance(msg, Message):
1134+
match msg.WhichOneof("value"):
1135+
case "double_value":
1136+
latest_value = msg.double_value
1137+
case "int64_value":
1138+
latest_value = msg.int64_value
1139+
case field_type:
1140+
raise ValueError(f"Unexpected metric value type: {field_type}")
11321141
else:
1133-
latest_value = float(point_value.int64_value)
1142+
raise TypeError("Failed to parse point value as a Protobuf Message")
11341143

11351144
logging.info(
11361145
f"Metric {name_str} for JobSet {jobset_name}: "

0 commit comments

Comments
 (0)