Add a new DAG that tests orbax emergency checkpointing manager saving feature to GCS Bucket feature. (GoogleCloudPlatform#870)

This change adds a new DAG which test and verifies the Orbax Emergency
Checkpointer manager saving functionality of checkpoints to GCS bucket.
It checks `checkpoint_steps` are stored into GCS bucket. 

This DAG automate the installation of the mtc driver, but requires an
existing cluster and an GCS bucket with HNS enabled.

---------

Co-authored-by: Alfred Yu <alfred.yu@cienet.com>

Author: camiloCienet

dags/orbax/maxtext_emc_restore_local.py

@@ -8,23 +8,15 @@
 
 from dags import composer_env
 from dags.common import test_owner
-from dags.common.vm_resource import DockerImage
 from dags.common.vm_resource import XpkClusters
 from dags.multipod.configs import gke_config
-from dags.multipod.configs.common import SetupMode
 from dags.orbax.util import checkpoint_util, test_config_util, validation_util
 from xlml.utils.gke import zone_to_region
 from xlml.utils.xpk import BRANCH_ABHINAV_MTC
 
 DAG_TEST_NAME = "maxtext_emc_orbax_res_local"
 SCHEDULE = "0 11 * * *" if composer_env.is_prod_env() else None
 
-# Only one version of the Docker image is supported at the moment.
-# Other versions (e.g., "stable") may be introduced later.
-DOCKER_IMAGES = [(
-    SetupMode.NIGHTLY,
-    DockerImage.MAXTEXT_TPU_JAX_ORBAX_HEAD,
-)]
 
 with models.DAG(
     dag_id=DAG_TEST_NAME,
@@ -90,7 +82,7 @@
 
   step_to_interrupt = 40
 
-  for mode, image in DOCKER_IMAGES:
+  for mode, image in test_config_util.DOCKER_IMAGES:
     for test_config in test_configs:
       for slice_num in test_config.slices:
         wait_delete_cpc = checkpoint_util.wait_for_cpc_deletion.override(

dags/orbax/maxtext_emc_save_gcs.py

@@ -0,0 +1,161 @@
+"""
+A DAG to run MaxText multi-tier checkpointing with replicator enabled
+validates the local checkpoints are replicated (copy) to bucket
+with HNS (Hierarchical Namespace)
+"""
+
+import datetime
+
+from airflow import models
+
+from dags import composer_env
+from dags.common import test_owner
+from dags.common.vm_resource import XpkClusters
+from dags.multipod.configs import gke_config
+from dags.orbax.util import checkpoint_util
+from dags.orbax.util import validation_util
+from xlml.utils.xpk import BRANCH_ABHINAV_MTC
+from xlml.utils.gke import zone_to_region
+from dags.orbax.util import test_config_util
+
+
+SCHEDULE = "0 14 * * *" if composer_env.is_prod_env() else None
+DAG_TEST_NAME = "maxtext_emc_save_gcs"
+
+
+with models.DAG(
+    dag_id=DAG_TEST_NAME,
+    start_date=datetime.datetime(2025, 6, 30),
+    schedule_interval=SCHEDULE,
+    catchup=False,
+    tags=[
+        "multipod_team",
+        "maxtext",
+        "emergency_checkpointing",
+        "nightly",
+        "orbax",
+    ],
+    description="DAG that verifies the orbax multi-tier checkpointing saving functionality with replicator to GCS bucket",
+    doc_md="""
+      # Multi-tier Checkpoint Validation DAG
+
+      ### Description
+      This DAG (Directed Acyclic Graph) automates the process of validating
+      checkpoint saving when using **Emergency Checkpointer Manager** features.
+      It will check that the checkpoints are being stored in the GCS bucket.
+      Also the steps flag controls how many steps the job will run.
+
+      ### Prerequisites
+      To run this test, you need an existing cluster with the Multi-tier
+      Checkpointing configuration enabled, as well as a bucket with HNS
+      (Hierarchical Namespace) enabled.
+
+      ### Procedures
+      1.  **Apply Configuration:** A Checkpoint Configuration YAML file is
+      applied to the cluster, enabling Multi-tier Checkpoint (MTC) features.
+      2.  **Run Maxtext Jobsets:** The DAG runs a Maxtext jobset.
+      3.  The DAG validates that **GCS checkpoints** are being saved correctly
+      in the `GCS bucket` by checking bucket and pod logs.
+    """,
+    concurrency=2,
+) as dag:
+  checkpointing = test_config_util.Checkpointing(
+      name="emc", enable_multi_tier_checkpointing=False
+  )
+  test_configs = [
+      test_config_util.TestConfig(
+          cluster=XpkClusters.TPU_V5P_128_CLUSTER,
+          machine_type="ct5p-hightpu-4t",
+          accelerator="v5p-128",
+          slices=[2],
+          model_name="llama2-7b",
+          short_id="max-sv-gcs",
+          replicator_backup_time=30,
+          step=75,
+          local_checkpoint_step=20,
+          checkpoint_step=25,
+          base_dir=test_config_util.DEFAULT_BUCKET,
+      ),
+  ]
+  for mode, image in test_config_util.DOCKER_IMAGES:
+    for test_config in test_configs:
+      for slice_num in test_config.slices:
+        # We conditionally set the trigger_rule on the first task.
+        # If first task group failed the next one can execute.
+        wait_delete_cpc = checkpoint_util.wait_for_cpc_deletion.override(
+            trigger_rule="all_done"
+        )(test_config.cpc_config)
+        apply_cpc = checkpoint_util.apply_cpc(test_config.cpc_config)
+
+        # Generate consistent run name for both training phases
+        run_name = validation_util.generate_run_name(
+            short_id=test_config.short_id,
+            checkpointing_type=checkpointing.name,
+            slice_number=slice_num,
+            accelerator=test_config.accelerator,
+        )
+
+        workload_command = test_config.generate_workload_command(
+            checkpoint_dir=test_config_util.DEFAULT_RAM_DISK,
+            run_name=run_name,
+            out_folder=f"maxtext_emc_orbax_save_gcs",
+            enable_multi_tier_checkp=checkpointing.enable_multi_tier_checkpointing,
+            slice_num=slice_num,
+        )
+
+        start_time = validation_util.generate_timestamp()
+
+        maxtext_chkpt_run_test = gke_config.get_gke_config(
+            num_slices=slice_num,
+            cluster=test_config.cluster,
+            time_out_in_min=60,
+            test_name=f"{test_config.short_id}-emc",
+            run_model_cmds=workload_command,
+            docker_image=image.value,
+            test_owner=test_owner.CAMILO_Q,
+        ).run(
+            ramdisk_directory=test_config_util.DEFAULT_RAM_DISK,
+            mtc_enabled=True,
+            xpk_branch=BRANCH_ABHINAV_MTC,
+            skip_post_process=True,
+        )
+
+        end_time = validation_util.generate_timestamp()
+
+        steps_to_validate = test_config.generate_step_to_validate(
+            is_local=False
+        )
+
+        validate_steps = validation_util.validate_checkpoint_at_steps_are_saved(
+            project_id=test_config.cluster.project,
+            location=zone_to_region(test_config.cluster.zone),
+            cluster_name=test_config.cluster.name,
+            ram_disk="gcs",
+            pod_pattern=".*-0-0",
+            start_time=start_time,
+            end_time=end_time,
+            steps_to_validate=steps_to_validate,
+        )
+
+        # Validate that GCS restore happened during the second training run
+        validate_checkpoints_steps_gcs = validation_util.validate_gcs_checkpoint_files(
+            bucket_path=f"{test_config_util.DEFAULT_BUCKET}/maxtext_emc_orbax_save_gcs/{run_name}",
+            steps_to_validate=steps_to_validate,
+        )
+
+        # Final CPC cleanup to ensure symmetric start/end
+        wait_delete_cpc_final = checkpoint_util.wait_for_cpc_deletion.override(
+            trigger_rule="all_done", task_id="wait_delete_cpc_final"
+        )(test_config.cpc_config).as_teardown(setups=apply_cpc)
+
+        (
+            wait_delete_cpc
+            >> apply_cpc
+            >> run_name
+            >> start_time
+            >> maxtext_chkpt_run_test
+            >> end_time
+            >> validate_steps
+            >> validate_checkpoints_steps_gcs
+            >> wait_delete_cpc_final
+        )

dags/orbax/maxtext_mtc_emergency_save_local.py

@@ -14,9 +14,8 @@
 
 from dags import composer_env
 from dags.common import test_owner
-from dags.common.vm_resource import DockerImage, XpkClusters
+from dags.common.vm_resource import XpkClusters
 from dags.multipod.configs import gke_config
-from dags.multipod.configs.common import SetupMode
 from dags.orbax.util import validation_util
 from dags.orbax.util import checkpoint_util
 from xlml.utils.xpk import BRANCH_ABHINAV_MTC
@@ -27,13 +26,6 @@
 SCHEDULE = "0 13 * * *" if composer_env.is_prod_env() else None
 DAG_TEST_NAME = "maxtext_emc_and_mtc_orbax_save_local"
 
-# Only one version of the Docker image is supported at the moment.
-# Other versions (e.g., "stable") may be introduced later.
-DOCKER_IMAGES = [(
-    SetupMode.NIGHTLY,
-    DockerImage.MAXTEXT_TPU_JAX_ORBAX_HEAD,
-)]
-
 
 with models.DAG(
     dag_id=DAG_TEST_NAME,
@@ -112,7 +104,7 @@
     with TaskGroup(
         group_id=f"maxtext_{checkpointing.name}_orbax_save_local",
     ) as group:
-      for mode, image in DOCKER_IMAGES:
+      for mode, image in test_config_util.DOCKER_IMAGES:
         for test_config in test_configs:
           for slice_num in test_config.slices:
             # We conditionally set the trigger_rule on the first task.

dags/orbax/maxtext_mtc_save_gcs.py

@@ -10,9 +10,8 @@
 
 from dags import composer_env
 from dags.common import test_owner
-from dags.common.vm_resource import DockerImage, XpkClusters
+from dags.common.vm_resource import XpkClusters
 from dags.multipod.configs import gke_config
-from dags.multipod.configs.common import SetupMode
 from dags.orbax.util import checkpoint_util
 from dags.orbax.util import validation_util
 from xlml.utils.xpk import BRANCH_ABHINAV_MTC
@@ -23,13 +22,6 @@
 SCHEDULE = "0 15 * * *" if composer_env.is_prod_env() else None
 DAG_TEST_NAME = "maxtext_mtc_orbax_save_gcs"
 
-# Only one version of the Docker image is supported at the moment.
-# Other versions (e.g., "stable") may be introduced later.
-DOCKER_IMAGES = [(
-    SetupMode.NIGHTLY,
-    DockerImage.MAXTEXT_TPU_JAX_ORBAX_HEAD,
-)]
-
 
 with models.DAG(
     dag_id=DAG_TEST_NAME,
@@ -87,7 +79,7 @@
           base_dir=test_config_util.DEFAULT_BUCKET,
       ),
   ]
-  for mode, image in DOCKER_IMAGES:
+  for mode, image in test_config_util.DOCKER_IMAGES:
     for test_config in test_configs:
       for slice_num in test_config.slices:
         # We conditionally set the trigger_rule on the first task.

dags/orbax/util/checkpoint_util.py

@@ -53,7 +53,6 @@ def load_yaml_and_parse_body(
       full parsed YAML body as a dictionary.
     """
 
-    logging.info(f"RAMDISK ==> {self.ramdisk_memory_in_mi}")
     cpc_yaml_template = f"""
       apiVersion: checkpointing.gke.io/v1
       kind: CheckpointConfiguration

dags/orbax/util/test_config_util.py

@@ -11,13 +11,21 @@
 
 from dags import gcs_bucket
 from xlml.utils.gke import zone_to_region
-from dags.common.vm_resource import XpkClusters
+from dags.common.vm_resource import XpkClusters, DockerImage
 from dags.orbax.util import checkpoint_util
+from dags.multipod.configs.common import SetupMode
 
 
 DEFAULT_BUCKET = gcs_bucket.MTC_AUTOMATION_BUCKET
 DEFAULT_RAM_DISK = "/local"
 
+# Only one version of the Docker image is supported at the moment.
+# Other versions (e.g., "stable") may be introduced later.
+DOCKER_IMAGES = [(
+    SetupMode.NIGHTLY,
+    DockerImage.MAXTEXT_TPU_JAX_ORBAX_HEAD,
+)]
+
 # Valid models and sizes for current Maxtext Repository.
 MODELS = {
     "deepseek2",
@@ -111,7 +119,7 @@ def __init__(
     self.step = step
     self.local_checkpoint_step = local_checkpoint_step
     self.checkpoint_step = checkpoint_step
-    self.ram_disk_size = f"{self._get_disk_size(slice_num=max(self.slices), mode='mib',multiplier=30)}Mi"
+    self.ram_disk_size = f"{self._get_disk_size(slice_num=max(self.slices), mode='mib',multiplier=60)}Mi"
     self.base_dir = base_dir
     self.cpc_config = checkpoint_util.CheckpointConfiguration(
         project_id=self.cluster.project,

dags/orbax/util/validation_util.py

@@ -248,6 +248,62 @@ def validate_log_with_gcs(
   return max(gcs_save_step_list), max(gcs_save_step_list_bucket)
 
 
+@task
+def validate_gcs_checkpoint_files(
+    bucket_path: str,
+    steps_to_validate: Optional[list] = None,
+) -> None:
+  """
+  Validates that checkpoint files exist in GCS bucket for expected steps.
+  This function uses the GCS utility to check that checkpoint files
+  are properly saved in the bucket for each expected step.
+  Args:
+    bucket_path: The full gs:// path to the GCS bucket
+    vali_step_list: Optional list of steps to validate
+  Returns:
+    None: Raises AirflowFailException if checkpoint validation fails
+  """
+  if steps_to_validate is None:
+    logging.info(
+        "No validation steps provided, skipping GCS checkpoint validation"
+    )
+    return
+
+  try:
+    checkpoint_files = get_gcs_checkpoint(bucket_path)
+    logging.info(f"Found checkpoint files in GCS: {checkpoint_files}")
+
+    # Extract step directories from checkpoint files
+    found_steps = set()
+    for file_path in checkpoint_files:
+      # Extract directory names that are numeric (step numbers)
+      path_parts = file_path.split("/")
+      for part in path_parts:
+        if part.isdigit():
+          found_steps.add(int(part))
+
+    expected_steps = set(steps_to_validate)
+    missing_steps = expected_steps - found_steps
+
+    logging.info(f"Expected steps: {sorted(expected_steps)}")
+    logging.info(f"Found steps: {sorted(found_steps)}")
+
+    if missing_steps:
+      raise AirflowFailException(
+          f"GCS checkpoint validation failed: Missing checkpoint files for steps {sorted(missing_steps)}. "
+          f"Expected steps: {sorted(steps_to_validate)}, Found steps: {sorted(found_steps)}"
+      )
+
+    logging.info(f"GCS checkpoint validation successful!")
+    logging.info(
+        f"All {len(steps_to_validate)} expected checkpoint files found in GCS"
+    )
+    logging.info(f"Validated steps: {sorted(found_steps)}")
+
+  except Exception as e:
+    raise AirflowFailException(f"Error validating GCS checkpoints: {str(e)}")
+
+
 def list_log_entries(
     project_id: str,
     location: str,