|
| 1 | +"""A DAG to run MaxText EMC. |
| 2 | +
|
| 3 | +Validates the GCS 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_emc_orbax_res_gcs" |
| 21 | +SCHEDULE = "0 10 * * *" 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=[ |
| 29 | + "multipod_team", |
| 30 | + "maxtext", |
| 31 | + "emergency_checkpointing", |
| 32 | + "nightly", |
| 33 | + "orbax", |
| 34 | + ], |
| 35 | + description="DAG to verify MaxText's emergency restore from GCS checkpoints after a full cluster interruption.", |
| 36 | + doc_md=""" |
| 37 | + # MaxText Emergency Restore from GCS Validation DAG |
| 38 | +
|
| 39 | + ### Description |
| 40 | + This DAG validates the emergency restore capability of MaxText from |
| 41 | + Google Cloud Storage (GCS) checkpoints. It simulates a full cluster |
| 42 | + failure during a training job and verifies that the job can successfully |
| 43 | + resume from the last saved GCS checkpoint. This test is critical for |
| 44 | + ensuring the resilience of long-running training jobs against catastrophic |
| 45 | + failures where local checkpoints are lost. |
| 46 | +
|
| 47 | + ### Prerequisites |
| 48 | + - An existing GKE cluster with the Multi-tier Checkpointing (MTC) |
| 49 | + configuration enabled, which provides the necessary CSI driver for |
| 50 | + RAM disk (`/local`). |
| 51 | + - A GCS bucket for storing logs and checkpoints. |
| 52 | +
|
| 53 | + ### Procedures |
| 54 | + 1. **Apply MTC Configuration:** A `CheckpointingPolicyConfiguration` (CPC) |
| 55 | + is applied to the cluster to set up the MTC environment. |
| 56 | + 2. **Run MaxText with Interruption:** A MaxText training job is initiated. |
| 57 | + During its execution, the last node is deleted to simulate a full |
| 58 | + cluster interruption, triggering the emergency restore mechanism from GCS. |
| 59 | + 3. **Validate Restore:** The DAG inspects the application logs to confirm |
| 60 | + that an `'emergency_restore'` event occurred from a GCS source. |
| 61 | + 4. **Validate Checkpoint Integrity:** It then verifies that the training job |
| 62 | + resumed and continued to save checkpoints correctly after the restore. |
| 63 | + """, |
| 64 | + concurrency=2, |
| 65 | +) as dag: |
| 66 | + checkpointing = test_config_util.Checkpointing( |
| 67 | + name="emc", enable_multi_tier_checkpointing=False |
| 68 | + ) |
| 69 | + test_configs = [ |
| 70 | + test_config_util.TestConfig( |
| 71 | + cluster=XpkClusters.TPU_V5P_128_CLUSTER, |
| 72 | + machine_type="ct5p-hightpu-4t", |
| 73 | + accelerator="v5p-128", |
| 74 | + slices=[2], |
| 75 | + model_name="llama2-7b", |
| 76 | + short_id="max-res-gcs", |
| 77 | + replicator_backup_time=30, |
| 78 | + step=200, |
| 79 | + checkpoint_step=30, |
| 80 | + local_checkpoint_step=200, |
| 81 | + base_dir=test_config_util.DEFAULT_BUCKET, |
| 82 | + ), |
| 83 | + ] |
| 84 | + |
| 85 | + step_to_interrupt = 60 |
| 86 | + |
| 87 | + for mode, image in test_config_util.DOCKER_IMAGES: |
| 88 | + for test_config in test_configs: |
| 89 | + for slice_num in test_config.slices: |
| 90 | + wait_delete_cpc = checkpoint_util.wait_for_cpc_deletion.override( |
| 91 | + trigger_rule="all_done" |
| 92 | + )(test_config.cpc_config) |
| 93 | + |
| 94 | + apply_cpc = checkpoint_util.apply_cpc(test_config.cpc_config) |
| 95 | + |
| 96 | + # Generate consistent run name for both training phases |
| 97 | + run_name = validation_util.generate_run_name( |
| 98 | + short_id=test_config.short_id, |
| 99 | + checkpointing_type=checkpointing.name, |
| 100 | + slice_number=slice_num, |
| 101 | + accelerator=test_config.accelerator, |
| 102 | + ) |
| 103 | + |
| 104 | + workload_command = test_config.generate_workload_command( |
| 105 | + checkpoint_dir=test_config_util.DEFAULT_RAM_DISK, |
| 106 | + run_name=run_name, |
| 107 | + slice_num=slice_num, |
| 108 | + out_folder="maxtext_emc_orbax_res_gcs", |
| 109 | + enable_multi_tier_checkp=checkpointing.enable_multi_tier_checkpointing, |
| 110 | + ) |
| 111 | + |
| 112 | + start_time = validation_util.generate_timestamp.override( |
| 113 | + task_id="generate_start_time" |
| 114 | + )() |
| 115 | + |
| 116 | + maxtext_chkpt_run_test = gke_config.get_gke_config( |
| 117 | + num_slices=slice_num, |
| 118 | + cluster=test_config.cluster, |
| 119 | + time_out_in_min=60, |
| 120 | + test_name=f"{test_config.short_id}-emc", |
| 121 | + run_model_cmds=workload_command, |
| 122 | + docker_image=image.value, |
| 123 | + test_owner=test_owner.DEPP_L, |
| 124 | + ).run_with_node_interruption( |
| 125 | + ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, |
| 126 | + mtc_enabled=True, |
| 127 | + xpk_branch=BRANCH_ABHINAV_MTC, |
| 128 | + skip_post_process=True, |
| 129 | + last_node=True, |
| 130 | + expect_reach_to_step=step_to_interrupt, |
| 131 | + ) |
| 132 | + |
| 133 | + end_time = validation_util.generate_timestamp.override( |
| 134 | + task_id="generate_end_time" |
| 135 | + )() |
| 136 | + |
| 137 | + validate_restore_step = ( |
| 138 | + validation_util.validate_restored_correct_checkpoint( |
| 139 | + project_id=test_config.cluster.project, |
| 140 | + location=zone_to_region(test_config.cluster.zone), |
| 141 | + cluster_name=test_config.cluster.name, |
| 142 | + pod_pattern=f"{test_config.short_id}-emc.*1-0", |
| 143 | + interrupt_at_step=step_to_interrupt, |
| 144 | + start_time=start_time, |
| 145 | + end_time=end_time, |
| 146 | + check_last_two_local_saves=False, |
| 147 | + ) |
| 148 | + ) |
| 149 | + |
| 150 | + log_filters = [ |
| 151 | + "jsonPayload.message:\"'event_type': 'emergency_restore'\"", |
| 152 | + "jsonPayload.message:\"'is_restoring_slice': True\"", |
| 153 | + "jsonPayload.message:\"'directory': 'gs://\"", |
| 154 | + ] |
| 155 | + validate_restored_source = validation_util.validate_log_exist.override( |
| 156 | + task_id="validate_emc_restored_from_gcs" |
| 157 | + )( |
| 158 | + project_id=test_config.cluster.project, |
| 159 | + location=zone_to_region(test_config.cluster.zone), |
| 160 | + cluster_name=test_config.cluster.name, |
| 161 | + text_filter=" AND ".join(log_filters), |
| 162 | + start_time=start_time, |
| 163 | + end_time=end_time, |
| 164 | + ) |
| 165 | + |
| 166 | + gcs_saved_steps_to_validate = test_config.generate_step_to_validate( |
| 167 | + is_local=False |
| 168 | + ) |
| 169 | + |
| 170 | + validate_saved_checkpoints_steps_gcs = validation_util.validate_gcs_checkpoint_files( |
| 171 | + bucket_path=( |
| 172 | + f"{test_config_util.DEFAULT_BUCKET}/{DAG_TEST_NAME}/{run_name}" |
| 173 | + ), |
| 174 | + steps_to_validate=gcs_saved_steps_to_validate, |
| 175 | + ) |
| 176 | + # Final CPC cleanup to ensure symmetric start/end |
| 177 | + wait_delete_cpc_final = checkpoint_util.wait_for_cpc_deletion.override( |
| 178 | + trigger_rule="all_done", task_id="wait_delete_cpc_final" |
| 179 | + )(test_config.cpc_config).as_teardown(setups=apply_cpc) |
| 180 | + |
| 181 | + ( |
| 182 | + wait_delete_cpc |
| 183 | + >> apply_cpc |
| 184 | + >> run_name |
| 185 | + >> start_time |
| 186 | + >> maxtext_chkpt_run_test |
| 187 | + >> end_time |
| 188 | + >> validate_restore_step |
| 189 | + >> validate_restored_source |
| 190 | + >> validate_saved_checkpoints_steps_gcs |
| 191 | + >> wait_delete_cpc_final |
| 192 | + ) |
0 commit comments