From 7cc92fadb83a740611ccc2eedda1d5e83c2b61d7 Mon Sep 17 00:00:00 2001 From: CamiloCienet Date: Mon, 2 Jun 2025 18:12:58 +0800 Subject: [PATCH 1/4] Add a new DAG maxtext_multi_tier_save_local. This testcase test both save locally with MTC and EMC enabled. --- dags/common/test_owner.py | 1 + dags/common/vm_resource.py | 7 + dags/gcs_bucket.py | 2 +- .../maxtext_multi_tier_p2_checkpointing.py | 2 +- ...axtext_multi_tier_chechpoint_save_local.py | 225 ++++++++++++++ dags/orbax/util/__init__.py | 0 dags/orbax/util/logging_mtc.py | 159 ++++++++++ dags/orbax/util/multi_tier_checkpoint_util.py | 285 ++++++++++++++++++ xlml/utils/xpk.py | 29 +- 9 files changed, 705 insertions(+), 5 deletions(-) create mode 100644 dags/orbax/maxtext_multi_tier_chechpoint_save_local.py create mode 100644 dags/orbax/util/__init__.py create mode 100644 dags/orbax/util/logging_mtc.py create mode 100644 dags/orbax/util/multi_tier_checkpoint_util.py diff --git a/dags/common/test_owner.py b/dags/common/test_owner.py index 74796c2bf..bc519da23 100644 --- a/dags/common/test_owner.py +++ b/dags/common/test_owner.py @@ -65,6 +65,7 @@ class Team(enum.Enum): # Multi-tier Checkpointing ABHINAV_S = "abhinavclemson" XUEFENG_G = "xuefgu" +CAMILO = "CAMILO P." # MLCompass ORTI_B = "ortibazar" diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index 2eba611f5..6950a5d2c 100644 --- a/dags/common/vm_resource.py +++ b/dags/common/vm_resource.py @@ -259,6 +259,13 @@ class XpkClusters: project=Project.CLOUD_TPU_MULTIPOD_DEV.value, zone=Zone.EUROPE_WEST4_B.value, ) + TPU_V5P_128_CLUSTER_ORBAX = XpkClusterConfig( + name="orbax-ml-auto-solutions", + 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, diff --git a/dags/gcs_bucket.py b/dags/gcs_bucket.py index d55761ee3..44f382822 100644 --- a/dags/gcs_bucket.py +++ b/dags/gcs_bucket.py @@ -26,4 +26,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" diff --git a/dags/multipod/maxtext_multi_tier_p2_checkpointing.py b/dags/multipod/maxtext_multi_tier_p2_checkpointing.py index b0cf8afff..704d98c13 100644 --- a/dags/multipod/maxtext_multi_tier_p2_checkpointing.py +++ b/dags/multipod/maxtext_multi_tier_p2_checkpointing.py @@ -39,7 +39,7 @@ concurrency=2, ) as dag: base_output_directory = ( - f"{gcs_bucket.MTC_BUCKET}/maxtext_multi_tier_p2_checkpointing" + f"{gcs_bucket.MTC_AUTOMATION_BUCKET}/maxtext_multi_tier_p2_checkpointing" ) dataset_path = gcs_bucket.MAXTEXT_DIR docker_images = [ diff --git a/dags/orbax/maxtext_multi_tier_chechpoint_save_local.py b/dags/orbax/maxtext_multi_tier_chechpoint_save_local.py new file mode 100644 index 000000000..27558268c --- /dev/null +++ b/dags/orbax/maxtext_multi_tier_chechpoint_save_local.py @@ -0,0 +1,225 @@ +""" +A DAG to run MaxText multi-tier checkpointing tests. + +This DAG performs a series of tests to save and validate checkpoints +for the MaxText model. It runs tests in two modes: one with the replicator +service enabled (Multi-tier Checkpointing). The tests are executed on a TPU multi-pod cluster. +""" + +import datetime +from dataclasses import dataclass +import posixpath + +from airflow import models +from airflow.utils.task_group import TaskGroup + +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 dags.orbax.util import logging_mtc as log +from dags.orbax.util import multi_tier_checkpoint_util as mtc +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. +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. + + 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. + """ + checkpointing_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. + """ + cluster: XpkClusters + machine_type: str + accelerator: str + slices: list[int] + model_name: str + name_prefix: str + replicator_min: int + step: int + local_checkpoint_step: int + checkpoint_step: int + ram_disk_mi: str + + +with models.DAG( + dag_id=DAG_TEST_NAME, + start_date=datetime.datetime(2025, 6, 12), + schedule_interval=SCHEDULE, + catchup=False, + tags=[ + "multipod_team", + "maxtext", + "emergency_checkpoint_manager", + "multitier_checkpointing", + "nightly", + "orbax", + ], + description="A DAG to run MaxText multi-tier checkpointing tests.", + doc_md="", + concurrency=2, + params={}, +) as 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", + name_prefix="max-sv", + replicator_min=30, + step=100, + local_checkpoint_step=20, + checkpoint_step=25, + ram_disk_mi="800000Mi", + ), + ] + 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), + ]: + 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: + 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, + 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) + wait_delete_cpc = mtc.wait_for_cpc_deletion.override( + trigger_rule="all_done" + )(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}" + 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={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 " + f"local_checkpoint_directory={RAM_DISK} local_checkpoint_period={test_config.local_checkpoint_step} " + f"use_replicator_service={tc.use_replicator} replicator_backup_interval_minutes={test_config.replicator_min} " + f"run_name={run_name}", + ) + + start_time = log.generate_timestamp() + maxtext_chkpt_run_test = gke_config.get_gke_config( + num_slices=slice_num, + cluster=test_config.cluster, + time_out_in_min=60, + test_name=f"{test_config.name_prefix}-{tc.checkpointing_name}", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.CAMILO_Q, + ).run( + ramdisk_directory=RAM_DISK, + mtc_enabled=True, + xpk_branch=BRANCH_ABHINAV_MTC, + skip_post_process=True, + ) + + cleanup_command = (f"rm -rf {RAM_DISK}/*",) + ram_disk_cleanup = gke_config.get_gke_config( + num_slices=slice_num, + cluster=test_config.cluster, + time_out_in_min=60, + test_name=f"{test_config.name_prefix}-cl", + run_model_cmds=cleanup_command, + docker_image=image.value, + test_owner=test_owner.CAMILO_Q, + ).run( + ramdisk_directory=RAM_DISK, + mtc_enabled=True, + xpk_branch=BRANCH_ABHINAV_MTC, + skip_post_process=True, + ) + + 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) + ] + 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( + project_id=test_config.cluster.project, + location=zone_to_region(test_config.cluster.zone), + cluster_name=test_config.cluster.name, + text_filter="(blocking + background).", + start_time=start_time, + end_time=end_time, + vali_step_list=vali_step_list, + ) + + ( + init_delete_cpc + >> wait_delete_cpc + >> apply_cpc + >> start_time + >> maxtext_chkpt_run_test + >> ram_disk_cleanup + >> end_time + >> validate_log + ) + # Add to a global list of test to be run in a sequential way + tests_to_run_seq.append(task) + + # Chain the task groups sequentially + for idx_test in range(len(tests_to_run_seq) - 1): + tests_to_run_seq[idx_test] >> tests_to_run_seq[idx_test + 1] diff --git a/dags/orbax/util/__init__.py b/dags/orbax/util/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/dags/orbax/util/logging_mtc.py b/dags/orbax/util/logging_mtc.py new file mode 100644 index 000000000..c482dedd6 --- /dev/null +++ b/dags/orbax/util/logging_mtc.py @@ -0,0 +1,159 @@ +"""Utilities to get workloads logs and some utils.""" + +from datetime import datetime, timezone, timedelta +from typing import Optional +from absl import logging + +from airflow.decorators import task +from airflow.exceptions import AirflowFailException +from google.cloud import logging as log_explorer + + +@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 diff --git a/dags/orbax/util/multi_tier_checkpoint_util.py b/dags/orbax/util/multi_tier_checkpoint_util.py new file mode 100644 index 000000000..b0bd981a4 --- /dev/null +++ b/dags/orbax/util/multi_tier_checkpoint_util.py @@ -0,0 +1,285 @@ +"""Utility functions for managing Multi-tier Cluster Configuration. + +This module provides tasks for creating, applying, and deleting +Multi-tier Driver cluster Configurations for enebale Multi Tier Checkpointing. +""" + +from absl import logging +import yaml +import time +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 xlml.utils import gke + + +# ramdisk_memory_in_mi: This should be Mi +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", + ): + """ + 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. + ramdisk_memory_in_mi (str): The size of the RAM disk in mebibytes (Mi). + Defaults to "2G". + 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 + 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""" + apiVersion: checkpointing.gke.io/v1 + kind: CheckpointConfiguration + metadata: + name: my-checkpointconfiguration # This name will be used for deletion + spec: + cloudStorageBucketName: {gcs_bucket} + nodeSelector: + node.kubernetes.io/instance-type: {machine_type} + tolerations: + - key: {toleration_key} + operator: Exists + effect: NoSchedule + inMemoryVolumeSize: {memory_size_in_mi} + """ + logging.info(f"Generated CPC YAML content: \n{cpc_yaml_template}") + return cpc_yaml_template + + +@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) + + 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" + + 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: + + # 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: + logging.info( + f"CheckpointConfiguration '{name}' already exists. Attempting to replace..." + ) + try: + custom_api.replace_cluster_custom_object( + group=group, + version=version, + plural=plural, + name=name, + body=cpc_body, + ) + logging.info(f"CheckpointConfiguration '{name}' replaced successfully.") + except ApiException as replace_error: + logging.error( + f"Error replacing CheckpointConfiguration:" + f" {replace_error.status} - {replace_error.reason} - {replace_error.body}" + ) + raise AirflowFailException( + f"Failed to replace CheckpointConfiguration: {replace_error.reason}" + ) + else: + raise AirflowFailException( + f"Failed to apply CheckpointConfiguration: {api_error.reason}" + ) + + +@task +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, + ) + ) + name_to_delete = cpc_body.get("metadata", {}).get("name") + + if not name_to_delete: + 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}" + ) + custom_api.delete_cluster_custom_object( + group=group, + version=version, + plural=plural, + name=name_to_delete, + body=delete_options, + ) + logging.info( + f"Delete request sent for CheckpointConfiguration '{name_to_delete}'." + ) + + 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}' not found. Already deleted or never existed." + ) + return + else: + raise AirflowFailException( + f"Unexpected error during initiate_cpc_deletion: {e}" + ) + + +@task.sensor(poke_interval=10) +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, + ) + ) + name_to_delete = cpc_body.get("metadata", {}).get("name") + + if not name_to_delete: + 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 + ) + logging.info( + f"CheckpointConfiguration '{name_to_delete}' 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" + ) + return True # Return True to tell the sensor to succeed + else: + logging.error( + f"API error while waiting for deletion: " + f"{e.status} - {e.reason} - {e.body}" + ) + raise AirflowFailException( + f"API error during CPC deletion wait: {e.reason}" + ) diff --git a/xlml/utils/xpk.py b/xlml/utils/xpk.py index 5b287d676..4dfe6c030 100644 --- a/xlml/utils/xpk.py +++ b/xlml/utils/xpk.py @@ -27,7 +27,13 @@ from dags.common.vm_resource import GpuVersion # b/411426745 - Setting branch to 0.4.1 till the depdency issue is resolved. -MAIN_BRANCH = "v0.4.1" +BRANCH_V0_4_1 = "v0.4.1" +MAIN_BRANCH = BRANCH_V0_4_1 + + +# b/437817546 - Orbax test need to use a branch to bypass the `validate_dependencies()` crash issue +BRANCH_ABHINAV_MTC = "abhinav-mtc" + # Duration = past 7 days LOGGING_URL_FORMAT = ( "https://pantheon.corp.google.com/logs/query;" @@ -114,12 +120,29 @@ def run_workload( f" --{multi_keyword}={num_slices} --docker-image={docker_image}" f" --project={cluster_project} --zone={zone}" f" --env {metric_config.SshEnvVars.GCS_OUTPUT.name}={gcs_path}" - " --restart-on-user-code-failure" ) + + # The `restart-on-user-code-failure` flag was supported in v0.4.1, + # but has been removed in later versions. + if xpk_branch == BRANCH_V0_4_1: + workload_create_cmd += " --restart-on-user-code-failure" + if ramdisk_directory: workload_create_cmd += f" --ramdisk-directory={ramdisk_directory}" + if mtc_enabled: - workload_create_cmd += " --mtc-enabled" + # b/437817546 - The flag is "mtc-enabled" (hyphen) for normal branches; + # on BRANCH_ABHINAV_MTC, it's "mtc_enabled" (underscore) instead. + flag = ( + "mtc-enabled" if xpk_branch != BRANCH_ABHINAV_MTC else "mtc_enabled" + ) + workload_create_cmd += f" --{flag}" + + # For Orbax DAG add flag '--max-restars=50' it is need it to test + # resiliency during Maxtext training with Emergency Checkpointer and + # Multi-tier Checkpointing. + if ramdisk_directory and mtc_enabled: + workload_create_cmd += " --max-restarts=50" # If using a valid GPU and the XPK branch is set to "main", then branch is switch to "v0.4.1". if is_valid_gpu_version(accelerator_type) and xpk_branch == MAIN_BRANCH: From c6050d552eadc4339a9722239025ca567b491e1a Mon Sep 17 00:00:00 2001 From: Camilo Quinones Date: Sun, 10 Aug 2025 08:10:22 +0000 Subject: [PATCH 2/4] 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 --- dags/common/test_owner.py | 2 +- dags/common/vm_resource.py | 2 +- ...py => maxtext_mtc_emergency_save_local.py} | 78 ++++++----- dags/orbax/util/multi_tier_checkpoint_util.py | 130 ++++++------------ 4 files changed, 82 insertions(+), 130 deletions(-) rename dags/orbax/{maxtext_multi_tier_chechpoint_save_local.py => maxtext_mtc_emergency_save_local.py} (76%) diff --git a/dags/common/test_owner.py b/dags/common/test_owner.py index bc519da23..2bf27f9b6 100644 --- a/dags/common/test_owner.py +++ b/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" diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index 6950a5d2c..e402068b2 100644 --- a/dags/common/vm_resource.py +++ b/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, diff --git a/dags/orbax/maxtext_multi_tier_chechpoint_save_local.py b/dags/orbax/maxtext_mtc_emergency_save_local.py similarity index 76% rename from dags/orbax/maxtext_multi_tier_chechpoint_save_local.py rename to dags/orbax/maxtext_mtc_emergency_save_local.py index 27558268c..42f4c229f 100644 --- a/dags/orbax/maxtext_multi_tier_chechpoint_save_local.py +++ b/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,14 +117,12 @@ 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( @@ -128,11 +130,10 @@ class Testconfig: 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) diff --git a/dags/orbax/util/multi_tier_checkpoint_util.py b/dags/orbax/util/multi_tier_checkpoint_util.py index b0bd981a4..93415b4cb 100644 --- a/dags/orbax/util/multi_tier_checkpoint_util.py +++ b/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: From 807d6fdf3382d9432e0e8726a9109b42fdd82746 Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Thu, 21 Aug 2025 17:39:24 +0800 Subject: [PATCH 3/4] Fixes for code review --- .../orbax/maxtext_mtc_emergency_save_local.py | 199 ++++++++++++------ dags/orbax/util/multi_tier_checkpoint_util.py | 31 ++- 2 files changed, 149 insertions(+), 81 deletions(-) diff --git a/dags/orbax/maxtext_mtc_emergency_save_local.py b/dags/orbax/maxtext_mtc_emergency_save_local.py index 42f4c229f..6e4ee07e3 100644 --- a/dags/orbax/maxtext_mtc_emergency_save_local.py +++ b/dags/orbax/maxtext_mtc_emergency_save_local.py @@ -3,7 +3,8 @@ This DAG performs a series of tests to save and validate checkpoints for the MaxText model. It runs tests in two modes: one with the replicator -service enabled (Multi-tier Checkpointing). The tests are executed on a TPU multi-pod cluster. +service enabled (Multi-tier Checkpointing). The tests are executed on a TPU +multi-pod cluster. """ import datetime @@ -37,13 +38,13 @@ @dataclass -class Testcase: +class Checkpointing: """ Represents the information of a checkpointing mechanism. Attributes: name: A unique name for the checkpointing configuration. - use_replicator: Indicates whether a replicator is enabled for this test case. + use_replicator: Indicates whether a replicator is enabled. """ name: str @@ -51,22 +52,8 @@ class Testcase: @dataclass -class Testconfig: - """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. - """ +class TestConfig: + """Holds the general configuration for a checkpointing test.""" cluster: XpkClusters machine_type: str @@ -74,11 +61,98 @@ class Testconfig: slices: list[int] model_name: str short_id: str - replicator_min: int + replicator_backup_time: int step: int local_checkpoint_step: int checkpoint_step: int - ram_disk_mi: str + ram_disk_size: str + cpc_config: mtc.CheckpointConfiguration + + def __init__( + self, + cluster: XpkClusters, + machine_type: str, + accelerator: str, + slices: list[int], + model_name: str, + short_id: str, + replicator_backup_time: int, + step: int, + local_checkpoint_step: int, + checkpoint_step: int, + ram_disk_size_in_mi: str, + ): + """ + Initializes the test configurations. + + Args: + 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: A short identifier for the test run. + replicator_backup_time: The allowed time for replicator takes to backup + and store checkpoint 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_size_in_mi: The size in mebibytes (Mi) about the RAM disk in the + CSI driver. 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"". + """ + + self.cluster = cluster + self.machine_type = machine_type + self.accelerator = accelerator + self.slices = slices + self.model_name = model_name + self.short_id = short_id + self.replicator_backup_time = replicator_backup_time + self.step = step + 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( + project_id=self.cluster.project, + region=zone_to_region(self.cluster.zone), + cluster_name=self.cluster.name, + gcs_bucket=gcs_bucket.MTC_AUTOMATION_BUCKET.removeprefix("gs://"), + ramdisk_memory_in_mi=self.ram_disk_size, + machine_type=self.machine_type, + ) + + def generate_workload_command( + self, + cp: Checkpointing, + checkpoint_dir: str, + out_dir: str, + slice_number: int, + ) -> str: + run_time = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M") + run_name = f"{self.short_id}-{cp.name}-{slice_number}x-{self.accelerator}-{run_time}" + return ( + "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 " + "global_parameter_scale=1 " + f"base_output_directory={out_dir} " + "dataset_type=synthetic " + f"steps={self.step} " + "per_device_batch_size=1 " + "max_target_length=256 " + f"model_name={self.model_name} " + "per_device_batch_size=2 " + "reuse_example_batch=1 " + "enable_emergency_checkpoint=true " + f"local_checkpoint_directory={checkpoint_dir} " + f"local_checkpoint_period={self.local_checkpoint_step} " + f"use_replicator_service={cp.use_replicator} " + f"replicator_backup_interval_minutes={self.replicator_backup_time} " + f"run_name={run_name}", + ) with models.DAG( @@ -97,63 +171,56 @@ class Testconfig: description="A DAG to run MaxText multi-tier checkpointing tests.", doc_md="", concurrency=2, - params={}, ) as dag: test_configs = [ - Testconfig( + 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", - replicator_min=30, + replicator_backup_time=30, step=100, local_checkpoint_step=20, checkpoint_step=25, - ram_disk_mi="800000Mi", + ram_disk_size_in_mi="800000Mi", ), ] - tests_to_run_seq = [] - # Individual test cases for Multi-tier Checkpointing and Emergency Checkpointing - for tc in [ - Testcase(name="mtc", use_replicator=True), - Testcase(name="emc", use_replicator=False), + + task_groups = [] + + for checkpointing in [ + Checkpointing( + name="mtc", # Multi-tier Checkpointing + use_replicator=True, + ), + Checkpointing( + name="emc", # Emergency Checkpointing + use_replicator=False, + ), ]: - 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: + 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: 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://"), - 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) + init_delete_cpc = mtc.initiate_cpc_deletion(test_config.cpc_config) wait_delete_cpc = mtc.wait_for_cpc_deletion.override( trigger_rule="all_done" - )(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.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={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 " - f"local_checkpoint_directory={RAM_DISK} local_checkpoint_period={test_config.local_checkpoint_step} " - f"use_replicator_service={tc.use_replicator} replicator_backup_interval_minutes={test_config.replicator_min} " - f"run_name={run_name}", + )(test_config.cpc_config) + apply_cpc = mtc.apply_cpc(test_config.cpc_config) + workload_command = test_config.generate_workload_command( + cp=checkpointing, + checkpoint_dir=RAM_DISK, + out_dir=output_dir, + slice_number=slice_num, ) start_time = log.generate_timestamp() @@ -161,7 +228,7 @@ class Testconfig: num_slices=slice_num, cluster=test_config.cluster, time_out_in_min=60, - test_name=f"{test_config.short_id}-{tc.name}", + test_name=f"{test_config.short_id}-{checkpointing.name}", run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.CAMILO_Q, @@ -197,8 +264,10 @@ class Testconfig: 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) + # 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( project_id=test_config.cluster.project, location=zone_to_region(test_config.cluster.zone), @@ -219,9 +288,9 @@ class Testconfig: >> end_time >> validate_log ) - # Add to a global list of test to be run in a sequential way - tests_to_run_seq.append(task) + # Add to a list of test to chain them sequentially. + task_groups.append(group) - # Chain the task groups sequentially - for idx_test in range(len(tests_to_run_seq) - 1): - tests_to_run_seq[idx_test] >> tests_to_run_seq[idx_test + 1] + # Chain all task groups sequentially. + for idx in range(len(task_groups) - 1): + task_groups[idx] >> task_groups[idx + 1] diff --git a/dags/orbax/util/multi_tier_checkpoint_util.py b/dags/orbax/util/multi_tier_checkpoint_util.py index 93415b4cb..10ee89e1a 100644 --- a/dags/orbax/util/multi_tier_checkpoint_util.py +++ b/dags/orbax/util/multi_tier_checkpoint_util.py @@ -1,12 +1,11 @@ """Utility functions for managing Multi-tier Cluster Configuration. This module provides tasks for creating, applying, and deleting -Multi-tier Driver cluster Configurations for enebale Multi Tier Checkpointing. +Multi-tier Driver cluster Configurations for enable Multi Tier Checkpointing. """ from absl import logging import yaml -import time from dataclasses import dataclass from kubernetes import client as k8s_client @@ -17,7 +16,7 @@ from xlml.utils import gke -# ramdisk_memory_in_mi: This should be Mi +@dataclass class CheckpointConfiguration: """A dataclass to hold attributes of a Cloud Public Compute (CPC) instance.""" @@ -26,7 +25,7 @@ class CheckpointConfiguration: cluster_name: str gcs_bucket: str machine_type: str - ramdisk_memory_in_mi: str + ramdisk_memory: str toleration_key: str cpc_yaml_template: str @@ -37,30 +36,30 @@ def __init__( cluster_name: str, gcs_bucket: str, machine_type: str, - ram_disk_memory_in_mi: 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. - ramdisk_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"". - toleration_key (str): The toleration key for the Kubernetes pod. - Defaults to "google.com/tpu". + 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., "2G" or "2048M". Defaults to "100G"". + 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 = ram_disk_memory_in_mi + self.ramdisk_memory = ramdisk_memory_in_mi self.toleration_key = toleration_key self.cpc_yaml_template = f""" apiVersion: checkpointing.gke.io/v1 From a40c57d6e4c4cfd9d68e9a3f31e44a280c68bbf6 Mon Sep 17 00:00:00 2001 From: Camilo Quinones Date: Thu, 21 Aug 2025 17:39:24 +0800 Subject: [PATCH 4/4] Fix: Adding the TODO comment on b/437817546 for Abhinov branch. 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 --- dags/common/vm_resource.py | 3 + .../orbax/maxtext_mtc_emergency_save_local.py | 137 +++++----- dags/orbax/util/checkpoint_util.py | 177 +++++++++++++ dags/orbax/util/logging_mtc.py | 159 ------------ dags/orbax/util/multi_tier_checkpoint_util.py | 234 ------------------ dags/orbax/util/validation_util.py | 157 ++++++++++++ xlml/utils/xpk.py | 13 +- 7 files changed, 411 insertions(+), 469 deletions(-) create mode 100644 dags/orbax/util/checkpoint_util.py delete mode 100644 dags/orbax/util/logging_mtc.py delete mode 100644 dags/orbax/util/multi_tier_checkpoint_util.py create mode 100644 dags/orbax/util/validation_util.py diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index e402068b2..bffd25e5d 100644 --- a/dags/common/vm_resource.py +++ b/dags/common/vm_resource.py @@ -328,6 +328,9 @@ class DockerImage(enum.Enum): """Common docker images.""" XPK_JAX_TEST = "gcr.io/cloud-ml-auto-solutions/xpk_jax_test:latest" + MAXTEXT_TPU_JAX_ORBAX_HEAD = ( + "gcr.io/cloud-tpu-multipod-dev/maxtext-orbax-custom:latest" + ) PYTORCH_NIGHTLY = ( "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/" f"xla:nightly_3.10_tpuvm_{datetime.datetime.today().strftime('%Y%m%d')}" diff --git a/dags/orbax/maxtext_mtc_emergency_save_local.py b/dags/orbax/maxtext_mtc_emergency_save_local.py index 6e4ee07e3..27e70e7e5 100644 --- 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,20 +20,20 @@ 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 +DEFAULT_BUCKET = 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, - DockerImage.MAXTEXT_TPU_JAX_NIGHTLY, + DockerImage.MAXTEXT_TPU_JAX_ORBAX_HEAD, )] RAM_DISK = "/local" @@ -44,11 +45,11 @@ class Checkpointing: Attributes: name: A unique name for the checkpointing configuration. - use_replicator: Indicates whether a replicator is enabled. + enable_multi_tier_checkpointing: Indicates whether multi-tier checkpointing is enabled. """ name: str - use_replicator: bool + enable_multi_tier_checkpointing: bool = True @dataclass @@ -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,11 +115,11 @@ 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, - gcs_bucket=gcs_bucket.MTC_AUTOMATION_BUCKET.removeprefix("gs://"), + gcs_bucket=DEFAULT_BUCKET.removeprefix("gs://"), ramdisk_memory_in_mi=self.ram_disk_size, machine_type=self.machine_type, ) @@ -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(DEFAULT_BUCKET, out_folder)} " "dataset_type=synthetic " f"steps={self.step} " "per_device_batch_size=1 " @@ -149,8 +150,8 @@ def generate_workload_command( "enable_emergency_checkpoint=true " f"local_checkpoint_directory={checkpoint_dir} " f"local_checkpoint_period={self.local_checkpoint_step} " - f"use_replicator_service={cp.use_replicator} " - f"replicator_backup_interval_minutes={self.replicator_backup_time} " + f"enable_multi_tier_checkpointing={cp.enable_multi_tier_checkpointing} " + f"multi_tier_checkpointing_backup_interval_minutes={self.replicator_backup_time} " f"run_name={run_name}", ) @@ -169,9 +170,38 @@ 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, @@ -179,11 +209,10 @@ def generate_workload_command( 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", ), ] @@ -193,16 +222,13 @@ def generate_workload_command( for checkpointing in [ Checkpointing( name="mtc", # Multi-tier Checkpointing - use_replicator=True, + enable_multi_tier_checkpointing=True, ), Checkpointing( name="emc", # Emergency Checkpointing - use_replicator=False, + enable_multi_tier_checkpointing=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, @@ -239,54 +264,32 @@ def generate_workload_command( skip_post_process=True, ) - cleanup_command = (f"rm -rf {RAM_DISK}/*",) - ram_disk_cleanup = gke_config.get_gke_config( - num_slices=slice_num, - cluster=test_config.cluster, - time_out_in_min=60, - test_name=f"{test_config.short_id}-cl", - run_model_cmds=cleanup_command, - docker_image=image.value, - test_owner=test_owner.CAMILO_Q, - ).run( - ramdisk_directory=RAM_DISK, - mtc_enabled=True, - xpk_branch=BRANCH_ABHINAV_MTC, - skip_post_process=True, - ) + end_time = validation_util.generate_timestamp() - 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) - ] - vali_step_list.append(vali_step) + total_steps = test_config.step + k = test_config.local_checkpoint_step + last_step = test_config.step - 1 + steps_to_validate = [*range(0, total_steps, k), last_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( - project_id=test_config.cluster.project, - location=zone_to_region(test_config.cluster.zone), - cluster_name=test_config.cluster.name, - text_filter="(blocking + background).", - start_time=start_time, - end_time=end_time, - vali_step_list=vali_step_list, + validate_local_check_steps = ( + validation_util.validate_checkpoint_at_steps_are_saved( + project_id=test_config.cluster.project, + location=zone_to_region(test_config.cluster.zone), + cluster_name=test_config.cluster.name, + ram_disk=RAM_DISK, + start_time=start_time, + end_time=end_time, + steps_to_validate=steps_to_validate, + ) ) ( - init_delete_cpc - >> wait_delete_cpc + wait_delete_cpc >> apply_cpc >> start_time >> maxtext_chkpt_run_test - >> ram_disk_cleanup >> end_time - >> validate_log + >> validate_local_check_steps ) # Add to a list of test to chain them sequentially. task_groups.append(group) diff --git a/dags/orbax/util/checkpoint_util.py b/dags/orbax/util/checkpoint_util.py new file mode 100644 index 000000000..16835e7fe --- /dev/null +++ b/dags/orbax/util/checkpoint_util.py @@ -0,0 +1,177 @@ +"""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() + + logging.info(f"Attempting to create CheckpointConfiguration '{name}'...") + custom_api.create_cluster_custom_object( + group=group, version=version, plural=plural, body=cpc_body + ) + + +def _delete_cpc(cpc: CheckpointConfiguration) -> bool: + """ + Sends the delete request for the CheckpointConfiguration. + Returns if the deletion was success in boolean. + """ + + 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: + # A `CheckpointConfiguration not found` error indicates that the deletion was successful. + 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) diff --git a/dags/orbax/util/logging_mtc.py b/dags/orbax/util/logging_mtc.py deleted file mode 100644 index c482dedd6..000000000 --- a/dags/orbax/util/logging_mtc.py +++ /dev/null @@ -1,159 +0,0 @@ -"""Utilities to get workloads logs and some utils.""" - -from datetime import datetime, timezone, timedelta -from typing import Optional -from absl import logging - -from airflow.decorators import task -from airflow.exceptions import AirflowFailException -from google.cloud import logging as log_explorer - - -@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 diff --git a/dags/orbax/util/multi_tier_checkpoint_util.py b/dags/orbax/util/multi_tier_checkpoint_util.py deleted file mode 100644 index 10ee89e1a..000000000 --- a/dags/orbax/util/multi_tier_checkpoint_util.py +++ /dev/null @@ -1,234 +0,0 @@ -"""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 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 - cpc_yaml_template: 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., "2G" or "2048M". Defaults to "100G"". - 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 - 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 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() - 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" - - 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: - # 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: - logging.info( - f"CheckpointConfiguration '{name}' already exists. Attempting to replace..." - ) - try: - custom_api.replace_cluster_custom_object( - group=group, - version=version, - plural=plural, - name=name, - body=cpc_body, - ) - logging.info(f"CheckpointConfiguration '{name}' replaced successfully.") - except ApiException as replace_error: - logging.error( - f"Error replacing CheckpointConfiguration:" - f" {replace_error.status} - {replace_error.reason} - {replace_error.body}" - ) - raise AirflowFailException( - f"Failed to replace CheckpointConfiguration: {replace_error.reason}" - ) - else: - raise AirflowFailException( - f"Failed to apply CheckpointConfiguration: {api_error.reason}" - ) - - -@task -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") - - if not name_to_delete: - 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}" - ) - custom_api.delete_cluster_custom_object( - group=group, - version=version, - plural=plural, - name=name_to_delete, - body=delete_options, - ) - logging.info( - f"Delete request sent for CheckpointConfiguration '{name_to_delete}'." - ) - - 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}' not found. Already deleted or never existed." - ) - return - else: - raise AirflowFailException( - f"Unexpected error during initiate_cpc_deletion: {e}" - ) - - -@task.sensor(poke_interval=10) -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") - - if not name_to_delete: - 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 - ) - logging.info( - f"CheckpointConfiguration '{name_to_delete}' 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" - ) - return True # Return True to tell the sensor to succeed - else: - logging.error( - f"API error while waiting for deletion: " - f"{e.status} - {e.reason} - {e.body}" - ) - raise AirflowFailException( - f"API error during CPC deletion wait: {e.reason}" - ) diff --git a/dags/orbax/util/validation_util.py b/dags/orbax/util/validation_util.py new file mode 100644 index 000000000..da75c932b --- /dev/null +++ b/dags/orbax/util/validation_util.py @@ -0,0 +1,157 @@ +"""Utilities to get workloads logs and some utils.""" + +from datetime import datetime, timezone, timedelta +from typing import Optional +from absl import logging +import re + +from airflow.decorators import task +from airflow.exceptions import AirflowFailException +from google.cloud import logging as logging_api + + +@task +def generate_timestamp(): + return datetime.now(timezone.utc) + + +@task +def validate_checkpoint_at_steps_are_saved( + project_id: str, + location: str, + cluster_name: str, + steps_to_validate: list, + ram_disk: str = "/local", + start_time: Optional[datetime] = None, + end_time: Optional[datetime] = None, +) -> None: + """ + Validates that a workload is training correctly by checking for specific log steps. + + This function queries logs from a specified GKE cluster and namespace. + It searches for a log entry containing the string '(blocking + background)' + and then compares the number of steps found against an expected list of steps. + + A mismatch in the number of steps will cause the validation to fail. This can + happen if, for example, a restore operation causes the step count to restart + from zero, leading to `len(vali_step_list) != len(found_steps)`. + + Args: + project_id: The Google Cloud project ID + location: GKE cluster location + cluster_name: GKE cluster name + start_time: Optional start time for log retrieval + (defaults to 12 hours ago) + end_time: Optional end time for log retrieval (defaults to now) + steps_to_validate: Optional to validate list of steps + Returns: + None: This function does not return a value. + """ + + log_pattern = ( + r"Finished async_save \(blocking \+ background\)\. " + rf"Time taken: \d+\.\d+s\. directory={ram_disk}/(\d+)" + ) + complied_pattern = re.compile(log_pattern) + entries = list_log_entries( + project_id=project_id, + location=location, + cluster_name=cluster_name, + text_filter=f'jsonPayload.message=~"{log_pattern}"', + start_time=start_time, + end_time=end_time, + ) + + steps_are_saved: set[int] = set() # Use a set for faster lookup. + for entry in entries: + if not isinstance(entry, logging_api.StructEntry): + raise AirflowFailException( + "Log entry must be contain a jsonPayload attribute." + ) + message = entry.payload.get("message") + if not message: + raise AirflowFailException(f"Failed to parse entry {entry}") + + m = complied_pattern.search(message) + if m: + steps_are_saved.add(int(m.group(1))) + + for step in steps_to_validate: + if step not in steps_are_saved: + logging.info(f"Found entries: {entries}") + raise AirflowFailException( + f"Failed to validate. Expect steps are saved: {steps_to_validate}; " + f"got: {steps_are_saved}" + ) + + +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[logging_api.LogEntry]: + """ + 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 + """ + + logging_client = logging_api.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") + + conditions = [ + 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}"', + ] + + if container_name: + conditions.append(f'resource.labels.container_name="{container_name}"') + if text_filter: + conditions.append(f"{text_filter}") + + log_filter = " AND ".join(conditions) + + logging.info(f"Log filter constructed: {log_filter}") + return list(logging_client.list_entries(filter_=log_filter)) diff --git a/xlml/utils/xpk.py b/xlml/utils/xpk.py index 4dfe6c030..4ee5d9e5d 100644 --- a/xlml/utils/xpk.py +++ b/xlml/utils/xpk.py @@ -27,11 +27,10 @@ from dags.common.vm_resource import GpuVersion # b/411426745 - Setting branch to 0.4.1 till the depdency issue is resolved. -BRANCH_V0_4_1 = "v0.4.1" -MAIN_BRANCH = BRANCH_V0_4_1 +MAIN_BRANCH = "v0.4.1" - -# b/437817546 - Orbax test need to use a branch to bypass the `validate_dependencies()` crash issue +# TODO(b/437817546): Switch back to the main branch after the issue is resolved. +# This branch includes changes fixing the `validate_dependencies()` crash issue. BRANCH_ABHINAV_MTC = "abhinav-mtc" # Duration = past 7 days @@ -120,13 +119,9 @@ def run_workload( f" --{multi_keyword}={num_slices} --docker-image={docker_image}" f" --project={cluster_project} --zone={zone}" f" --env {metric_config.SshEnvVars.GCS_OUTPUT.name}={gcs_path}" + " --restart-on-user-code-failure" ) - # The `restart-on-user-code-failure` flag was supported in v0.4.1, - # but has been removed in later versions. - if xpk_branch == BRANCH_V0_4_1: - workload_create_cmd += " --restart-on-user-code-failure" - if ramdisk_directory: workload_create_cmd += f" --ramdisk-directory={ramdisk_directory}"