add maxtext_multi_tier_checkpointing_recover DAG

dags/common/test_owner.py

@@ -54,6 +54,7 @@ class Team(enum.Enum):
 # Multi-tier Checkpointing
 ABHINAV_S = "ABHINAV S."
 XUEFENG_G = "XUEFENG G."
+JACKY_F = "JACKY F."
 
 # MLCompass
 ORTI_B = "Orti B."

dags/map_reproducibility/internal_runs/a3ultra_maxtext_benchmarking_dags.py

@@ -19,9 +19,9 @@
 
 from airflow import models
 from dags import composer_env
-from dags.map_reproducibility.utils.constants import Image
+from dags.map_reproducibility.utils.constants import Image, WorkloadLauncher
 from dags.map_reproducibility.internal_runs.dag_configs import DAG_CONFIGS_ULTRA
-from dags.map_reproducibility.utils.internal_aotc_workload import run_internal_aotc_workload
+from dags.map_reproducibility.utils.internal_aotc_workload import run_internal_dag_united_workload
 
 
 # Configuration parameters
@@ -38,7 +38,7 @@
     "reproducibility",
     "experimental",
     "xlml",
-    "v1.17",
+    "v2.0",
     "internal",
     "regressiontests",
     "a3ultra",
@@ -69,12 +69,13 @@
       start_date=datetime.datetime(2025, 4, 3),
       catchup=False,
   ) as dag:
-    run_internal_aotc_workload(
+    run_internal_dag_united_workload(
         relative_config_yaml_path=config_path,
         test_run=TEST_RUN,
         backfill=BACKFILL,
         timeout=timeout,
         image_version=NIGHTLY_IMAGE,
+        workload_launcher=WorkloadLauncher.MAXTEXT_LAUNCHER,
     )
 
   # Create DAG for stable release
@@ -86,10 +87,11 @@
       start_date=datetime.datetime(2025, 4, 3),
       catchup=False,
   ) as dag:
