Skip to content

Commit 29e6ffd

Browse files
ooops678Depp Lee
ooops678
and
Depp Lee
authored
Add a new DAG that tests MaxText MTC restoring from peer feature (GoogleCloudPlatform#880)
Add a new DAG that tests orbax MTC restoring from peer feature --------- Co-authored-by: Depp Lee <depp.lee@cienet.com>
1 parent d627d3a commit 29e6ffd

3 files changed

Lines changed: 200 additions & 4 deletions

File tree

‎dags/orbax/maxtext_mtc_restore_local.py‎

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
1+
"""A DAG to run MaxText MTC.
2+
3+
Validates the local checkpoints are restored as expected
4+
"""
5+
6+
import datetime
7+
8+
from airflow import models
9+
10+
from dags import composer_env
11+
from dags.common import test_owner
12+
from dags.common.vm_resource import XpkClusters
13+
from dags.multipod.configs import gke_config
14+
from dags.orbax.util import checkpoint_util
15+
from dags.orbax.util import test_config_util
16+
from dags.orbax.util import validation_util
17+
from xlml.utils.gke import zone_to_region
18+
from xlml.utils.xpk import BRANCH_ABHINAV_MTC
19+
20+
DAG_TEST_NAME = "maxtext_mtc_orbax_res_local"
21+
SCHEDULE = "0 14 * * *" if composer_env.is_prod_env() else None
22+
23+
with models.DAG(
24+
dag_id=DAG_TEST_NAME,
25+
start_date=datetime.datetime(2025, 6, 30),
26+
schedule_interval=SCHEDULE,
27+
catchup=False,
28+
tags=["multipod_team", "maxtext", "nightly", "orbax", "mtc"],
29+
description="DAG to verify MaxText's multi-tier restore from local checkpoints after a node interruption.",
30+
doc_md="""
31+
# MaxText Multi-tier Restore from Local Checkpoint Validation DAG
32+
33+
### Description
34+
This DAG validates the multi-tier restore capability of MaxText from local
35+
checkpoints. It simulates a node failure during a training job and verifies
36+
that the job can successfully resume from the last saved local checkpoint
37+
replicated from a peer. This test is critical for ensuring the resilience
38+
of long-running training jobs against hardware failures when using
39+
Multi-tier Checkpointing (MTC).
40+
41+
### Prerequisites
42+
- An existing GKE cluster with the Multi-tier Checkpointing (MTC)
43+
configuration enabled, which provides the necessary CSI driver for
44+
RAM disk (`/local`).
45+
- A GCS bucket for storing logs and base model outputs.
46+
47+
### Procedures
48+
1. **Apply MTC Configuration:** A `CheckpointingPolicyConfiguration` (CPC)
49+
is applied to the cluster to set up the MTC environment.
50+
2. **Run MaxText with Interruption:** A MaxText training job is initiated.
51+
During its execution, a node interruption is simulated to trigger the
52+
restore mechanism.
53+
3. **Validate Restore:** The DAG inspects the application logs to confirm
54+
that a `'restore'` event occurred and that the checkpoint was copied
55+
from a peer node.
56+
4. **Validate Checkpoint Integrity:** It then verifies that the training job
57+
resumed and continued to save checkpoints correctly after the restore,
58+
ensuring no data was lost.
59+
""",
60+
concurrency=2,
61+
) as dag:
62+
test_configs = [
63+
test_config_util.TestConfig(
64+
cluster=XpkClusters.TPU_V5P_128_CLUSTER,
65+
machine_type="ct5p-hightpu-4t",
66+
accelerator="v5p-128",
67+
slices=[2],
68+
model_name="llama2-7b",
69+
short_id="max-res-loc",
70+
replicator_backup_time=30,
71+
step=300,
72+
checkpoint_step=100,
73+
local_checkpoint_step=20,
74+
base_dir=test_config_util.DEFAULT_BUCKET,
75+
),
76+
]
77+
checkpointing = test_config_util.Checkpointing(
78+
name="mtc", enable_multi_tier_checkpointing=True
79+
)
80+
81+
step_to_interrupt = 60
82+
83+
for mode, image in test_config_util.DOCKER_IMAGES:
84+
for test_config in test_configs:
85+
for slice_num in test_config.slices:
86+
wait_delete_cpc = checkpoint_util.wait_for_cpc_deletion.override(
87+
trigger_rule="all_done"
88+
)(test_config.cpc_config)
89+
apply_cpc = checkpoint_util.apply_cpc(test_config.cpc_config)
90+
91+
run_name = validation_util.generate_run_name(
92+
short_id=test_config.short_id,
93+
checkpointing_type=checkpointing.name,
94+
slice_number=slice_num,
95+
accelerator=test_config.accelerator,
96+
)
97+
98+
workload_command = test_config.generate_workload_command(
99+
checkpoint_dir=test_config_util.DEFAULT_RAM_DISK,
100+
run_name=run_name,
101+
slice_num=slice_num,
102+
out_folder="maxtext_mtc_orbax_res_local",
103+
enable_multi_tier_checkp=checkpointing.enable_multi_tier_checkpointing,
104+
)
105+
106+
start_time = validation_util.generate_timestamp.override(
107+
task_id="generate_start_time"
108+
)()
109+
110+
maxtext_chkpt_run_test = gke_config.get_gke_config(
111+
num_slices=slice_num,
112+
cluster=test_config.cluster,
113+
time_out_in_min=60,
114+
test_name=f"{test_config.short_id}-mtc",
115+
run_model_cmds=workload_command,
116+
docker_image=image.value,
117+
test_owner=test_owner.DEPP_L,
118+
).run_with_node_interruption(
119+
ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
120+
mtc_enabled=True,
121+
xpk_branch=BRANCH_ABHINAV_MTC,
122+
skip_post_process=True,
123+
expect_reach_to_step=step_to_interrupt,
124+
)
125+
126+
end_time = validation_util.generate_timestamp.override(
127+
task_id="generate_end_time"
128+
)()
129+
130+
validate_restore_step = (
131+
validation_util.validate_restored_correct_checkpoint(
132+
project_id=test_config.cluster.project,
133+
location=zone_to_region(test_config.cluster.zone),
134+
cluster_name=test_config.cluster.name,
135+
pod_pattern=f"{test_config.short_id}-mtc.*0-0",
136+
interrupt_at_step=step_to_interrupt,
137+
start_time=start_time,
138+
end_time=end_time,
139+
)
140+
)
141+
142+
log_filters = [
143+
"jsonPayload.message:\"'event_type': 'restore'\"",
144+
"jsonPayload.message:\"'directory': '/local\"",
145+
]
146+
validate_restored_source = validation_util.validate_log_exist.override(
147+
task_id="validate_restore_copy_from_peer"
148+
)(
149+
project_id=test_config.cluster.project,
150+
location=zone_to_region(test_config.cluster.zone),
151+
cluster_name=test_config.cluster.name,
152+
text_filter=" AND ".join(log_filters),
153+
start_time=start_time,
154+
end_time=end_time,
155+
)
156+
157+
steps_to_validate = test_config.generate_step_to_validate(is_local=True)
158+
159+
validate_local_saved_steps = (
160+
validation_util.validate_checkpoint_at_steps_are_saved(
161+
project_id=test_config.cluster.project,
162+
location=zone_to_region(test_config.cluster.zone),
163+
cluster_name=test_config.cluster.name,
164+
start_time=start_time,
165+
end_time=end_time,
166+
steps_to_validate=steps_to_validate,
167+
)
168+
)
169+
170+
# Final CPC cleanup to ensure symmetric start/end
171+
wait_delete_cpc_final = checkpoint_util.wait_for_cpc_deletion.override(
172+
trigger_rule="all_done", task_id="wait_delete_cpc_final"
173+
)(test_config.cpc_config).as_teardown(setups=apply_cpc)
174+
175+
(
176+
wait_delete_cpc
177+
>> apply_cpc
178+
>> run_name
179+
>> start_time
180+
>> maxtext_chkpt_run_test
181+
>> end_time
182+
>> validate_restore_step
183+
>> validate_restored_source
184+
>> validate_local_saved_steps
185+
>> wait_delete_cpc_final
186+
)

‎plugins/allow_list.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,3 +23,4 @@ validate_interruption_count_gke_other
2323
maxtext_emc_and_mtc_orbax_save_local
2424
maxtext_mtc_orbax_save_gcs
2525
maxtext_emc_orbax_res_local
26+
maxtext_mtc_orbax_res_local

‎xlml/utils/xpk.py‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -393,10 +393,19 @@ def wait_for_workload_reach_step(
393393
logs = core_api.read_namespaced_pod_log(
394394
name=pod.metadata.name, namespace=pod.metadata.namespace
395395
)
396-
# Check if the workload completed step reached to the expected step
397-
if f"completed step: {expect_reach_to_step}" in logs:
398-
logging.info("Reached to the expected step %s.", expect_reach_to_step)
399-
return True
396+
# Check if the workload completed step reached over the expected step
397+
completed_step_matches = re.findall(r"completed step: (\d+)", logs)
398+
if completed_step_matches:
399+
current_step = int(completed_step_matches[-1])
400+
if current_step >= int(expect_reach_to_step):
401+
logging.info(
402+
"Reached to the expected step %s. Current step is %s.",
403+
expect_reach_to_step,
404+
current_step,
405+
)
406+
return True
407+
408+
logging.info("Waiting for reaching expected step %s.", expect_reach_to_step)
400409

401410
return False
402411

0 commit comments

Comments
 (0)