Fix: Change cluster test to "ml-auto-solutions-orbax".

Fix: Refactor cpc dict into a CheckpointConfiguration class, that holds relevant info about cpc metadata.

Fix: Change name of DAG to be mtc and emc. mainly because both of these test are being tested.

Fix: Fix some try and catch exceptions that seems to be unnecessary.

Fix: Refactor whole Cpc delete task into two task initiate_cpc and wait_to_delete. wait_to_delete task is a sensor task that will check every 10 s if the delition already happened.

Fix: Two Dags testing mtc and emc (with replicator=true and replicator=false)

Fix: Change loggin.py library to logging_mtc.py names conflict.

Fix: Use abhinav-mtc xpk branch instead of main branch since xpk version > 0.4.1 crash airflow pod. b/437817546

Fix: Correct name of DAG to emtc (Emergency multitier Checkpointer).

Fix: Change name_id of the test to be related to Orbax. Create new tag with Orbax. Remove --restart-on-user-code-failure flag since is already outdated in both main and develop branch

Fix: Backup interval to longer. Only testing local saves. GCS testing is in DAG002

mend

Author: camiloCienet

dags/common/test_owner.py

@@ -65,7 +65,7 @@ class Team(enum.Enum):
 # Multi-tier Checkpointing
 ABHINAV_S = "abhinavclemson"
 XUEFENG_G = "xuefgu"
-CAMILO = "CAMILO P."
+CAMILO_Q = "camiloCienet"
 
 # MLCompass
 ORTI_B = "ortibazar"

dags/common/vm_resource.py

@@ -260,7 +260,7 @@ class XpkClusters:
       zone=Zone.EUROPE_WEST4_B.value,
   )
   TPU_V5P_128_CLUSTER_ORBAX = XpkClusterConfig(
-      name="orbax-ml-auto-solutions",
+      name="ml-auto-solutions-orbax",
       device_version=TpuVersion.V5P,
       core_count=128,
       project=Project.CLOUD_TPU_MULTIPOD_DEV.value,

…text_multi_tier_chechpoint_save_local.py …rbax/maxtext_mtc_emergency_save_local.pydags/orbax/maxtext_multi_tier_chechpoint_save_local.py renamed to dags/orbax/maxtext_mtc_emergency_save_local.py dags/orbax/maxtext_multi_tier_chechpoint_save_local.py renamed to dags/orbax/maxtext_mtc_emergency_save_local.py

@@ -23,53 +23,57 @@
 from xlml.utils.xpk import BRANCH_ABHINAV_MTC
 from xlml.utils.gke import zone_to_region
 
-# Global variables across test configurations.
 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
 
-# The variable DOCKER_IMAGES will hold multiple configurations e.g nightly , stable.
+# For this DAG, 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,
     DockerImage.MAXTEXT_TPU_JAX_NIGHTLY,
 )]
 RAM_DISK = "/local"
 
+
 @dataclass
 class Testcase:
-  """Holds a single test case configuration for multi-tier checkpointing.
+  """
+  Represents the information of a checkpointing mechanism.
 
-  This class is used to define the properties of a specific test run,
-  including the name of the checkpoint and whether to use a replicator.
+  Attributes:
+    name: A unique name for the checkpointing configuration.
+    use_replicator: Indicates whether a replicator is enabled for this test case.
   """
-  checkpointing_name: str
+
+  name: str
   use_replicator: bool
 
+
 @dataclass
 class Testconfig:
-  """Holds the general configuration for a multi-tier checkpointing test.
-
-  This class defines the environment and parameters for a single test run,
-  including cluster details, model settings, and checkpointing intervals.
-
-  cluster (XpkClusters): The cluster to be used for the test.
-  machine_type (str): The type of machine to run the test on.
-  accelerator (str): The type of accelerator (e.g., GPU, TPU) to use.
-  slices (list[int]): The number of slices to be used.
-  model_name (str): The name of the model being tested.
-  name_prefix (str): A prefix to be used for naming the test run.
-  replicator_min (int): The minimum number of replicators required.
-  step (int): The current step of the training process.
-  local_checkpoint_step (int): The step interval for local checkpoints.
-  checkpoint_step (int): The step interval for global checkpoints.
-  ram_disk_mi (str): Information about the RAM disk.
+  """Holds the general configuration for a checkpointing test.
+
+  Attributes:
+    cluster: The specified cluster to be used for the test.
+    machine_type: The type of machine (e.g., GPU, TPU).
+    accelerator: The type of accelerator (e.g., GPU, TPU) to use.
+    slices: The number of slices to be used.
+    model_name: The name of the model being tested.
+    short_id (str): A short id to be used for naming the test run.
+    replicator_min: The time the replicator takes to backup checkpoints to bucket.
+    step: The current step of the training process.
+    local_checkpoint_step: The step interval for local checkpoints.
+    checkpoint_step: The step interval for the checkpoints store in the bucket.
+    ram_disk_mi: The size about the RAM disk in the CSI driver, in Mi.
   """
