diff --git a/dags/common/test_owner.py b/dags/common/test_owner.py index 4789e4a4c..24bc9e0aa 100644 --- a/dags/common/test_owner.py +++ b/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." diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index 885fdbfa8..5e0dead83 100644 --- a/dags/common/vm_resource.py +++ 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, diff --git a/dags/gcs_bucket.py b/dags/gcs_bucket.py index d55761ee3..ccc6795ae 100644 --- a/dags/gcs_bucket.py +++ b/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" 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..d0f6a07e8 --- /dev/null +++ b/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 + ) diff --git a/xlml/utils/log_explorer.py b/xlml/utils/log_explorer.py new file mode 100644 index 000000000..c4cdde096 --- /dev/null +++ b/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 diff --git a/xlml/utils/orbax.py b/xlml/utils/orbax.py new file mode 100644 index 000000000..67cf571f9 --- /dev/null +++ b/xlml/utils/orbax.py @@ -0,0 +1,227 @@ +from absl import logging +from airflow.decorators import task +from airflow.exceptions import AirflowFailException +from kubernetes import client as k8s_client +from kubernetes.client.rest import ApiException +from xlml.utils import gke +import yaml +import time + +# --- Utility Functions --- + + +def _get_custom_objects_api_client( + project_id: str, region: str, cluster_name: str +) -> k8s_client.CustomObjectsApi: + """Create a CustomObjectsApi client for the given cluster.""" + client = gke.get_authenticated_client(project_id, region, cluster_name) + custom_api = k8s_client.CustomObjectsApi(client) + logging.info( + "Successful initialize k8s CustomObjectsApi client from cluster response." + ) + return custom_api + + +def create_cpc_content( + gcs_bucket: str, + machine_type: str, + toleration_key: str, + memory_size: str, +) -> str: + """ + Creates the CheckpointConfiguration YAML content string with placeholders. + Returns 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} +""" + logging.info(f"Generated CPC YAML content: \n{cpc_yaml_template}") + return cpc_yaml_template + + +# --- Airflow Tasks --- + + +@task +def apply_cpc( + project_id: str, + region: str, + cluster_name: str, + gcs_bucket: str, + machine_type: str, + toleration_key: str, + memory_size: str, +) -> None: + """Applies the CheckpointConfiguration to the cluster (create or replace).""" + custom_api = _get_custom_objects_api_client(project_id, region, cluster_name) + + cpc_yaml_string = create_cpc_content( + gcs_bucket, + machine_type, + toleration_key, + memory_size, + ) + 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: + logging.info(f"Checking if CheckpointConfiguration '{name}' exists...") + custom_api.get_cluster_custom_object( + group=group, version=version, plural=plural, name=name + ) + logging.info( + f"CheckpointConfiguration '{name}' found. Attempting to replace." + ) + 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 api_error: + if api_error.status == 404: + logging.info( + f"CheckpointConfiguration '{name}' not found. Attempting to create." + ) + custom_api.create_cluster_custom_object( + group=group, version=version, plural=plural, body=cpc_body + ) + logging.info(f"CheckpointConfiguration '{name}' created successfully.") + else: + logging.error( + f"Error applying CheckpointConfiguration:" + f"{api_error.status} - {api_error.reason} - {api_error.body}" + ) + raise AirflowFailException( + f"Failed to apply CheckpointConfiguration: {api_error.reason}" + ) + except Exception as e: + logging.error(f"An unexpected error occurred during apply_cpc: {e}") + raise AirflowFailException(f"Unexpected error during apply_cpc: {e}") + + +@task +def delete_cpc( + project_id: str, + region: str, + cluster_name: str, + gcs_bucket: str, + machine_type: str, + toleration_key: str, + memory_size: str, + poll_interval_seconds: int = 10, +) -> None: + """ + Deletes the CheckpointConfiguration from the cluster and waits indefinitely + for its complete deletion. + """ + custom_api = _get_custom_objects_api_client(project_id, region, cluster_name) + + cpc_yaml_string = create_cpc_content( + gcs_bucket, + machine_type, + toleration_key, + memory_size, + ) + cpc_body = yaml.safe_load(cpc_yaml_string) + 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"Sending delete request for 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}'." + ) + + # --- Wait indefinitely for deletion to complete --- + logging.info( + f"Waiting indefinitely for CheckpointConfiguration" + f" '{name_to_delete}' to be deleted." + ) + while True: + try: + custom_api.get_cluster_custom_object( + group=group, version=version, plural=plural, name=name_to_delete + ) + logging.info( + f"CheckpointConfiguration '{name_to_delete}' " + f"still exists. Polling again in {poll_interval_seconds}s..." + ) + time.sleep(poll_interval_seconds) + except ApiException as e: + if e.status == 404: + logging.info( + f"CheckpointConfiguration: {name_to_delete} successfully deleted" + ) + return + 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}" + ) + except Exception as e: + logging.error( + f"An unexpected error occurred while waiting for CPC deletion: {e}" + ) + raise AirflowFailException( + f"Unexpected error during CPC deletion wait: {e}" + ) + + except ApiException as e: + if e.status == 404: + logging.info( + f"CheckpointConfiguration '{name_to_delete}' not found." + f"Already deleted or never existed. Skipping deletion." + ) + return + else: + logging.error( + f"Error during initial delete request for CheckpointConfiguration: " + f"{e.status} - {e.reason} - {e.body}" + ) + raise AirflowFailException( + f"Failed to delete CheckpointConfiguration: {e.reason}" + ) + except Exception as e: + logging.error(f"An unexpected error occurred during delete_cpc: {e}") + raise AirflowFailException(f"Unexpected error during delete_cpc: {e}") diff --git a/xlml/utils/xpk.py b/xlml/utils/xpk.py index 5b287d676..12f4b2a79 100644 --- a/xlml/utils/xpk.py +++ b/xlml/utils/xpk.py @@ -51,8 +51,12 @@ def get_xpk_setup_cmd(tmpdir, branch: str = MAIN_BRANCH): cmds = [ "set -xue", clone_branch, + f"cd {tmpdir}/xpk/src/xpk", + "sed -i '/validate_dependencies()/s/^/#/' main.py || true", + "cat main.py", "pip install ruamel.yaml docker", ] + return cmds @@ -121,6 +125,13 @@ def run_workload( if mtc_enabled: workload_create_cmd += " --mtc-enabled" + # For Orbax DAG + if ramdisk_directory and mtc_enabled: + workload_create_cmd = workload_create_cmd.replace( + " --restart-on-user-code-failure", "" + ) + 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: xpk_branch = "v0.4.1"