Fix: Change name logging_mtc.py to validate_util.py. Since validate_util as the filename make more sense, we are validating the logs.

Fix: checkpoint_steps --> to Optional. Since this value exist but is not use in this test.

Fix: Change name mx-sv --> max-sv-loc. Helps with clarity in bucket checkpoints.

Fixes for code review

Author: alfredyu-cienet

dags/orbax/maxtext_mtc_emergency_save_local.py

@@ -10,6 +10,7 @@
 import datetime
 from dataclasses import dataclass
 import posixpath
+from typing import Optional
 
 from airflow import models
 from airflow.utils.task_group import TaskGroup
@@ -19,16 +20,16 @@
 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 dags.orbax.util import validation_util
+from dags.orbax.util import checkpoint_util
 from xlml.utils.xpk import BRANCH_ABHINAV_MTC
 from xlml.utils.gke import zone_to_region
 
 SCHEDULE = "0 10 * * *" if composer_env.is_prod_env() else None
 DAG_TEST_NAME = "maxtext_emc_and_mtc_orbax_save_local"
-BASE_OUTPUT_DIRECTORY = gcs_bucket.MTC_AUTOMATION_BUCKET
+BASE_OUTPUT_DIR = gcs_bucket.MTC_AUTOMATION_BUCKET
 
