Fix: Change name of DAG to be mtc and emc. mainly because both of these test are being tested.

Fix: Fix some try and catch exceptions that seems to be unnecessary.

Fix: Refactor whole Cpc delete task into two task initiate_cpc and wait_to_delete. wait_to_delete task is a sensor task that will check every 10 s if the delition already happened.

Fix: Two Dags testing mtc and emc (with replicator=true and replicator=false)

Fix: Change loggin.py library to logging_mtc.py names conflict.

Fix: Use abhinav-mtc xpk branch instead of main branch since xpk version > 0.4.1 crash airflow pod. b/437817546

Fix: Correct name of DAG to emtc (Emergency multitier Checkpointer).

Fix: Change name_id of the test to be related to Orbax. Create new tag with Orbax. Remove --restart-on-user-code-failure flag since is already outdated in both main and develop branch

Fix: Backup interval to longer. Only testing local saves. GCS testing is in DAG002

mend

Author: camiloCienet

dags/common/test_owner.py

@@ -55,7 +55,7 @@ class Team(enum.Enum):
 # Multi-tier Checkpointing
 ABHINAV_S = "ABHINAV S."
 XUEFENG_G = "XUEFENG G."
-CAMILO = "CAMILO P."
+CAMILO = "camiloCienet"
 
 # MLCompass
 ORTI_B = "Orti B."

dags/common/vm_resource.py