+
   cluster: XpkClusters
   machine_type: str
   accelerator: str
   slices: list[int]
   model_name: str
-  name_prefix: str
+  short_id: str
   replicator_min: int
   step: int
   local_checkpoint_step: int
@@ -102,7 +106,7 @@ class Testconfig:
           accelerator="v5p-128",
           slices=[2],
           model_name="llama2-7b",
-          name_prefix="max-sv",
+          short_id="max-sv",
           replicator_min=30,
           step=100,
           local_checkpoint_step=20,
@@ -113,26 +117,23 @@ class Testconfig:
   tests_to_run_seq = []
   # Individual test cases for Multi-tier Checkpointing and  Emergency Checkpointing
   for tc in [
-      Testcase(checkpointing_name="mtc", use_replicator=True),
-      Testcase(checkpointing_name="emc", use_replicator=False),
+      Testcase(name="mtc", use_replicator=True),
+      Testcase(name="emc", use_replicator=False),
   ]:
-    folder = f"maxtext_{tc.checkpointing_name}_orbax_save_local"
-    BASE_OUTPUT_DIRECTORY = posixpath.join(BASE_OUTPUT_DIRECTORY, folder)
-    with TaskGroup(
-        group_id=f"maxtext_{tc.checkpointing_name}_orbax_save_local"
-    ) as task:
+    folder = f"maxtext_{tc.name}_orbax_save_local"
+    run_dir_out = posixpath.join(BASE_OUTPUT_DIRECTORY, folder)
+    with TaskGroup(group_id=f"maxtext_{tc.name}_orbax_save_local") as task:
       for mode, image in DOCKER_IMAGES:
         for test_config in test_configs:
           cpc_conf = mtc.CheckpointConfiguration(
               project_id=test_config.cluster.project,
               region=zone_to_region(test_config.cluster.zone),
               cluster_name=test_config.cluster.name,
               gcs_bucket=gcs_bucket.MTC_AUTOMATION_BUCKET.removeprefix("gs://"),
-              ramdisk_memory=test_config.ram_disk_mi,
+              ram_disk_memory_in_mi=test_config.ram_disk_mi,
               machine_type=test_config.machine_type,
           )
           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(cpc_conf)
@@ -141,12 +142,12 @@ class Testconfig:
             )(cpc_conf)
             apply_cpc = mtc.apply_cpc(cpc_conf)
             run_time = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")
