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

alfredyu-cienet · camiloCienet · commit 8d82f174adfa · 2025-08-25T06:46:33.000Z
Fix: checkpoint_steps --&gt; to Optional. Since this value exist but is not use in this test.

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

Fixes for code review
diff --git a/dags/orbax/maxtext_mtc_emergency_save_local.py b/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,14 +20,14 @@
 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 as orbax_logging_util
+from dags.orbax.util import checkpoint_util as mtc
 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.
 # Other versions (e.g., "stable") may be introduced later.
@@ -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.
@@ -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:
+  # TODO: In the near future this test configs will be add it in the
+  # params {} attributes inside models.DAG
   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,7 @@ def generate_workload_command(
           use_replicator=False,
       ),
   ]:
-    folder = f"maxtext_{checkpointing.name}_orbax_save_local"
-    output_dir = posixpath.join(BASE_OUTPUT_DIRECTORY, folder)
-
+    output_folder = f"maxtext_{checkpointing.name}_orbax_save_local"
     with TaskGroup(
         group_id=f"maxtext_{checkpointing.name}_orbax_save_local",
     ) as group:
@@ -219,11 +246,11 @@ def generate_workload_command(
             workload_command = test_config.generate_workload_command(
                 cp=checkpointing,
                 checkpoint_dir=RAM_DISK,
-                out_dir=output_dir,
+                out_folder=output_folder,
                 slice_number=slice_num,
             )
 
-            start_time = log.generate_timestamp()
+            start_time = orbax_logging_util.generate_timestamp()
             maxtext_chkpt_run_test = gke_config.get_gke_config(
                 num_slices=slice_num,
                 cluster=test_config.cluster,
@@ -255,20 +282,15 @@ def generate_workload_command(
                 skip_post_process=True,
             )
 
-            end_time = log.generate_timestamp()
+            end_time = orbax_logging_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 = orbax_logging_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,
diff --git a/dags/orbax/util/checkpoint_util.py b/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,33 @@ 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]]:
+    """doc"""
+
+    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."""
@@ -90,13 +102,7 @@ def create_custom_objects_api_client(self) -> k8s_client.CustomObjectsApi:
 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"
+  group, version, plural, name, cpc_body = cpc.load_yaml_and_parse_body()
 
   try:
     # Here we first try to create a reasource
@@ -108,8 +114,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 +147,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 +186,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(
diff --git a/dags/orbax/util/validation_util.py b/dags/orbax/util/validation_util.py
