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

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

@@ -56,7 +56,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/orbax/maxtext_mtc_emergency_save_local.py

@@ -0,0 +1,159 @@
+"""A DAG to run MaxText multi-tier checkpointing tests (phase2: save & validate)."""
+
+import datetime
+
+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 as log
+from dags.orbax.util import multi_tier_checkpoint_util as mtc
+from xlml.utils.xpk import BRANCH_ABHINAV_MTC
+
+SCHEDULE = "0 10 * * *" if composer_env.is_prod_env() else None
+
+with models.DAG(
+    dag_id="maxtext_emtc_orbax_sav01_save_local",
+    schedule_interval=SCHEDULE,
+    tags=[
+        "multipod_team",
+        "maxtext",
+        "multi_tier_p2_chkpt_save_local",
+        "nightly",
+        "orbax",
+    ],
+    start_date=datetime.datetime(2025, 6, 12),
+    catchup=False,
+    concurrency=2,
+) as dag:
+  base_output_directory = (
+      f"{gcs_bucket.MTC_AUTOMATION_BUCKET}/maxtext_emtc_orbax_sav01_save_local"
+  )
+  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-p2-sv"
+  model_name = "llama2-7b"
+  tests_to_run_seq = []
+  # Define the flags for replicator usage to create two separate test cases
+  use_replicator_flags = ["True", "False"]
+
+  for i, use_replicator_flag in enumerate(use_replicator_flags):
+    with TaskGroup(
+        group_id=f"maxtext_emtc_orbax_sav01_save_local_replicator_{use_replicator_flag.lower()}"
+    ) as task:
+      for mode, image in docker_images:
+        for accelerator, slices in test_configs.items():
+          for slice_num in slices:
+            cpc = (
+                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",
+            )
+            # This is the key change. We conditionally set the trigger_rule
+            # on the very first task of each task group (except the first one).
+            # Start from fresh apply cpc yaml mtc driver
+            delete_cpc = (
+                mtc.delete_cpc.override(trigger_rule="all_done")(*cpc)
+                if i > 0
+                else mtc.delete_cpc(*cpc)
+            )
+            apply_cpc = mtc.apply_cpc(*cpc)
+            run_time = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")
+            run_name = f"{name_prefix}-{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_phase2_chkpt_test = gke_config.get_gke_config(
+                num_slices=slice_num,
+                cluster=clusters[accelerator],
+                time_out_in_min=60,
+                test_name=f"{name_prefix}-{use_replicator_flag.lower()}",
+                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,
+            )
+
+            (
+                delete_cpc
+                >> apply_cpc
+                >> start_time
+                >> maxtext_phase2_chkpt_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,12 @@
 """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():

dags/orbax/util/__init__.py

@@ -1,13 +1,19 @@
+
+"""Utility functions for managing Multi-tier Cluster Configuration.
+
+This module provides tasks for creating, applying, and deleting
+Multi-tier Driver cluster Configurations for enebale Multi Tier Checkpointing.
+"""
 from absl import logging
-from airflow.decorators import task
-from airflow.exceptions import AirflowFailException
-from kubernetes import client as k8s_client
-from kubernetes.client.rest import ApiException
-from xlml.utils import gke
 import yaml
 import time
 
-# --- Utility Functions ---
+from kubernetes import client as k8s_client
+from kubernetes.client.rest import ApiException
+from airflow.decorators import task
+from airflow.exceptions import AirflowFailException
+
+from xlml.utils import gke
 
 
 def _get_custom_objects_api_client(