-# For this DAG, only one version of the Docker image is supported at the moment.
+# Only one version of the Docker image is supported at the moment.
 # Other versions (e.g., "stable") may be introduced later.
 DOCKER_IMAGES = [(
     SetupMode.NIGHTLY,
@@ -66,7 +67,7 @@ class TestConfig:
   local_checkpoint_step: int
   checkpoint_step: int
   ram_disk_size: str
-  cpc_config: mtc.CheckpointConfiguration
+  cpc_config: checkpoint_util.CheckpointConfiguration
 
   def __init__(
       self,
@@ -79,8 +80,8 @@ def __init__(
       replicator_backup_time: int,
       step: int,
       local_checkpoint_step: int,
-      checkpoint_step: int,
       ram_disk_size_in_mi: str,
+      checkpoint_step: Optional[int] = None,
   ):
     """
     Initializes the test configurations.
@@ -114,7 +115,7 @@ def __init__(
     self.local_checkpoint_step = local_checkpoint_step
     self.checkpoint_step = checkpoint_step
     self.ram_disk_size = ram_disk_size_in_mi
-    self.cpc_config = mtc.CheckpointConfiguration(
+    self.cpc_config = checkpoint_util.CheckpointConfiguration(
         project_id=self.cluster.project,
         region=zone_to_region(self.cluster.zone),
         cluster_name=self.cluster.name,
@@ -127,7 +128,7 @@ def generate_workload_command(
       self,
       cp: Checkpointing,
       checkpoint_dir: str,
-      out_dir: str,
+      out_folder: str,
       slice_number: int,
   ) -> str:
     run_time = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")
@@ -138,7 +139,7 @@ def generate_workload_command(
         "python3 -m MaxText.train MaxText/configs/base.yml "
         "remat_policy=full "
         "global_parameter_scale=1 "
-        f"base_output_directory={out_dir} "
+        f"base_output_directory={posixpath.join(BASE_OUTPUT_DIR, out_folder)} "
         "dataset_type=synthetic "
         f"steps={self.step} "
         "per_device_batch_size=1 "
@@ -169,21 +170,49 @@ def generate_workload_command(
         "orbax",
     ],
     description="A DAG to run MaxText multi-tier checkpointing tests.",
-    doc_md="",
+    doc_md="""
+      # Emergency Checkpoint Manager and Multi-tier Checkpoint Validation DAG
+
+      ### Description
+      This DAG (Directed Acyclic Graph) automates the process of validating
+      checkpoint saving when using both **Emergency Checkpoint Manager**
+      and **Multi-tier Checkpoint** features. The flag that controls the
+      behaviour of using MTC or EMC is **user_replicatior**.
+      Also the steps flag controls how many steps the job will run.
+
+      ### Prerequisites
+      To run this test, you need an existing cluster with the Multi-tier
+      Checkpointing configuration enabled, as well as a bucket with HNS
+      (Hierarchical Namespace) enabled.
+
+      ### Procedures
+      1.  **Apply Configuration:** A Checkpoint Configuration YAML file is
+      applied to the cluster, enabling Multi-tier Checkpoint (MTC) features.
+      2.  **Run Maxtext Jobsets:** The DAG runs a Maxtext jobset twice.
+      One run has `replicator_enabled` set to `True` (for MTC), and the
+      other has it set to `False` (for Emergency Checkpoint Manager).
+      3.  **Validate Checkpoints:** The DAG validates that **local checkpoints**
+      are being saved correctly in the `local/` directory for both job runs.
+
+      4.  The validation logic is the same for both the MTC and Emergency
+      Checkpoint Manager job runs because their local checkpoint saving
+      behavior is similar. This is why a single validation task is used for both.
+    """,
     concurrency=2,
 ) as dag:
+  # Only one set of test configurations (e.g., v5p-128) is supported at the moment.
+  # Other configurations (e.g., v5e and/or v6e) may be introduced later.
   test_configs = [
       TestConfig(
           cluster=XpkClusters.TPU_V5P_128_CLUSTER_ORBAX,
           machine_type="ct5p-hightpu-4t",
           accelerator="v5p-128",
           slices=[2],
           model_name="llama2-7b",
-          short_id="max-sv",
+          short_id="max-sv-loc",
           replicator_backup_time=30,
           step=100,
           local_checkpoint_step=20,
-          checkpoint_step=25,
           ram_disk_size_in_mi="800000Mi",
       ),
   ]
@@ -200,9 +229,6 @@ def generate_workload_command(
           use_replicator=False,
       ),
   ]:
-    folder = f"maxtext_{checkpointing.name}_orbax_save_local"
-    output_dir = posixpath.join(BASE_OUTPUT_DIRECTORY, folder)
-
     with TaskGroup(
         group_id=f"maxtext_{checkpointing.name}_orbax_save_local",
     ) as group:
@@ -211,19 +237,18 @@ def generate_workload_command(
           for slice_num in test_config.slices:
             # 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(test_config.cpc_config)
-            wait_delete_cpc = mtc.wait_for_cpc_deletion.override(
+            wait_delete_cpc = checkpoint_util.wait_for_cpc_deletion.override(
                 trigger_rule="all_done"
             )(test_config.cpc_config)
-            apply_cpc = mtc.apply_cpc(test_config.cpc_config)
+            apply_cpc = checkpoint_util.apply_cpc(test_config.cpc_config)
             workload_command = test_config.generate_workload_command(
                 cp=checkpointing,
                 checkpoint_dir=RAM_DISK,
-                out_dir=output_dir,
+                out_folder=group.group_id,
                 slice_number=slice_num,
             )
 
-            start_time = log.generate_timestamp()
+            start_time = validation_util.generate_timestamp()
             maxtext_chkpt_run_test = gke_config.get_gke_config(
                 num_slices=slice_num,
                 cluster=test_config.cluster,
@@ -255,20 +280,15 @@ def generate_workload_command(
                 skip_post_process=True,
             )
 
-            end_time = log.generate_timestamp()
+            end_time = validation_util.generate_timestamp()
             vali_step = test_config.step - 1
             vali_step_list = [
                 i
                 for i in range(0, vali_step, test_config.local_checkpoint_step)
             ]
             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(
+            validate_log = validation_util.validate_log_with_step(
                 project_id=test_config.cluster.project,
                 location=zone_to_region(test_config.cluster.zone),
                 cluster_name=test_config.cluster.name,
@@ -279,8 +299,7 @@ def generate_workload_command(
             )
 
             (
-                init_delete_cpc
-                >> wait_delete_cpc
+                wait_delete_cpc
                 >> apply_cpc
                 >> start_time
                 >> maxtext_chkpt_run_test

dags/orbax/util/checkpoint_util.py

@@ -0,0 +1,183 @@
+"""Utility functions for managing Multi-tier Cluster Configuration.
+
+This module provides tasks for creating, applying, and deleting
+Multi-tier Driver cluster Configurations for enable Multi Tier Checkpointing.
+"""
+
+from absl import logging
+import yaml
+from dataclasses import dataclass
+
+from kubernetes import client as k8s_client
+from kubernetes.client.rest import ApiException
+from airflow.decorators import task
+from airflow.exceptions import AirflowFailException
+from http import HTTPStatus
+
+from xlml.utils import gke
+
+
+@dataclass
+class CheckpointConfiguration:
+  """A dataclass to hold attributes of a Cloud Public Compute (CPC) instance."""
+
+  project_id: str
+  region: str
+  cluster_name: str
+  gcs_bucket: str
+  machine_type: str
+  ramdisk_memory: str
+  toleration_key: str
+
+  def __init__(
+      self,
+      project_id: str,
+      region: str,
+      cluster_name: str,
+      gcs_bucket: str,
+      machine_type: str,
+      ramdisk_memory_in_mi: str,
+      toleration_key: str = "google.com/tpu",
+  ):
+    """
+    Initializes the CheckpointConfiguration.
+
+    Args:
+      project_id (str): The Google Cloud project ID.
+      region (str): The Google Cloud region.
+      cluster_name (str): The name of the GKE cluster.
+      gcs_bucket (str): The name of the GCS bucket for checkpoints.
+      machine_type (str): The machine type for the instance.
+      ram_disk_memory_in_mi (str): The size of the RAM disk in mebibytes (Mi).
+        The unit is in mebibytes (Mi) but the value should be passed as a
+        string with the unit, e.g., "1G" or "1024Mi".
+      toleration_key (str): The toleration key for the Kubernetes pod.
+        Defaults to "google.com/tpu".
+    """
+    self.project_id = project_id
+    self.region = region
+    self.cluster_name = cluster_name
+    self.gcs_bucket = gcs_bucket
+    self.machine_type = machine_type
+    self.ramdisk_memory = ramdisk_memory_in_mi
+    self.toleration_key = toleration_key
+
+  def load_yaml_and_parse_body(
+      self,
+  ) -> tuple[str, str, str, str, dict[str, any]]:
+    """
+    Loads a YAML string template, populates it with class attributes, and parses the resulting body.
+
+    This method constructs a CheckpointConfiguration YAML manifest as a string,
+    using class attributes such as `self.gcs_bucket`, `self.machine_type`,
+    `self.toleration_key`, and `self.ramdisk_memory` to fill in the
+    placeholders. It then uses `yaml.safe_load` to convert this YAML string
+    into a Python dictionary.
+
+    Finally, it extracts key fields—group, version, plural, and name—from the
+    loaded dictionary for use in API requests or other operations.
+
+    Returns:
+      tuple[str, str, str, str, dict[str, any]]: A tuple containing the
+      extracted API group, API version, plural name, resource name, and the
+      full parsed YAML body as a dictionary.
+    """
+
+    cpc_yaml_template = f"""
+      apiVersion: checkpointing.gke.io/v1
+      kind: CheckpointConfiguration
+      metadata:
+        name: my-checkpointconfiguration # This name will be used for deletion
+      spec:
+        cloudStorageBucketName: {self.gcs_bucket}
+        nodeSelector:
+          node.kubernetes.io/instance-type: {self.machine_type}
+        tolerations:
+        - key: {self.toleration_key}
+          operator: Exists
+          effect: NoSchedule
+        inMemoryVolumeSize: {self.ramdisk_memory}
+    """
+    cpc_body = yaml.safe_load(cpc_yaml_template)
+    group, version = cpc_body.get("apiVersion").split("/", 1)
+    plural = cpc_body.get("kind").lower() + "s"
+    name = cpc_body.get("metadata", {}).get("name")
+
+    return group, version, plural, name, cpc_body
+
+  def create_custom_objects_api_client(self) -> k8s_client.CustomObjectsApi:
+    """Create a CustomObjectsApi client for the given cluster."""
+    return k8s_client.CustomObjectsApi(
+        gke.get_authenticated_client(
+            self.project_id, self.region, self.cluster_name
+        )
+    )
+
+
+@task
+def apply_cpc(cpc: CheckpointConfiguration) -> None:
+  """Applies the CheckpointConfiguration to the cluster (create or replace)."""
+
+  custom_api = cpc.create_custom_objects_api_client()
+  group, version, plural, name, cpc_body = cpc.load_yaml_and_parse_body()
+
+  try:
+    # Here we first try to create a reasource
+    logging.info(f"Attempting to create CheckpointConfiguration '{name}'...")
+    custom_api.create_cluster_custom_object(
+        group=group, version=version, plural=plural, body=cpc_body
+    )
+    logging.info(f"CheckpointConfiguration '{name}' created successfully.")
+
+  except ApiException as api_error:
+    raise AirflowFailException(
+        f"Failed to apply CheckpointConfiguration: {api_error.reason}"
+    )
+
+
+def _delete_cpc(cpc: CheckpointConfiguration) -> bool:
+  """
+  Sends the delete request for the CheckpointConfiguration.
+  """
+
+  custom_api = cpc.create_custom_objects_api_client()
+  group, version, plural, name, _ = cpc.load_yaml_and_parse_body()
+
+  if not name:
+    logging.error(
+        "Could not determine CheckpointConfiguration name for deletion."
+    )
+    raise AirflowFailException("Failed to determine CPC name for deletion.")
+
+  delete_options = k8s_client.V1DeleteOptions(propagation_policy="Foreground")
+
+  try:
+    logging.info(f"Attempting to delete CheckpointConfiguration: {name}")
+    custom_api.delete_cluster_custom_object(
+        group=group,
+        version=version,
+        plural=plural,
+        name=name,
+        body=delete_options,
+    )
+    logging.info(f"Delete request sent for CheckpointConfiguration '{name}'.")
+  except Exception as e:
+    logging.info(
+        f"An warning has occurred while deleting CPC. Please take a look: {e}"
+    )
+    pass
+
+  try:
+    custom_api.get_cluster_custom_object(
+        group=group, version=version, plural=plural, name=name
+    )
+  except ApiException as e:
+    return e.status == HTTPStatus.NOT_FOUND
+
+
+@task.sensor(poke_interval=10)
+def wait_for_cpc_deletion(cpc: CheckpointConfiguration) -> bool:
+  """
+  A sensor that waits for the CheckpointConfiguration to be completely deleted.
+  """
+  return _delete_cpc(cpc)