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,21 @@ 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(
+            init_delete_cpc = checkpoint_util.initiate_cpc_deletion(
+                test_config.cpc_config
+            )
+            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 +283,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,

…orbax/util/multi_tier_checkpoint_util.py dags/orbax/util/checkpoint_util.pydags/orbax/util/multi_tier_checkpoint_util.py renamed to dags/orbax/util/checkpoint_util.py dags/orbax/util/multi_tier_checkpoint_util.py renamed to dags/orbax/util/checkpoint_util.py

@@ -12,6 +12,7 @@
 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
 
@@ -27,7 +28,6 @@ class CheckpointConfiguration:
   machine_type: str
   ramdisk_memory: str
   toleration_key: str
-  cpc_yaml_template: str
 
   def __init__(
       self,
@@ -50,7 +50,7 @@ def __init__(
       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., "2G" or "2048M". Defaults to "100G"".
+        string with the unit, e.g., "1G" or "1024Mi".
       toleration_key (str): The toleration key for the Kubernetes pod.
         Defaults to "google.com/tpu".
     """
@@ -61,21 +61,49 @@ def __init__(
     self.machine_type = machine_type
     self.ramdisk_memory = ramdisk_memory_in_mi
     self.toleration_key = toleration_key
-    self.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}
-  """
+
+  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."""
@@ -89,14 +117,9 @@ def create_custom_objects_api_client(self) -> k8s_client.CustomObjectsApi:
 @task
 def apply_cpc(cpc: CheckpointConfiguration) -> None:
   """Applies the CheckpointConfiguration to the cluster (create or replace)."""
-  custom_api = cpc.create_custom_objects_api_client()
-  cpc_body = yaml.safe_load(cpc.cpc_yaml_template)
-  api_version = cpc_body.get("apiVersion")
-  kind = cpc_body.get("kind")
-  name = cpc_body.get("metadata", {}).get("name")
 
-  group, version = api_version.split("/", 1)
-  plural = f"{kind.lower()}s"
+  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
@@ -108,8 +131,7 @@ def apply_cpc(cpc: CheckpointConfiguration) -> None:
 
   except ApiException as api_error:
     # If it already exists (409 Conflict), then try to replace it
-    # See https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/409
-    if api_error.status == 409:
+    if api_error.status == HTTPStatus.CONFLICT:
       logging.info(
           f"CheckpointConfiguration '{name}' already exists. Attempting to replace..."
       )
@@ -142,43 +164,31 @@ def initiate_cpc_deletion(cpc: CheckpointConfiguration) -> None:
   Sends the delete request for the CheckpointConfiguration.
   """
   custom_api = cpc.create_custom_objects_api_client()
-  cpc_body = yaml.safe_load(cpc.cpc_yaml_template)
-  name_to_delete = cpc_body.get("metadata", {}).get("name")
+  group, version, plural, name, _ = cpc.load_yaml_and_parse_body()
 
-  if not name_to_delete:
+  if not name:
     logging.error(
         "Could not determine CheckpointConfiguration name for deletion."
     )
     raise AirflowFailException("Failed to determine CPC name for deletion.")
 
-  api_version = cpc_body.get("apiVersion")
-  kind = cpc_body.get("kind")
-  group, version = api_version.split("/", 1)
-  plural = f"{kind.lower()}s"
-
   delete_options = k8s_client.V1DeleteOptions(propagation_policy="Foreground")
-
   try:
-    logging.info(
-        f"Attempting to delete CheckpointConfiguration: {name_to_delete}"
-    )
+    logging.info(f"Attempting to delete CheckpointConfiguration: {name}")
     custom_api.delete_cluster_custom_object(
         group=group,
         version=version,
         plural=plural,
-        name=name_to_delete,
+        name=name,
         body=delete_options,
     )
-    logging.info(
-        f"Delete request sent for CheckpointConfiguration '{name_to_delete}'."
-    )
+    logging.info(f"Delete request sent for CheckpointConfiguration '{name}'.")
 
   except ApiException as e:
     # The resource is already gone (404), so we can exit successfully
-    # See https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/409
-    if e.status == 404:
+    if e.status == HTTPStatus.NOT_FOUND:
       logging.info(
-          f"CheckpointConfiguration '{name_to_delete}' not found. Already deleted or never existed."
+          f"CheckpointConfiguration '{name}' not found. Already deleted or never existed."
       )
       return
     else:
@@ -193,36 +203,28 @@ def wait_for_cpc_deletion(cpc: CheckpointConfiguration) -> bool:
   A sensor that waits for the CheckpointConfiguration to be completely deleted.
   """
   custom_api = cpc.create_custom_objects_api_client()
-  cpc_body = yaml.safe_load(cpc.cpc_yaml_template)
-  name_to_delete = cpc_body.get("metadata", {}).get("name")
+  group, version, plural, name, _ = cpc.load_yaml_and_parse_body()
 
-  if not name_to_delete:
+  if not name:
     logging.error(
         "Could not determine CheckpointConfiguration name for deletion."
     )
     raise AirflowFailException("Failed to determine CPC name for deletion.")
 
-  api_version = cpc_body.get("apiVersion")
-  kind = cpc_body.get("kind")
-  group, version = api_version.split("/", 1)
-  plural = f"{kind.lower()}s"
-
   try:
     custom_api.get_cluster_custom_object(
-        group=group, version=version, plural=plural, name=name_to_delete
+        group=group, version=version, plural=plural, name=name
     )
     logging.info(
-        f"CheckpointConfiguration '{name_to_delete}' still exists. "
+        f"CheckpointConfiguration '{name}' still exists. "
         f"Polling again in 10s..."
     )
     return False
   except ApiException as e:
     # The resource is already gone (404), so we can exit successfully
     # See https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/409
     if e.status == 404:
-      logging.info(
-          f"CheckpointConfiguration: {name_to_delete} successfully deleted"
-      )
+      logging.info(f"CheckpointConfiguration: {name} successfully deleted")
       return True  # Return True to tell the sensor to succeed
     else:
       logging.error(