@@ -257,7 +257,7 @@ class XpkClusters:
       zone=Zone.EUROPE_WEST4_B.value,
   )
   TPU_V5P_128_CLUSTER_ORBAX = XpkClusterConfig(
-      name="b428061876-jf-v5p-128-2",
+      name="orbax-ml-auto-solutions",
       device_version=TpuVersion.V5P,
       core_count=128,
       project=Project.CLOUD_TPU_MULTIPOD_DEV.value,

dags/gcs_bucket.py

@@ -26,5 +26,4 @@
 
 # Multi-tier checkpointing need special permission for GCS Bucket
 # For further question reach out to  Multi-tier Checkpointing Owners.
-MTC_BUCKET = "gs://mtc-bucket-us-east5/output"
 MTC_AUTOMATION_BUCKET = "gs://mtc-automation-bucket"

dags/multipod/maxtext_multi_tier_p2_checkpointing.py

@@ -39,7 +39,7 @@
     concurrency=2,
 ) as dag:
   base_output_directory = (
-      f"{gcs_bucket.MTC_BUCKET}/maxtext_multi_tier_p2_checkpointing"
+      f"{gcs_bucket.MTC_AUTOMATION_BUCKET}/maxtext_multi_tier_p2_checkpointing"
   )
   dataset_path = gcs_bucket.MAXTEXT_DIR
   docker_images = [

dags/orbax/maxtext_mtc_emergency_save_local.py

@@ -0,0 +1,190 @@
+"""
+A DAG to run MaxText multi-tier and emergency checkpointing tests.
+
+This DAG performs a series of tests to save and validate checkpoints
+for the MaxText model. It runs tests in two modes: one with the replicator
+service enabled (Multi-tier Checkpointing) and one without
+(Emergency Checkpointing). The tests are executed on a TPU multi-pod cluster.
+
+Variables:
+  SCHEDULE (str): The cron schedule.
+  DAG_TEST_NAME (str): Unique DAG name.
+  BASE_OUTPUT_DIRECTORY (str): GCS path for saved checkpoints.
+  DOCKER_IMAGES (list): Tuples of setup mode and Docker image.
+  RAM_DISK (str): Local directory mounted on each machine.
+  TEST_CONFIGS (dict): Accelerator type and slice count.
+  CLUSTERS (dict): Maps accelerator to TPU cluster configs.
+  STEP (int): Total training steps.
+  LOCAL_CHECKPOINT_PERIOD (int): Frequency (in steps) for local checkpoints.
+  REPLICATOR_BACKUP_INTERVAL_MINUTES (int): Interval (in minutes) for replicator.
+  NAME_PREFIX (str): Prefix for test run names.
+  MODEL_NAME (str): Name of the model under test.
+"""
+
+import datetime
+from dataclasses import asdict
+
+from airflow import models
+from airflow.utils.task_group import TaskGroup
+
+from dags import composer_env, gcs_bucket
+from dags.common import test_owner
+from dags.common.vm_resource import DockerImage, XpkClusters
+from dags.multipod.configs import gke_config
+from dags.multipod.configs.common import SetupMode
+from dags.orbax.util import logging_mtc as log
+from dags.orbax.util import multi_tier_checkpoint_util as mtc
+from xlml.utils.xpk import BRANCH_ABHINAV_MTC
+from absl import logging
+
+SCHEDULE = "0 10 * * *" if composer_env.is_prod_env() else None
+DAG_TEST_NAME = "maxtext_emc_and_mtc_orbax_save_local"
+BASE_OUTPUT_DIRECTORY = (
+    f"{gcs_bucket.MTC_AUTOMATION_BUCKET}/"
+)
+DOCKER_IMAGES = [(
+    SetupMode.NIGHTLY,
+    DockerImage.MAXTEXT_TPU_JAX_NIGHTLY,
+)]
+RAM_DISK = "/local"
+TEST_CONFIGS = {"v5p-128": [2]}
+CLUSTERS = {"v5p-128": XpkClusters.TPU_V5P_128_CLUSTER_ORBAX}
+STEP = 100
+LOCAL_CHECKPOINT_PERIOD = 20
+REPLICATOR_BACKUP_INTERVAL_MINUTES = 30
+NAME_PREFIX = "max-sv"
+MODEL_NAME = "llama2-7b"
+
+with models.DAG(
+    dag_id=DAG_TEST_NAME,
+    schedule_interval=SCHEDULE,
+    tags=[
+        "multipod_team",
+        "maxtext",
+        "emergency_checkpoint_manager",
+        "multitier_checkpointing",
+        "nightly",
+        "orbax",
+    ],
+    start_date=datetime.datetime(2025, 6, 12),
+    catchup=False,
+    concurrency=2,
+) as dag:
+  tests_to_run_seq = []
+  for use_replicator_flag in [True, False]:
+
+    # These are the two test cases we are testing
+    # use_replicator = True --> Multi-tier Checkpointing
+    # use_replicator = False --> Emergency Checkpointing
+    prefix_replicator = "mtc" if use_replicator_flag else "emc"
+    folder = f"maxtext_{prefix_replicator}_orbax_save_local"
+    BASE_OUTPUT_DIRECTORY = f"{gcs_bucket.MTC_AUTOMATION_BUCKET}/{folder}"
+    with TaskGroup(
+        group_id=f"maxtext_{prefix_replicator}_orbax_save_local"
+    ) as task:
+      for mode, image in DOCKER_IMAGES:
+        for accelerator, slices in TEST_CONFIGS.items():
+          for slice_num in slices:
+
+            # Multi-tier Checkpointing Configuration
+            # TODO(Camilo): Make this configurable based on the accelerator type.
+            # Currently, it is hardcoded for v5p.
+            cpc_conf = (
+                CLUSTERS[accelerator].project,
+                CLUSTERS[accelerator].zone[:-2],
+                CLUSTERS[accelerator].name,
+                gcs_bucket.MTC_AUTOMATION_BUCKET.split("gs://")[1],
+                "ct5p-hightpu-4t",
+                "google.com/tpu",
+                "800000Mi",
+            )
+
+            # We conditionally set the trigger_rule on the first task.
+            # If first task group failed the next one can execute.
+            init_delete_cpc = mtc.initiate_cpc_deletion(*cpc_conf)
+            wait_delete_cpc = mtc.wait_for_cpc_deletion.override(trigger_rule="all_done")(*cpc_conf)
+            apply_cpc = mtc.apply_cpc(*cpc_conf)
+            run_time = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")
+            run_name = f"{NAME_PREFIX}-{prefix_replicator}-{slice_num}x-{accelerator}-{run_time}"
+            workload_command = (
+                "export TPU_PREMAPPED_BUFFER_SIZE=52428800000 && "
+                "export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES=52428800000 && "
+                "python3 -m MaxText.train MaxText/configs/base.yml remat_policy=full "
+                f"global_parameter_scale=1 base_output_directory={BASE_OUTPUT_DIRECTORY} "
+                f"dataset_type=synthetic steps={STEP} per_device_batch_size=1 "
+                f"max_target_length=256 model_name={MODEL_NAME} per_device_batch_size=2 "
+                "reuse_example_batch=1 enable_emergency_checkpoint=true "
+                f"local_checkpoint_directory={RAM_DISK} local_checkpoint_period={LOCAL_CHECKPOINT_PERIOD} "
+                f"use_replicator_service={use_replicator_flag} replicator_backup_interval_minutes={REPLICATOR_BACKUP_INTERVAL_MINUTES} "
+                f"run_name={run_name}",
+            )
+
+            start_time = log.generate_timestamp()
+
+            maxtext_chkpt_run_test = gke_config.get_gke_config(
+                num_slices=slice_num,
+                cluster=CLUSTERS[accelerator],
+                time_out_in_min=60,
+                test_name=f"{NAME_PREFIX}-{prefix_replicator}",
+                run_model_cmds=workload_command,
+                docker_image=image.value,
+                test_owner=test_owner.CAMILO,
+            ).run(
+                ramdisk_directory=RAM_DISK,
+                mtc_enabled=True,
+                xpk_branch=BRANCH_ABHINAV_MTC,
+                skip_post_process=True,
+            )
+
+            cleanup_command = (f"rm -rf {RAM_DISK}/*",)
+            ram_disk_cleanup = gke_config.get_gke_config(
+                num_slices=slice_num,
+                cluster=CLUSTERS[accelerator],
+                time_out_in_min=60,
+                test_name=f"{NAME_PREFIX}-cl",
+                run_model_cmds=cleanup_command,
+                docker_image=image.value,
+                test_owner=test_owner.CAMILO,
+            ).run(
+                ramdisk_directory=RAM_DISK,
+                mtc_enabled=True,
+                xpk_branch=BRANCH_ABHINAV_MTC,
+                skip_post_process=True,
+            )
+
+            end_time = log.generate_timestamp()
+            vali_step = STEP- 1
+            vali_step_list = [
+                i for i in range(0, vali_step, LOCAL_CHECKPOINT_PERIOD)
+            ]
+            vali_step_list.append(vali_step)
+
+            # Here we are looking for the string '(blocking + background)'.
+            # We will compare expected steps with the ones we found when query this regex. Should be the same
+            # If for some reason the restore start from 0 this task will fail because len(valid_step_list) != len(founded_steps)
+            validate_log = log.validate_log_with_step(
+                project_id=CLUSTERS[accelerator].project,
+                location=CLUSTERS[accelerator].zone[:-2],
+                cluster_name=CLUSTERS[accelerator].name,
+                text_filter="(blocking + background).",
+                start_time=start_time,
+                end_time=end_time,
+                vali_step_list=vali_step_list,
+            )
+
+            (
+                init_delete_cpc
+                >> wait_delete_cpc
+                >> apply_cpc
+                >> start_time
+                >> maxtext_chkpt_run_test
+                >> ram_disk_cleanup
+                >> end_time
+                >> validate_log
+            )
+    # Add to a global list of test to be run in a sequential way
+    tests_to_run_seq.append(task)
+
+  # Chain the task groups sequentially
+  for idx_test in range(len(tests_to_run_seq) - 1):
+    tests_to_run_seq[idx_test] >> tests_to_run_seq[idx_test + 1]

dags/orbax/maxtext_multi_tier_chechpoint_save_local.py

@@ -1,12 +1,13 @@
 """Utilities to get workloads logs and some utils."""
 
-from airflow.decorators import task
-from airflow.exceptions import AirflowFailException
-from google.cloud import logging as log_explorer
 from datetime import datetime, timezone, timedelta
 from typing import Optional
 from absl import logging
 
+from airflow.decorators import task
+from airflow.exceptions import AirflowFailException
+from google.cloud import logging as log_explorer
+
 
 @task
 def generate_timestamp():