diff --git a/dags/common/test_owner.py b/dags/common/test_owner.py index 74796c2bf..2bf27f9b6 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_Q = "camiloCienet" # MLCompass ORTI_B = "ortibazar" diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py index 2eba611f5..bffd25e5d 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="ml-auto-solutions-orbax", + 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, @@ -321,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/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_mtc_emergency_save_local.py b/dags/orbax/maxtext_mtc_emergency_save_local.py new file mode 100644 index 000000000..27e70e7e5 --- /dev/null +++ b/dags/orbax/maxtext_mtc_emergency_save_local.py @@ -0,0 +1,299 @@ +""" +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 typing import Optional + +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 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" +DEFAULT_BUCKET = gcs_bucket.MTC_AUTOMATION_BUCKET + +# 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_ORBAX_HEAD, +)] +RAM_DISK = "/local" + + +@dataclass +class Checkpointing: + """ + Represents the information of a checkpointing mechanism. + + Attributes: + name: A unique name for the checkpointing configuration. + enable_multi_tier_checkpointing: Indicates whether multi-tier checkpointing is enabled. + """ + + name: str + enable_multi_tier_checkpointing: bool = True + + +@dataclass +class TestConfig: + """Holds the general configuration for a checkpointing test.""" + + 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: str + cpc_config: checkpoint_util.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, + ram_disk_size_in_mi: str, + checkpoint_step: Optional[int] = None, + ): + """ + 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 = checkpoint_util.CheckpointConfiguration( + project_id=self.cluster.project, + region=zone_to_region(self.cluster.zone), + cluster_name=self.cluster.name, + gcs_bucket=DEFAULT_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_folder: 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={posixpath.join(DEFAULT_BUCKET, out_folder)} " + "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"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}", + ) + + +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=""" + # 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, + machine_type="ct5p-hightpu-4t", + accelerator="v5p-128", + slices=[2], + model_name="llama2-7b", + short_id="max-sv-loc", + replicator_backup_time=30, + step=100, + local_checkpoint_step=20, + ram_disk_size_in_mi="800000Mi", + ), + ] + + task_groups = [] + + for checkpointing in [ + Checkpointing( + name="mtc", # Multi-tier Checkpointing + enable_multi_tier_checkpointing=True, + ), + Checkpointing( + name="emc", # Emergency Checkpointing + enable_multi_tier_checkpointing=False, + ), + ]: + 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: + 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. + wait_delete_cpc = checkpoint_util.wait_for_cpc_deletion.override( + trigger_rule="all_done" + )(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_folder=group.group_id, + slice_number=slice_num, + ) + + start_time = validation_util.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.short_id}-{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, + ) + + end_time = validation_util.generate_timestamp() + + 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] + + 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, + ) + ) + + ( + wait_delete_cpc + >> apply_cpc + >> start_time + >> maxtext_chkpt_run_test + >> end_time + >> validate_local_check_steps + ) + # Add to a list of test to chain them sequentially. + task_groups.append(group) + + # 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/__init__.py b/dags/orbax/util/__init__.py new file mode 100644 index 000000000..e69de29bb 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/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 5b287d676..4ee5d9e5d 100644 --- a/xlml/utils/xpk.py +++ b/xlml/utils/xpk.py @@ -28,6 +28,11 @@ # b/411426745 - Setting branch to 0.4.1 till the depdency issue is resolved. MAIN_BRANCH = "v0.4.1" + +# 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 LOGGING_URL_FORMAT = ( "https://pantheon.corp.google.com/logs/query;" @@ -116,10 +121,23 @@ def run_workload( f" --env {metric_config.SshEnvVars.GCS_OUTPUT.name}={gcs_path}" " --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: