add a new DAG 'maxtext_multi_tier_sav01_save_local'

dags/common/test_owner.py

@@ -55,6 +55,7 @@ class Team(enum.Enum):
 # Multi-tier Checkpointing
 ABHINAV_S = "ABHINAV S."
 XUEFENG_G = "XUEFENG G."
+CAMILO = "CAMILO P."
 
 # MLCompass
 ORTI_B = "Orti B."

dags/common/vm_resource.py

@@ -256,6 +256,13 @@ class XpkClusters:
       project=Project.CLOUD_TPU_MULTIPOD_DEV.value,
       zone=Zone.EUROPE_WEST4_B.value,
   )
+  TPU_V5P_128_CLUSTER_ORBAX = XpkClusterConfig(
+      name="b428061876-jf-v5p-128-2",
+      device_version=TpuVersion.V5P,
+      core_count=128,
+      project=Project.CLOUD_TPU_MULTIPOD_DEV.value,
+      zone=Zone.EUROPE_WEST4_B.value,
+  )
   TPU_V5E_256_CLUSTER = XpkClusterConfig(
       name="v5e-256-bodaborg-europe-west4",
       device_version=TpuVersion.V5E,

dags/gcs_bucket.py

@@ -27,3 +27,4 @@
 # Multi-tier checkpointing need special permission for GCS Bucket
 # For further question reach out to  Multi-tier Checkpointing Owners.
 MTC_BUCKET = "gs://mtc-bucket-us-east5/output"
+MTC_AUTOMATION_BUCKET = "gs://mtc-automation-bucket"

dags/orbax/maxtext_multi_tier_chechpoint_save_local.py

@@ -0,0 +1,136 @@
+"""Add commentMore actions
+A DAG to run MaxText multi-tier checkpointing tests (phase2: save & validate).
+"""
+
+import datetime
+from airflow import models
+from dags import composer_env, gcs_bucket
+from dags.common import test_owner
+from dags.common.vm_resource import DockerImage, XpkClusters
+from dags.multipod.configs import gke_config
+from dags.multipod.configs.common import SetupMode
+from xlml.utils import log_explorer
+from xlml.utils import orbax
+
+SCHEDULE = "0 10 * * *" if composer_env.is_prod_env() else None
+
+with models.DAG(
+    dag_id="maxtext_multi_tier_sav01_save_local",
+    schedule_interval=SCHEDULE,
+    tags=[
+        "multipod_team",
+        "maxtext",
+        "multi_tier_p2_chkpt_save_local",
+        "nightly",
+    ],
+    start_date=datetime.datetime(2025, 6, 12),
+    catchup=False,
+    concurrency=2,
+) as dag:
+  base_output_directory = (
+      f"{gcs_bucket.MTC_AUTOMATION_BUCKET}/maxtext_multi_tier_sav01_save_local"
+  )
+  docker_images = [(
+      SetupMode.NIGHTLY,
+      DockerImage.MAXTEXT_TPU_JAX_NIGHTLY,
+  )]
+  ram_disk = "/local"
+  test_configs = {"v5p-128": [2]}
+  clusters = {"v5p-128": XpkClusters.TPU_V5P_128_CLUSTER_ORBAX}
+  step = 100
+  local_checkpoint_period = 20
+  replicator_backup_interval_minutes = 1
+  use_replicator = "True"
+  name_prefix = "maxtext-p2-cpt-sv"
+  model_name = "llama2-7b"
+
+  for mode, image in docker_images:
+    for accelerator, slices in test_configs.items():
+      for slice_num in slices:
+        cpc = (
+            clusters[accelerator].project,
+            clusters[accelerator].zone[:-2],
+            clusters[accelerator].name,
+            gcs_bucket.MTC_AUTOMATION_BUCKET.split("gs://")[1],
+            "ct5p-hightpu-4t",
+            "google.com/tpu",
+            "800000Mi",
+        )
+        delete_cpc = orbax.delete_cpc(*cpc)
+        apply_cpc = orbax.apply_cpc(*cpc)
+        run_time = datetime.datetime.now().strftime("%Y-%m-%d-%H")
+        run_name = f"{name_prefix}-{slice_num}x-{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"dataset_type=synthetic steps={step} per_device_batch_size=1 "
+            f"max_target_length=256 model_name={model_name} per_device_batch_size=2 "
+            "reuse_example_batch=1 enable_emergency_checkpoint=true "
+            f"local_checkpoint_directory={ram_disk} local_checkpoint_period={local_checkpoint_period} "
+            f"use_replicator_service={use_replicator} replicator_backup_interval_minutes={replicator_backup_interval_minutes} "
+            f"run_name={run_name}",
+        )
+
+        start_time = log_explorer.generate_timestamp()
+
+        # make launch test_name unique
+        maxtext_phase2_chkpt_test = gke_config.get_gke_config(
+            num_slices=slice_num,
+            cluster=clusters[accelerator],
+            time_out_in_min=60,
+            test_name=f"{name_prefix}",
+            run_model_cmds=workload_command,
+            docker_image=image.value,
+            test_owner=test_owner.CAMILO,
+        ).run(
+            ramdisk_directory=ram_disk,
+            mtc_enabled=True,
+            xpk_branch="main",
+            skip_post_process=True,
+        )
+
+        # cleanup run: unique test_name
+        cleanup_command = (f"rm -rf {ram_disk}/*",)
+        ram_disk_cleanup = gke_config.get_gke_config(
+            num_slices=slice_num,
+            cluster=clusters[accelerator],
+            time_out_in_min=60,
+            test_name=f"{name_prefix}-cleanup",
+            run_model_cmds=cleanup_command,
+            docker_image=image.value,
+            test_owner=test_owner.CAMILO,
+        ).run(
+            ramdisk_directory=ram_disk,
+            mtc_enabled=True,
+            xpk_branch="main",
+            skip_post_process=True,
+        )
+
+        end_time = log_explorer.generate_timestamp()
+        vali_step = step - 1
+        vali_step_list = [
+            i for i in range(0, vali_step, local_checkpoint_period)
+        ]
+        vali_step_list.append(vali_step)
+
+        validate_log = log_explorer.validate_log_with_step(
+            project_id=clusters[accelerator].project,
+            location=clusters[accelerator].zone[:-2],
+            cluster_name=clusters[accelerator].name,
+            text_filter="(blocking + background).",
+            start_time=start_time,
+            end_time=end_time,
+            vali_step_list=vali_step_list,
+        )
+
+        (
+            delete_cpc
+            >> apply_cpc
+            >> start_time
+            >> maxtext_phase2_chkpt_test
+            >> ram_disk_cleanup
+            >> end_time
+            >> validate_log
+        )

xlml/utils/log_explorer.py

@@ -0,0 +1,158 @@
+"""Utilities to get workloads logs and some utils."""
+
+from airflow.decorators import task
+from airflow.exceptions import AirflowFailException
+from google.cloud import logging as log_explorer
+from datetime import datetime, timezone, timedelta
+from typing import Optional
+from absl import logging
+
+
+@task
+def generate_timestamp():
+  return datetime.now(timezone.utc)
+
+
+@task
+def validate_log_with_step(
+    project_id: str,
+    location: str,
+    cluster_name: str,
+    namespace: str = "default",
+    pod_pattern: str = "*",
+    container_name: Optional[str] = None,
+    text_filter: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+    vali_step_list: Optional[list] = None,
+) -> bool:
+  """
+  Validate the workload log is training correct
+  Args:
+      project_id: The Google Cloud project ID
+      location: GKE cluster location
+      cluster_name: GKE cluster name
+      namespace: Kubernetes namespace (defaults to "default")
+      pod_pattern: Pattern to match pod names (defaults to "*")
+      container_name: Optional container name to filter logs
+      text_filter: Optional comma-separated string to
+      filter log entries by textPayload content
+      start_time: Optional start time for log retrieval
+      (defaults to 12 hours ago)
+      end_time: Optional end time for log retrieval (defaults to now)
+      vali_step_list: optional to validate list of steps
+  Returns:
+      bool: validate success or not
+  """
+  entries = list_log_entries(
+      project_id=project_id,
+      location=location,
+      cluster_name=cluster_name,
+      namespace=namespace,
+      pod_pattern=pod_pattern,
+      container_name=container_name,
+      text_filter=text_filter,
+      start_time=start_time,
+      end_time=end_time,
+  )
+  if vali_step_list is None:
+    return False
+  new_step_list = []
+  for entry in entries:
+    if not entry.payload:
+      continue
+    payload_str = str(entry.payload)
+    for line in payload_str.split("\n"):
+      if vali_step_list is not None:
+        for step in vali_step_list:
+          vali_str = "directory=/local/" + str(step)
+          if vali_str in line and step not in new_step_list:
+            logging.info(f"├─ Timestamp: {entry.timestamp}")
+            logging.info("└─ Payload:")
+            logging.info(f"   {line}")
+            new_step_list.append(step)
+  if len(vali_step_list) == len(new_step_list):
+    logging.info("Validate success")
+    return True
+  else:
+    raise AirflowFailException(
+        f"{len(vali_step_list)} saves are expected,"
+        f"but got {len(new_step_list)}"
+    )
+
+
+def list_log_entries(
+    project_id: str,
+    location: str,
+    cluster_name: str,
+    namespace: str = "default",
+    pod_pattern: str = "*",
+    container_name: Optional[str] = None,
+    text_filter: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+) -> list:
+  """
+  List log entries for the specified Google Cloud project.
+  This function connects to Google Cloud Logging,
+  constructs a filter for Kubernetes container logs
+  within a specific project, location, cluster, namespace,
+  and pod name pattern, and retrieves log
+  entries from the specified time range.
+  It prints the timestamp, severity, resource information,
+  and payload for each log entry found.
+  Args:
+      project_id: The Google Cloud project ID
+      location: GKE cluster location
+      cluster_name: GKE cluster name
+      namespace: Kubernetes namespace (defaults to "default")
+      pod_pattern: Pattern to match pod names (defaults to "*")
+      container_name: Optional container name to filter logs
+      text_filter: Optional comma-separated string to
+      filter log entries by textPayload content
+      start_time: Optional start time for log retrieval
+      (defaults to 12 hours ago)
+      end_time: Optional end time for log retrieval (defaults to now)
+  Returns:
+      bool: Number of log entries found
+  """
+
+  # Create a Logging Client for the specified project
+  logging_client = log_explorer.Client(project=project_id)
+
+  # Set the time window for log retrieval:
+  # default to last 12 hours if not provided
+  if end_time is None:
+    end_time = datetime.now(timezone.utc)
+  if start_time is None:
+    start_time = end_time - timedelta(hours=12)
+
+  # Format times as RFC3339 UTC "Zulu" format required by the Logging API
+  start_time_str = start_time.strftime("%Y-%m-%dT%H:%M:%SZ")
+  end_time_str = end_time.strftime("%Y-%m-%dT%H:%M:%SZ")
+
+  # Construct the log filter
+  log_filter = (
+      f'resource.labels.project_id="{project_id}" '
+      f'resource.labels.location="{location}" '
+      f'resource.labels.cluster_name="{cluster_name}" '
+      f'resource.labels.namespace_name="{namespace}" '
+      f'resource.labels.pod_name:"{pod_pattern}" '
+      "severity>=DEFAULT "
+      f'timestamp>="{start_time_str}" '
+      f'timestamp<="{end_time_str}"'
+  )
+
+  # Add container name filter if provided
+  if container_name:
+    log_filter += f' resource.labels.container_name="{container_name}"'
+
+  # Add text content filter if provided
+  if text_filter:
+    log_filter += f' SEARCH("{text_filter}")'
+
+  # Retrieve log entries matching the filter
+  logging.info(f"Log filter constructed: {log_filter}")
+  entries = logging_client.list_entries(filter_=log_filter)
+
+  return entries