add maxtext_multi_tier_checkpointing_recover DAG

severus-ho · severus-ho · commit 12251fe13818 · 2025-05-20T11:04:23.000+08:00
diff --git a/dags/common/test_owner.py b/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."
diff --git a/dags/multipod/maxtext_multi_tier_checkpointing_recover.py b/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)
diff --git a/xlml/apis/task.py b/xlml/apis/task.py
@@ -201,6 +201,41 @@ def run(
 
     return group
 
+  def run_with_interruption(
+      self,
+      *,
+      gcs_location: Optional[airflow.XComArg] = None,
+      use_vertex_tensorboard: bool = False,
+      use_pathways: bool = False,
+      skip_post_process: bool = False,
+      ramdisk_directory: str = "",
+      mtc_enabled: bool = False,
+      xpk_branch: str = xpk.MAIN_BRANCH,
+  ) -> DAGNode:
+    """Run a test job within a docker image.
+
+    Attributes:
+      gcs_location: GCS path for all artifacts of the test.
+      use_vertex_tensorboard: Set to True to view workload data on
+        Vertex AI Tensorboard.
+
+    Returns:
+      A task group with the following tasks chained: run_model and
+      post_process.
+    """
+    with TaskGroup(group_id=self.task_test_config.benchmark_id) as group:
+      run_model, gcs_path = self.run_model_with_interruption(
+          gcs_location,
+          use_vertex_tensorboard,
+          use_pathways,
+          ramdisk_directory,
+          mtc_enabled,
+          xpk_branch,
+      )
+      if not skip_post_process:
+        run_model >> self.post_process(gcs_path)
+    return group
+
   def run_with_name_gen_and_quarantine(
       self,
       quarantine_task_group,
@@ -333,6 +368,72 @@ def run_model(
       )
       return group, gcs_path
 
+  def run_model_with_interruption(
+      self,
+      gcs_location: Optional[airflow.XComArg] = None,
+      use_vertex_tensorboard: bool = False,
+      use_pathways: bool = False,
+      ramdisk_directory: str = "",
+      mtc_enabled: bool = False,
+      xpk_branch: str = xpk.MAIN_BRANCH,
+  ) -> DAGNode:
+    """Run the TPU/GPU test in `task_test_config` using xpk.
+      Different behaviour for testing interruption.
+
+    Attributes:
+      gcs_location: GCS path for all artifacts of the test.
+      use_vertex_tensorboard: Set to True to view workload data on
+        Vertex AI Tensorboard.
+
+    Returns:
+      A DAG node that executes the model test.
+    """
+    with TaskGroup(group_id="run_model") as group:
+      workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id)
+      if gcs_location:
+        gcs_path = gcs_location
+      else:
+        gcs_path = name_format.generate_gcs_folder_location(
+            self.task_test_config.gcs_subfolder,
+            self.task_test_config.benchmark_id,
+        )
+
+      launch_workload_with_interruption = (
+          self.launch_workload_with_interruption(
+              workload_id,
+              gcs_path,
+              use_vertex_tensorboard,
+              use_pathways,
+              ramdisk_directory,
+              mtc_enabled,
+              xpk_branch,
+          )
+      )
+
+      wait_for_workload_completion = xpk.wait_for_workload_completion.override(
+          timeout=int(self.task_test_config.timeout.total_seconds()),
+      )(
+          workload_id=workload_id,
+          project_id=self.task_gcp_config.project_name,
+          region=self.task_gcp_config.zone[:-2],
+          cluster_name=self.task_test_config.cluster_name,
+      )
+
+      clean_up_workload = xpk.clean_up_workload(
+          workload_id=workload_id,
+          project_id=self.task_gcp_config.project_name,
+          zone=self.task_gcp_config.zone,
+          cluster_name=self.task_test_config.cluster_name,
+      )
+
+      (
+          (workload_id, gcs_path)
+          >> launch_workload_with_interruption
+          >> wait_for_workload_completion
+          >> clean_up_workload
+      )
+    return group, gcs_path
+
   def launch_workload(
       self,
       workload_id: str,
@@ -376,6 +477,61 @@ def launch_workload(
       run_workload >> wait_for_workload_start
       return group
 
+  def launch_workload_with_interruption(
+      self,
+      workload_id: str,
+      gcs_path: str,
+      use_vertex_tensorboard: bool,
+      use_pathways: bool = False,
+      ramdisk_directory: str = "",
+      mtc_enabled: bool = False,
+      xpk_branch: str = xpk.MAIN_BRANCH,
+  ) -> DAGNode:
+    """Create the workload and wait for it to provision."""
+    with TaskGroup(group_id="launch_workload_with_interruption") as group:
+      run_workload = xpk.run_workload.override(
+          owner=self.task_test_config.task_owner
+      )(
+          task_id="run_workload",
+          cluster_project=self.task_gcp_config.project_name,
+          zone=self.task_gcp_config.zone,
+          cluster_name=self.task_test_config.cluster_name,
+          benchmark_id=self.task_test_config.benchmark_id,
+          workload_id=workload_id,
+          gcs_path=gcs_path,
+          docker_image=self.task_test_config.docker_image,
+          accelerator_type=self.task_test_config.accelerator.name,
+          run_cmds=self.task_test_config.test_script,
+          num_slices=self.task_test_config.num_slices,
+          use_vertex_tensorboard=use_vertex_tensorboard,
+          use_pathways=use_pathways,
+          ramdisk_directory=ramdisk_directory,
+          mtc_enabled=mtc_enabled,
+          xpk_branch=xpk_branch,
+      )
+
+      wait_for_workload_start = xpk.wait_for_workload_start.override(
+          timeout=self.workload_provision_timeout.total_seconds()
+      )(
+          workload_id=workload_id,
+          project_id=self.task_gcp_config.project_name,
+          region=self.task_gcp_config.zone[:-2],
+          cluster_name=self.task_test_config.cluster_name,
+      )
+
+      run_interruption_workload = xpk.run_interruption_cmd.override(
+          owner=self.task_test_config.task_owner
+      )(
+          task_id="run_interruption_cmd",
+          workload_id=workload_id,
+          project_id=self.task_gcp_config.project_name,
+          region=self.task_gcp_config.zone[:-2],
+          cluster_name=self.task_test_config.cluster_name,
+      )
+
+      run_workload >> wait_for_workload_start >> run_interruption_workload
+      return group
+
   def post_process(self, result_location: Optional[str] = None) -> DAGNode:
     """Process metrics and metadata, and insert them into BigQuery tables.
 
diff --git a/xlml/utils/xpk.py b/xlml/utils/xpk.py
@@ -120,6 +120,12 @@ def run_workload(
       workload_create_cmd += f" --ramdisk-directory={ramdisk_directory}"
     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:
@@ -284,6 +290,51 @@ def wait_for_workload_completion(
   return True
 
 
+@task.sensor(poke_interval=120, timeout=1200, mode="reschedule")
+def run_interruption_cmd(
+    task_id: str,
+    project_id: str,
+    region: str,
+    cluster_name: str,
+    workload_id: str,
+) -> bool:
+  """Run command to interrupt pod ."""
+  core_api = _get_core_api_client(project_id, region, cluster_name)
+  pods = _list_workload_pods(core_api, workload_id)
+
+  if any(pod.status.phase in ["Pending"] for pod in pods.items):
+    logging.info("Some of the pods is still pending. Waiting to start")
+    return False
+
+  try:
+    for pod in pods.items:
+      if pod.status.phase == "Failed":
+        # Don't keep retrying if the pod has failed
+        raise AirflowFailException(f"Bad pod phase: {pod.status.phase}")
+      elif pod.status.phase in ["Unknown"]:
+        raise RuntimeError(f"Bad pod phase: {pod.status.phase}")
+  finally:
+    if all(pod.status.phase in ["Running"] for pod in pods.items):
+      # Pick last one running pod
+      pod = pods.items[len(pods.items) - 1]
+      logs = core_api.read_namespaced_pod_log(
+          name=pod.metadata.name, namespace=pod.metadata.namespace
+      )
+      logging.info(f"Logs for pod {pod.metadata.name}:")
+      for line in logs.split("\n"):
+        logging.info(line)
+
+      # TODO --> More sophisticated way to know when the pod start training.
+      if "completed step:" in logs:
+        # Here where regex expresion to kill pod will go. First test with killing the  pod
+        result = core_api.delete_namespaced_pod(
+            name=pod.metadata.name, namespace=pod.metadata.namespace
+        )
+        logging.info("The {pod.metadata.name} pod was successfully deleted.")
+        return True
+  return False
+
+
 @task(trigger_rule="all_done")
 def clean_up_workload(
     workload_id: str,