-            run_name = f"{test_config.name_prefix}-{tc.checkpointing_name}-{slice_num}x-{test_config.accelerator}-{run_time}"
+            run_name = f"{test_config.short_id}-{tc.name}-{slice_num}x-{test_config.accelerator}-{run_time}"
             workload_command = (
                 "export TPU_PREMAPPED_BUFFER_SIZE=52428800000 && "
                 "export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES=52428800000 && "
                 "python3 -m MaxText.train MaxText/configs/base.yml remat_policy=full "
-                f"global_parameter_scale=1 base_output_directory={BASE_OUTPUT_DIRECTORY} "
+                f"global_parameter_scale=1 base_output_directory={run_dir_out} "
                 f"dataset_type=synthetic steps={test_config.step} per_device_batch_size=1 "
                 f"max_target_length=256 model_name={test_config.model_name} per_device_batch_size=2 "
                 "reuse_example_batch=1 enable_emergency_checkpoint=true "
@@ -160,7 +161,7 @@ class Testconfig:
                 num_slices=slice_num,
                 cluster=test_config.cluster,
                 time_out_in_min=60,
-                test_name=f"{test_config.name_prefix}-{tc.checkpointing_name}",
+                test_name=f"{test_config.short_id}-{tc.name}",
                 run_model_cmds=workload_command,
                 docker_image=image.value,
                 test_owner=test_owner.CAMILO_Q,
@@ -176,7 +177,7 @@ class Testconfig:
                 num_slices=slice_num,
                 cluster=test_config.cluster,
                 time_out_in_min=60,
-                test_name=f"{test_config.name_prefix}-cl",
+                test_name=f"{test_config.short_id}-cl",
                 run_model_cmds=cleanup_command,
                 docker_image=image.value,
                 test_owner=test_owner.CAMILO_Q,
@@ -190,7 +191,8 @@ class Testconfig:
             end_time = log.generate_timestamp()
             vali_step = test_config.step - 1
             vali_step_list = [
-                i for i in range(0, vali_step, test_config.local_checkpoint_step)
+                i
+                for i in range(0, vali_step, test_config.local_checkpoint_step)
             ]
             vali_step_list.append(vali_step)
 

dags/orbax/util/multi_tier_checkpoint_util.py

@@ -21,14 +21,24 @@
 class CheckpointConfiguration:
   """A dataclass to hold attributes of a Cloud Public Compute (CPC) instance."""
 
-  def __init__(self,
-    project_id: str,
-    region: str,
-    cluster_name: str,
-    gcs_bucket: str,
-    machine_type: str,
-    ramdisk_memory: str,
-    toleration_key: str = "google.com/tpu",
+  project_id: str
+  region: str
+  cluster_name: str
+  gcs_bucket: str
+  machine_type: str
+  ramdisk_memory_in_mi: str
+  toleration_key: str
+  cpc_yaml_template: str
+
+  def __init__(
+      self,
+      project_id: str,
+      region: str,
+      cluster_name: str,
+      gcs_bucket: str,
+      machine_type: str,
+      ram_disk_memory_in_mi: str,
+      toleration_key: str = "google.com/tpu",
   ):
     """
     Initializes the CheckpointConfiguration.
@@ -40,7 +50,8 @@ def __init__(self,
         gcs_bucket (str): The name of the GCS bucket for checkpoints.
         machine_type (str): The machine type for the instance.
         ramdisk_memory_in_mi (str): The size of the RAM disk in mebibytes (Mi).
-            Defaults to "2G".
+            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"".
         toleration_key (str): The toleration key for the Kubernetes pod.
             Defaults to "google.com/tpu".
     """
@@ -49,80 +60,38 @@ def __init__(self,
     self.cluster_name = cluster_name
     self.gcs_bucket = gcs_bucket
     self.machine_type = machine_type
-    self.ramdisk_memory = ramdisk_memory
+    self.ramdisk_memory = ram_disk_memory_in_mi
     self.toleration_key = toleration_key
-
-  def _get_custom_objects_api_client(
-      self, project_id: str, region: str, cluster_name: str
-  ) -> 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
-      )
-    )
-
-  # memory_size: This should be Mi
-  def _create_cpc_content(
-      self,
-      gcs_bucket: str,
-      machine_type: str,
-      toleration_key: str,
-      memory_size_in_mi: str,
-  ) -> str:
-    """Creates the CheckpointConfiguration YAML content as a string.
-
-    This method generates a templated YAML string for a GKE CheckpointConfiguration
-    resource, including parameters for the GCS bucket, machine type, toleration,
-    and RAM disk size.
-
-    Args:
-        gcs_bucket (str): The name of the GCS bucket for storing checkpoints.
-        machine_type (str): The instance type of the node.
-        toleration_key (str): The toleration key to allow scheduling on specific nodes.
-        memory_size_in_mi (str): The desired size of the in-memory volume.
-          The unit is in mebibytes (Mi) but the value should be passed as a string
-          with the unit, e.g., "2G" or "1024M".
-
-    Returns:
-        str: The templated YAML content as a string.
-    """
-    cpc_yaml_template = f"""
+    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: {gcs_bucket}
+    cloudStorageBucketName: {self.gcs_bucket}
     nodeSelector:
-      node.kubernetes.io/instance-type: {machine_type}
+      node.kubernetes.io/instance-type: {self.machine_type}
     tolerations:
-    - key: {toleration_key}
+    - key: {self.toleration_key}
       operator: Exists
       effect: NoSchedule
-    inMemoryVolumeSize: {memory_size_in_mi}
+    inMemoryVolumeSize: {self.ramdisk_memory}
   """
-    logging.info(f"Generated CPC YAML content: \n{cpc_yaml_template}")
-    return cpc_yaml_template
+
+  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._get_custom_objects_api_client(
-      cpc.project_id, cpc.region, cpc.cluster_name
-  )
-
-  cpc_yaml_string = cpc._create_cpc_content(
-      cpc.gcs_bucket,
-      cpc.machine_type,
-      cpc.toleration_key,
-      cpc.ramdisk_memory,
-  )
-  cpc_body = yaml.safe_load(cpc_yaml_string)
-
+  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")
@@ -139,7 +108,6 @@ def apply_cpc(cpc: CheckpointConfiguration) -> None:
     logging.info(f"CheckpointConfiguration '{name}' created successfully.")
 
   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:
@@ -174,17 +142,8 @@ def initiate_cpc_deletion(cpc: CheckpointConfiguration) -> None:
   """
   Sends the delete request for the CheckpointConfiguration.
   """
-  custom_api = cpc._get_custom_objects_api_client(
-      cpc.project_id, cpc.region, cpc.cluster_name
-  )
-  cpc_body = yaml.safe_load(
-      cpc._create_cpc_content(
-          cpc.gcs_bucket,
-          cpc.machine_type,
-          cpc.toleration_key,
-          cpc.ramdisk_memory,
-      )
-  )
+  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")
 
   if not name_to_delete:
@@ -234,17 +193,8 @@ def wait_for_cpc_deletion(cpc: CheckpointConfiguration) -> bool:
   """
   A sensor that waits for the CheckpointConfiguration to be completely deleted.
   """
-  custom_api = cpc._get_custom_objects_api_client(
-      cpc.project_id, cpc.region, cpc.cluster_name
-  )
-  cpc_body = yaml.safe_load(
-      cpc._create_cpc_content(
-          cpc.gcs_bucket,
-          cpc.machine_type,
-          cpc.toleration_key,
-          cpc.ramdisk_memory,
-      )
-  )
+  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")
 
   if not name_to_delete: