add a new DAG 'maxtext_multi_tier_sav01_save_local' GoogleCloudPlatform#804

Author: Camilo QUINONES

dags/common/test_owner.py

@@ -54,7 +54,7 @@ class Team(enum.Enum):
 # Multi-tier Checkpointing
 ABHINAV_S = "ABHINAV S."
 XUEFENG_G = "XUEFENG G."
-ERNIE_C = "ERNIE C."
+CAMILO= "CAMILO Q."
 
 # MLCompass
 ORTI_B = "Orti B."

dags/common/vm_resource.py

@@ -256,6 +256,13 @@ class XpkClusters:
       project=Project.CLOUD_TPU_MULTIPOD_DEV.value,
       zone=Zone.EUROPE_WEST4_B.value,
   )
+  TPU_V5P_128_CLUSTER_ORBAX = XpkClusterConfig(
+      name="b428061876-jf-v5p-128-2",
+      device_version=TpuVersion.V5P,
+      core_count=128,
+      project=Project.CLOUD_TPU_MULTIPOD_DEV.value,
+      zone=Zone.EUROPE_WEST4_B.value,
+  )
   TPU_V5E_256_CLUSTER = XpkClusterConfig(
       name="v5e-256-bodaborg-europe-west4",
       device_version=TpuVersion.V5E,

dags/orbax/maxtext_multi_tier_chechpoint_save_local.py

@@ -36,7 +36,7 @@
   )]
   ram_disk = "/local"
   test_configs = {"v5p-128": [2]}
-  clusters = {"v5p-128": XpkClusters.TPU_V5P_128_CLUSTER}
+  clusters = {"v5p-128": XpkClusters.TPU_V5P_128_CLUSTER_ORBAX}
   step = 100
   local_checkpoint_period = 20
   replicator_backup_interval_minutes = 1
@@ -83,7 +83,7 @@
             test_name=f"{name_prefix}",
             run_model_cmds=workload_command,
             docker_image=image.value,
-            test_owner=test_owner.ERNIE_C,
+            test_owner=test_owner.CAMILO,
         ).run(
             ramdisk_directory=ram_disk,
             mtc_enabled=True,
@@ -100,7 +100,7 @@
             test_name=f"{name_prefix}-cleanup",
             run_model_cmds=cleanup_command,
             docker_image=image.value,
-            test_owner=test_owner.ERNIE_C,
+            test_owner=test_owner.CAMILO,
         ).run(
             ramdisk_directory=ram_disk,
             mtc_enabled=True,

xlml/utils/xpk.py

@@ -51,8 +51,12 @@ def get_xpk_setup_cmd(tmpdir, branch: str = MAIN_BRANCH):
   cmds = [
       "set -xue",
       clone_branch,
+      f"cd {tmpdir}/xpk/src/xpk",
+      "sed -i '/validate_dependencies()/s/^/#/' main.py || true",
+      "cat main.py",
       "pip install ruamel.yaml docker",
   ]
+
   return cmds
 
 
@@ -121,6 +125,13 @@ def run_workload(
     if mtc_enabled:
       workload_create_cmd += " --mtc-enabled"
 
+    # For Orbax DAG
+    if ramdisk_directory and mtc_enabled:
+      workload_create_cmd = workload_create_cmd.replace(
+          " --restart-on-user-code-failure", ""
+      )
+      workload_create_cmd += " --max-restarts=50"
+
     # If using a valid GPU and the XPK branch is set to "main", then branch is switch to "v0.4.1".
     if is_valid_gpu_version(accelerator_type) and xpk_branch == MAIN_BRANCH:
       xpk_branch = "v0.4.1"