-    run_internal_aotc_workload(
+    run_internal_dag_united_workload(
         relative_config_yaml_path=config_path,
         test_run=TEST_RUN,
         backfill=BACKFILL,
         timeout=timeout,
         image_version=RELEASE_IMAGE,
+        workload_launcher=WorkloadLauncher.MAXTEXT_LAUNCHER,
     )

dags/map_reproducibility/internal_runs/a3ultra_nemo_benchmarking_dags.py

@@ -0,0 +1,77 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""DAGs to run Aotc reproducibility benchmarks."""
+
+import datetime
+import os
+
+from airflow import models
+from dags import composer_env
+from dags.map_reproducibility.utils.constants import Image, WorkloadLauncher
+from dags.map_reproducibility.internal_runs.dag_configs import DAG_CONFIGS_ULTRA_NEMO
+from dags.map_reproducibility.utils.internal_aotc_workload import run_internal_dag_united_workload
+
+
+# Configuration parameters
+TEST_RUN = False if composer_env.is_prod_env() else True
+BACKFILL = False
+
+# Get current date for image tags
+RELEASE_IMAGE = Image.NEMO_STABLE_RELEASE_A3U
+
+# Common DAG tags
+DAG_TAGS = [
+    "reproducibility",
+    "experimental",
+    "xlml",
+    "v2.0",
+    "internal",
+    "regressiontests",
+    "a3ultra",
+]
+
+
+# Create DAGs for each configuration
+for config_path, config_info in DAG_CONFIGS_ULTRA_NEMO.items():
+  # Extract config name for the DAG ID
+  config_name = os.path.basename(config_path).replace(".yaml", "")
+  schedule = config_info["schedule"] if not TEST_RUN else None
+  timeout = config_info["timeout_minutes"]
+  # Set retry parameter based on timeout
+  retries = 1 if timeout <= 15 else 0
+  retry_delay = datetime.timedelta(minutes=1)
+
+  dag_default_args = {
+      "retries": retries,
+      "retry_delay": retry_delay,
+  }
+
+  # Create DAG for stable release
+  with models.DAG(
+      dag_id=f"new_internal_stable_release_{config_name}",
+      default_args=dag_default_args,
+      schedule=schedule,
+      tags=DAG_TAGS,
+      start_date=datetime.datetime(2025, 4, 3),
+      catchup=False,
+  ) as dag:
+    run_internal_dag_united_workload(
+        relative_config_yaml_path=config_path,
+        test_run=TEST_RUN,
+        backfill=BACKFILL,
+        timeout=timeout,
+        image_version=RELEASE_IMAGE,
+        workload_launcher=WorkloadLauncher.NEMO_TEN_LAUNCHER,
+    )

dags/map_reproducibility/internal_runs/dag_configs.py

@@ -225,18 +225,16 @@
 }
 
 DAG_CONFIGS_ULTRA_NEMO = {
-    "recipes/a3ultra/nemo/a3ultra_llama3.1-70b_256gpus_fp8_nemo.yaml": {
-        "timeout_minutes": 25,
-        "backfill_group_nightly": 1,
-        "backfill_group_release": 1,
-        "nightly_schedule": Schedule.SATURDAY_PDT_12AM,
+    "recipes/a3ultra/nemo/a3ultra_llama3.1-8b_8gpus_fp8_nemo.yaml": {
+        "timeout_minutes": 20,
         "release_schedule": Schedule.SATURDAY_PDT_12AM,
     },
     "recipes/a3ultra/nemo/a3ultra_llama3.1-8b_8gpus_bf16_nemo.yaml": {
         "timeout_minutes": 20,
-        "backfill_group_nightly": 1,
-        "backfill_group_release": 1,
-        "nightly_schedule": Schedule.SATURDAY_PDT_12AM,
         "release_schedule": Schedule.SATURDAY_PDT_12AM,
     },
+    "recipes/a3ultra/nemo/a3ultra_llama3.1-70b_256gpus_fp8_nemo.yaml": {
+        "timeout_minutes": 25,
+        "release_schedule": Schedule.SATURDAY_PDT_3AM,
+    },
 }

dags/map_reproducibility/internal_runs/sample_a3ultra_maxtext_single_run.py

@@ -84,7 +84,7 @@ def main():
       base_helm_repo_root=base_helm_repo_root,
       timeout=timeout,
       image_version=RELEASE_IMAGE,
-      sample_run_bucket_name=SAMPLE_RUN_BUCKET_NAME,
+      gcs_bucket_name=SAMPLE_RUN_BUCKET_NAME,
       workload_launcher=WorkloadLauncher.MAXTEXT_LAUNCHER,
   )
 

dags/map_reproducibility/internal_runs/sample_a3ultra_nemo_single_run.py

@@ -82,7 +82,7 @@ def main():
 
   # Setup configuration
   relative_config_yaml_path = (
-      "recipes/a3ultra/nemo/a3ultra_llama3.1-8b_8gpus_bf16_nemo.yaml"
+      "recipes/a3ultra/nemo/a3ultra_llama3.1-8b_8gpus_fp8_nemo.yaml"
   )
   timeout = DAG_CONFIGS_ULTRA_NEMO[relative_config_yaml_path]["timeout_minutes"]
 
@@ -92,7 +92,7 @@ def main():
       base_helm_repo_root=base_helm_repo_root,
       timeout=timeout,
       image_version=RELEASE_IMAGE,
-      sample_run_bucket_name=SAMPLE_RUN_BUCKET_NAME,
+      gcs_bucket_name=SAMPLE_RUN_BUCKET_NAME,
       workload_launcher=WorkloadLauncher.NEMO_TEN_LAUNCHER,
   )
 

dags/map_reproducibility/utils/common_utils.py

@@ -411,12 +411,8 @@ def helm_apply_cmds_internal_run(
     bucket_name=BUCKET_NAME,
 ):
   gcs_cmd = ""
-  if hypercomputer in ("a3ultra", "a4"):
-    if framework != "maxtext" and kueue_name:
-      gcs_cmd = f" --set queue={kueue_name}"
-    gcs_cmd += f" --set volumes.gcsMounts[0].bucketName={bucket_name}"
-  else:
-    gcs_cmd = f" --set workload.gcsBucketForDataCataPath={bucket_name}"
+  if framework == "maxtext":
+    gcs_cmd += f" --set volumes.gcsMounts[0].bucketName={bucket_name} "
 
   cluster_cmd = ""
   if framework == "nemo" and hypercomputer == "a3ultra":

dags/map_reproducibility/utils/constants.py

@@ -1,6 +1,7 @@
 KUEUE_NAME = "a3-ultra"
 OPTIMIZER = "adam"
 NUM_STEPS = 15
+BUCKET_NAME = "regression-testing-xlml"
 
 
 class Optimizer:
@@ -39,6 +40,7 @@ class Schedule:
   WEEKDAY_PDT_6AM_7AM_EXCEPT_THURSDAY = "0 13,14 * * 2,3,4,6"
 
   SATURDAY_PDT_12AM = "0 7 * * 6"
+  SATURDAY_PDT_3AM = "0 10 * * 6"
 
 
 class Image:

dags/map_reproducibility/utils/internal_aotc_workload.py

@@ -61,7 +61,7 @@
     get_internal_run_type_and_comment,
 )
 from dags.map_reproducibility.utils.sample_workload_utils import handle_profiler, assemble_sample_united_workload_commands, execute_workload_commands
-from dags.map_reproducibility.utils.constants import Optimizer, KUEUE_NAME, NUM_STEPS
+from dags.map_reproducibility.utils.constants import Optimizer, KUEUE_NAME, NUM_STEPS, BUCKET_NAME
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -309,7 +309,7 @@ def run_internal_united_workload(
     base_helm_repo_root: str,
     timeout: int,
     image_version: str,
-    sample_run_bucket_name: str,
+    gcs_bucket_name: str,
     workload_launcher: str,
     is_dag_run: bool = False,
     backfill: bool = False,
@@ -393,18 +393,18 @@ def run_internal_united_workload(
         base_helm_repo_root,
         full_config_path,
         values_file_path,
-        sample_run_bucket_name,
+        gcs_bucket_name,
         container_timeout,
         tmpdir,
     )
 
-    success, error = execute_workload_commands(commands, tmpdir)
+    success, error = execute_workload_commands(commands, cwd=tmpdir)
     if not success:
       return {"success": False, "error": error}
 
     gcs_bucket = get_job_gcs_bucket_folder(
         job_name,
-        bucket_name=sample_run_bucket_name,
+        bucket_name=gcs_bucket_name,
         framework=config.FRAMEWORK,
         gcs_experiment_folder_name="maxtext-experiments",
     )
@@ -464,7 +464,7 @@ def run_internal_dag_united_workload(
       base_helm_repo_root=None,
       timeout=timeout,
       image_version=image_version,
-      sample_run_bucket_name=None,
+      gcs_bucket_name=BUCKET_NAME,
       workload_launcher=workload_launcher,
       is_dag_run=True,
       backfill=backfill,

dags/map_reproducibility/utils/sample_workload_utils.py

@@ -48,6 +48,7 @@
     wait_for_jobsets_cmds,
     cleanup_existing_metrics_cmd,
     helm_apply_cmds_workload,
+    get_values_file_path,
 )
 
 from dags.map_reproducibility.utils.benchmarkdb_utils import write_run

dags/multipod/maxtext_multi_tier_checkpointing_recover.py

@@ -0,0 +1,95 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+A DAG to run MaxText multi-tier checkpointing tests.
+"""
+import datetime
+from airflow import models
+from dags import composer_env, gcs_bucket
+from dags.common.vm_resource import DockerImage, XpkClusters
+from dags.multipod.configs import gke_config
+from dags.multipod.configs.common import SetupMode  # Run once a day at 10 am UTC (2 am PST)
+from dags.common import test_owner
+
+SCHEDULED_TIME = "0 10 * * *" if composer_env.is_prod_env() else None
+
+with models.DAG(
+    dag_id="maxtext_multi_tier_checkpointing_recover",
+    schedule=SCHEDULED_TIME,
+    tags=[
+        "multipod_team",
+        "maxtext",
+        "multi_tier_checkpointing_recover",
+    ],
+    start_date=datetime.datetime(2025, 5, 15),
+    catchup=False,
+    concurrency=2,
+) as dag:
+  base_output_directory = (
+      f"{gcs_bucket.MTC_BUCKET}/maxtext_multi_tier_checkpointing_recover"
+  )
+  dataset_path = gcs_bucket.MTC_BUCKET
+  docker_images = [
+      (SetupMode.STABLE, DockerImage.MAXTEXT_TPU_JAX_NIGHTLY),
+  ]
+  test_configs = {
+      # accelerator: list of slices to test
+      "v5p-8": [2],
+  }
+  clusters = {
+      # accelerator: cluster name
+      "v5p-8": XpkClusters.TPU_V5P_8_CLUSTER,
+  }
+
+  for mode, image in docker_images:
+    for accelerator, slices in test_configs.items():
+      for slice_num in slices:
+        command = (
+            "bash end_to_end/test_mtc_phase_2_save_path.sh"
+            f" multitiercheckpointing-{slice_num}x-{accelerator}"
+            f" {base_output_directory} {dataset_path}",
+        )
+        maxtext_save_checkpoint = gke_config.get_gke_config(
+            num_slices=slice_num,
+            cluster=clusters[accelerator],
+            time_out_in_min=60,
+            test_name="maxtext-multi-tier-checkpointing-recover",
+            run_model_cmds=command,
+            docker_image=image.value,
+            test_owner=test_owner.JACKY_F,
+        ).run_with_interruption(
+            ramdisk_directory="local",
+            xpk_branch="main",
+            skip_post_process=True,
+            mtc_enabled=True,
+        )
+
+        clean_cmd = (f"rm -rf /local/*",)
+        clean_ramdisk_one = gke_config.get_gke_config(
+            num_slices=slice_num,
+            cluster=clusters[accelerator],
+            time_out_in_min=60,
+            test_name="clean-ramdisk",
+            run_model_cmds=clean_cmd,
+            docker_image=image.value,
+            test_owner=test_owner.JACKY_F,
+        ).run(
+            ramdisk_directory="local",
+            xpk_branch="main",
+            skip_post_process=True,
+            mtc_enabled=True,
+        )
+
+        (maxtext_save_checkpoint >> clean_ramdisk_one)