diff --git a/dags/common/test_owner.py b/dags/common/test_owner.py index fe2890d9..979f8559 100644 --- a/dags/common/test_owner.py +++ b/dags/common/test_owner.py @@ -54,6 +54,7 @@ class Team(enum.Enum): # Multi-tier Checkpointing ABHINAV_S = "ABHINAV S." XUEFENG_G = "XUEFENG G." +JACKY_F = "JACKY F." # MLCompass ORTI_B = "Orti B." diff --git a/dags/map_reproducibility/internal_runs/a3ultra_maxtext_benchmarking_dags.py b/dags/map_reproducibility/internal_runs/a3ultra_maxtext_benchmarking_dags.py index 8339548d..0d2e03d5 100644 --- a/dags/map_reproducibility/internal_runs/a3ultra_maxtext_benchmarking_dags.py +++ b/dags/map_reproducibility/internal_runs/a3ultra_maxtext_benchmarking_dags.py @@ -19,9 +19,9 @@ from airflow import models from dags import composer_env -from dags.map_reproducibility.utils.constants import Image +from dags.map_reproducibility.utils.constants import Image, WorkloadLauncher from dags.map_reproducibility.internal_runs.dag_configs import DAG_CONFIGS_ULTRA -from dags.map_reproducibility.utils.internal_aotc_workload import run_internal_aotc_workload +from dags.map_reproducibility.utils.internal_aotc_workload import run_internal_dag_united_workload # Configuration parameters @@ -38,7 +38,7 @@ "reproducibility", "experimental", "xlml", - "v1.17", + "v2.0", "internal", "regressiontests", "a3ultra", @@ -69,12 +69,13 @@ start_date=datetime.datetime(2025, 4, 3), catchup=False, ) as dag: - run_internal_aotc_workload( + run_internal_dag_united_workload( relative_config_yaml_path=config_path, test_run=TEST_RUN, backfill=BACKFILL, timeout=timeout, image_version=NIGHTLY_IMAGE, + workload_launcher=WorkloadLauncher.MAXTEXT_LAUNCHER, ) # Create DAG for stable release @@ -86,10 +87,11 @@ start_date=datetime.datetime(2025, 4, 3), catchup=False, ) as dag: - run_internal_aotc_workload( + run_internal_dag_united_workload( relative_config_yaml_path=config_path, test_run=TEST_RUN, backfill=BACKFILL, timeout=timeout, image_version=RELEASE_IMAGE, + workload_launcher=WorkloadLauncher.MAXTEXT_LAUNCHER, ) diff --git a/dags/map_reproducibility/internal_runs/a3ultra_nemo_benchmarking_dags.py b/dags/map_reproducibility/internal_runs/a3ultra_nemo_benchmarking_dags.py new file mode 100644 index 00000000..94f29b7d --- /dev/null +++ b/dags/map_reproducibility/internal_runs/a3ultra_nemo_benchmarking_dags.py @@ -0,0 +1,77 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DAGs to run Aotc reproducibility benchmarks.""" + +import datetime +import os + +from airflow import models +from dags import composer_env +from dags.map_reproducibility.utils.constants import Image, WorkloadLauncher +from dags.map_reproducibility.internal_runs.dag_configs import DAG_CONFIGS_ULTRA_NEMO +from dags.map_reproducibility.utils.internal_aotc_workload import run_internal_dag_united_workload + + +# Configuration parameters +TEST_RUN = False if composer_env.is_prod_env() else True +BACKFILL = False + +# Get current date for image tags +RELEASE_IMAGE = Image.NEMO_STABLE_RELEASE_A3U + +# Common DAG tags +DAG_TAGS = [ + "reproducibility", + "experimental", + "xlml", + "v2.0", + "internal", + "regressiontests", + "a3ultra", +] + + +# Create DAGs for each configuration +for config_path, config_info in DAG_CONFIGS_ULTRA_NEMO.items(): + # Extract config name for the DAG ID + config_name = os.path.basename(config_path).replace(".yaml", "") + schedule = config_info["schedule"] if not TEST_RUN else None + timeout = config_info["timeout_minutes"] + # Set retry parameter based on timeout + retries = 1 if timeout <= 15 else 0 + retry_delay = datetime.timedelta(minutes=1) + + dag_default_args = { + "retries": retries, + "retry_delay": retry_delay, + } + + # Create DAG for stable release + with models.DAG( + dag_id=f"new_internal_stable_release_{config_name}", + default_args=dag_default_args, + schedule=schedule, + tags=DAG_TAGS, + start_date=datetime.datetime(2025, 4, 3), + catchup=False, + ) as dag: + run_internal_dag_united_workload( + relative_config_yaml_path=config_path, + test_run=TEST_RUN, + backfill=BACKFILL, + timeout=timeout, + image_version=RELEASE_IMAGE, + workload_launcher=WorkloadLauncher.NEMO_TEN_LAUNCHER, + ) diff --git a/dags/map_reproducibility/internal_runs/dag_configs.py b/dags/map_reproducibility/internal_runs/dag_configs.py index 7ccd7865..a423c4dc 100644 --- a/dags/map_reproducibility/internal_runs/dag_configs.py +++ b/dags/map_reproducibility/internal_runs/dag_configs.py @@ -225,18 +225,16 @@ } DAG_CONFIGS_ULTRA_NEMO = { - "recipes/a3ultra/nemo/a3ultra_llama3.1-70b_256gpus_fp8_nemo.yaml": { - "timeout_minutes": 25, - "backfill_group_nightly": 1, - "backfill_group_release": 1, - "nightly_schedule": Schedule.SATURDAY_PDT_12AM, + "recipes/a3ultra/nemo/a3ultra_llama3.1-8b_8gpus_fp8_nemo.yaml": { + "timeout_minutes": 20, "release_schedule": Schedule.SATURDAY_PDT_12AM, }, "recipes/a3ultra/nemo/a3ultra_llama3.1-8b_8gpus_bf16_nemo.yaml": { "timeout_minutes": 20, - "backfill_group_nightly": 1, - "backfill_group_release": 1, - "nightly_schedule": Schedule.SATURDAY_PDT_12AM, "release_schedule": Schedule.SATURDAY_PDT_12AM, }, + "recipes/a3ultra/nemo/a3ultra_llama3.1-70b_256gpus_fp8_nemo.yaml": { + "timeout_minutes": 25, + "release_schedule": Schedule.SATURDAY_PDT_3AM, + }, } diff --git a/dags/map_reproducibility/internal_runs/sample_a3ultra_maxtext_single_run.py b/dags/map_reproducibility/internal_runs/sample_a3ultra_maxtext_single_run.py index f66e20a0..a974353c 100644 --- a/dags/map_reproducibility/internal_runs/sample_a3ultra_maxtext_single_run.py +++ b/dags/map_reproducibility/internal_runs/sample_a3ultra_maxtext_single_run.py @@ -84,7 +84,7 @@ def main(): base_helm_repo_root=base_helm_repo_root, timeout=timeout, image_version=RELEASE_IMAGE, - sample_run_bucket_name=SAMPLE_RUN_BUCKET_NAME, + gcs_bucket_name=SAMPLE_RUN_BUCKET_NAME, workload_launcher=WorkloadLauncher.MAXTEXT_LAUNCHER, ) diff --git a/dags/map_reproducibility/internal_runs/sample_a3ultra_nemo_single_run.py b/dags/map_reproducibility/internal_runs/sample_a3ultra_nemo_single_run.py index 1b00e94f..1dc1da7b 100644 --- a/dags/map_reproducibility/internal_runs/sample_a3ultra_nemo_single_run.py +++ b/dags/map_reproducibility/internal_runs/sample_a3ultra_nemo_single_run.py @@ -82,7 +82,7 @@ def main(): # Setup configuration relative_config_yaml_path = ( - "recipes/a3ultra/nemo/a3ultra_llama3.1-8b_8gpus_bf16_nemo.yaml" + "recipes/a3ultra/nemo/a3ultra_llama3.1-8b_8gpus_fp8_nemo.yaml" ) timeout = DAG_CONFIGS_ULTRA_NEMO[relative_config_yaml_path]["timeout_minutes"] @@ -92,7 +92,7 @@ def main(): base_helm_repo_root=base_helm_repo_root, timeout=timeout, image_version=RELEASE_IMAGE, - sample_run_bucket_name=SAMPLE_RUN_BUCKET_NAME, + gcs_bucket_name=SAMPLE_RUN_BUCKET_NAME, workload_launcher=WorkloadLauncher.NEMO_TEN_LAUNCHER, ) diff --git a/dags/map_reproducibility/utils/common_utils.py b/dags/map_reproducibility/utils/common_utils.py index ed16a9db..7331e0e1 100644 --- a/dags/map_reproducibility/utils/common_utils.py +++ b/dags/map_reproducibility/utils/common_utils.py @@ -411,12 +411,8 @@ def helm_apply_cmds_internal_run( bucket_name=BUCKET_NAME, ): gcs_cmd = "" - if hypercomputer in ("a3ultra", "a4"): - if framework != "maxtext" and kueue_name: - gcs_cmd = f" --set queue={kueue_name}" - gcs_cmd += f" --set volumes.gcsMounts[0].bucketName={bucket_name}" - else: - gcs_cmd = f" --set workload.gcsBucketForDataCataPath={bucket_name}" + if framework == "maxtext": + gcs_cmd += f" --set volumes.gcsMounts[0].bucketName={bucket_name} " cluster_cmd = "" if framework == "nemo" and hypercomputer == "a3ultra": diff --git a/dags/map_reproducibility/utils/constants.py b/dags/map_reproducibility/utils/constants.py index b9536187..e700078c 100644 --- a/dags/map_reproducibility/utils/constants.py +++ b/dags/map_reproducibility/utils/constants.py @@ -1,6 +1,7 @@ KUEUE_NAME = "a3-ultra" OPTIMIZER = "adam" NUM_STEPS = 15 +BUCKET_NAME = "regression-testing-xlml" class Optimizer: @@ -39,6 +40,7 @@ class Schedule: WEEKDAY_PDT_6AM_7AM_EXCEPT_THURSDAY = "0 13,14 * * 2,3,4,6" SATURDAY_PDT_12AM = "0 7 * * 6" + SATURDAY_PDT_3AM = "0 10 * * 6" class Image: diff --git a/dags/map_reproducibility/utils/internal_aotc_workload.py b/dags/map_reproducibility/utils/internal_aotc_workload.py index ddd569d8..a4edd111 100644 --- a/dags/map_reproducibility/utils/internal_aotc_workload.py +++ b/dags/map_reproducibility/utils/internal_aotc_workload.py @@ -61,7 +61,7 @@ get_internal_run_type_and_comment, ) from dags.map_reproducibility.utils.sample_workload_utils import handle_profiler, assemble_sample_united_workload_commands, execute_workload_commands -from dags.map_reproducibility.utils.constants import Optimizer, KUEUE_NAME, NUM_STEPS +from dags.map_reproducibility.utils.constants import Optimizer, KUEUE_NAME, NUM_STEPS, BUCKET_NAME # Configure logging logger = logging.getLogger(__name__) @@ -309,7 +309,7 @@ def run_internal_united_workload( base_helm_repo_root: str, timeout: int, image_version: str, - sample_run_bucket_name: str, + gcs_bucket_name: str, workload_launcher: str, is_dag_run: bool = False, backfill: bool = False, @@ -393,18 +393,18 @@ def run_internal_united_workload( base_helm_repo_root, full_config_path, values_file_path, - sample_run_bucket_name, + gcs_bucket_name, container_timeout, tmpdir, ) - success, error = execute_workload_commands(commands, tmpdir) + success, error = execute_workload_commands(commands, cwd=tmpdir) if not success: return {"success": False, "error": error} gcs_bucket = get_job_gcs_bucket_folder( job_name, - bucket_name=sample_run_bucket_name, + bucket_name=gcs_bucket_name, framework=config.FRAMEWORK, gcs_experiment_folder_name="maxtext-experiments", ) @@ -464,7 +464,7 @@ def run_internal_dag_united_workload( base_helm_repo_root=None, timeout=timeout, image_version=image_version, - sample_run_bucket_name=None, + gcs_bucket_name=BUCKET_NAME, workload_launcher=workload_launcher, is_dag_run=True, backfill=backfill, diff --git a/dags/map_reproducibility/utils/sample_workload_utils.py b/dags/map_reproducibility/utils/sample_workload_utils.py index dd3dcd36..c76be8de 100644 --- a/dags/map_reproducibility/utils/sample_workload_utils.py +++ b/dags/map_reproducibility/utils/sample_workload_utils.py @@ -48,6 +48,7 @@ wait_for_jobsets_cmds, cleanup_existing_metrics_cmd, helm_apply_cmds_workload, + get_values_file_path, ) from dags.map_reproducibility.utils.benchmarkdb_utils import write_run diff --git a/dags/multipod/maxtext_multi_tier_checkpointing_recover.py b/dags/multipod/maxtext_multi_tier_checkpointing_recover.py new file mode 100644 index 00000000..54838dce --- /dev/null +++ b/dags/multipod/maxtext_multi_tier_checkpointing_recover.py @@ -0,0 +1,95 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +A DAG to run MaxText multi-tier checkpointing tests. +""" +import datetime +from airflow import models +from dags import composer_env, gcs_bucket +from dags.common.vm_resource import DockerImage, XpkClusters +from dags.multipod.configs import gke_config +from dags.multipod.configs.common import SetupMode # Run once a day at 10 am UTC (2 am PST) +from dags.common import test_owner + +SCHEDULED_TIME = "0 10 * * *" if composer_env.is_prod_env() else None + +with models.DAG( + dag_id="maxtext_multi_tier_checkpointing_recover", + schedule=SCHEDULED_TIME, + tags=[ + "multipod_team", + "maxtext", + "multi_tier_checkpointing_recover", + ], + start_date=datetime.datetime(2025, 5, 15), + catchup=False, + concurrency=2, +) as dag: + base_output_directory = ( + f"{gcs_bucket.MTC_BUCKET}/maxtext_multi_tier_checkpointing_recover" + ) + dataset_path = gcs_bucket.MTC_BUCKET + docker_images = [ + (SetupMode.STABLE, DockerImage.MAXTEXT_TPU_JAX_NIGHTLY), + ] + test_configs = { + # accelerator: list of slices to test + "v5p-8": [2], + } + clusters = { + # accelerator: cluster name + "v5p-8": XpkClusters.TPU_V5P_8_CLUSTER, + } + + for mode, image in docker_images: + for accelerator, slices in test_configs.items(): + for slice_num in slices: + command = ( + "bash end_to_end/test_mtc_phase_2_save_path.sh" + f" multitiercheckpointing-{slice_num}x-{accelerator}" + f" {base_output_directory} {dataset_path}", + ) + maxtext_save_checkpoint = gke_config.get_gke_config( + num_slices=slice_num, + cluster=clusters[accelerator], + time_out_in_min=60, + test_name="maxtext-multi-tier-checkpointing-recover", + run_model_cmds=command, + docker_image=image.value, + test_owner=test_owner.JACKY_F, + ).run_with_interruption( + ramdisk_directory="local", + xpk_branch="main", + skip_post_process=True, + mtc_enabled=True, + ) + + clean_cmd = (f"rm -rf /local/*",) + clean_ramdisk_one = gke_config.get_gke_config( + num_slices=slice_num, + cluster=clusters[accelerator], + time_out_in_min=60, + test_name="clean-ramdisk", + run_model_cmds=clean_cmd, + docker_image=image.value, + test_owner=test_owner.JACKY_F, + ).run( + ramdisk_directory="local", + xpk_branch="main", + skip_post_process=True, + mtc_enabled=True, + ) + + (maxtext_save_checkpoint >> clean_ramdisk_one) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index fc68d7be..bec66f7a 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -201,6 +201,41 @@ def run( return group + def run_with_interruption( + self, + *, + gcs_location: Optional[airflow.XComArg] = None, + use_vertex_tensorboard: bool = False, + use_pathways: bool = False, + skip_post_process: bool = False, + ramdisk_directory: str = "", + mtc_enabled: bool = False, + xpk_branch: str = xpk.MAIN_BRANCH, + ) -> DAGNode: + """Run a test job within a docker image. + + Attributes: + gcs_location: GCS path for all artifacts of the test. + use_vertex_tensorboard: Set to True to view workload data on + Vertex AI Tensorboard. + + Returns: + A task group with the following tasks chained: run_model and + post_process. + """ + with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: + run_model, gcs_path = self.run_model_with_interruption( + gcs_location, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + ) + if not skip_post_process: + run_model >> self.post_process(gcs_path) + return group + def run_with_name_gen_and_quarantine( self, quarantine_task_group, @@ -333,6 +368,72 @@ def run_model( ) return group, gcs_path + def run_model_with_interruption( + self, + gcs_location: Optional[airflow.XComArg] = None, + use_vertex_tensorboard: bool = False, + use_pathways: bool = False, + ramdisk_directory: str = "", + mtc_enabled: bool = False, + xpk_branch: str = xpk.MAIN_BRANCH, + ) -> DAGNode: + """Run the TPU/GPU test in `task_test_config` using xpk. + Different behaviour for testing interruption. + + Attributes: + gcs_location: GCS path for all artifacts of the test. + use_vertex_tensorboard: Set to True to view workload data on + Vertex AI Tensorboard. + + Returns: + A DAG node that executes the model test. + """ + with TaskGroup(group_id="run_model") as group: + workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id) + if gcs_location: + gcs_path = gcs_location + else: + gcs_path = name_format.generate_gcs_folder_location( + self.task_test_config.gcs_subfolder, + self.task_test_config.benchmark_id, + ) + + launch_workload_with_interruption = ( + self.launch_workload_with_interruption( + workload_id, + gcs_path, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + ) + ) + + wait_for_workload_completion = xpk.wait_for_workload_completion.override( + timeout=int(self.task_test_config.timeout.total_seconds()), + )( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + region=self.task_gcp_config.zone[:-2], + cluster_name=self.task_test_config.cluster_name, + ) + + clean_up_workload = xpk.clean_up_workload( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + zone=self.task_gcp_config.zone, + cluster_name=self.task_test_config.cluster_name, + ) + + ( + (workload_id, gcs_path) + >> launch_workload_with_interruption + >> wait_for_workload_completion + >> clean_up_workload + ) + return group, gcs_path + def launch_workload( self, workload_id: str, @@ -376,6 +477,61 @@ def launch_workload( run_workload >> wait_for_workload_start return group + def launch_workload_with_interruption( + self, + workload_id: str, + gcs_path: str, + use_vertex_tensorboard: bool, + use_pathways: bool = False, + ramdisk_directory: str = "", + mtc_enabled: bool = False, + xpk_branch: str = xpk.MAIN_BRANCH, + ) -> DAGNode: + """Create the workload and wait for it to provision.""" + with TaskGroup(group_id="launch_workload_with_interruption") as group: + run_workload = xpk.run_workload.override( + owner=self.task_test_config.task_owner + )( + task_id="run_workload", + cluster_project=self.task_gcp_config.project_name, + zone=self.task_gcp_config.zone, + cluster_name=self.task_test_config.cluster_name, + benchmark_id=self.task_test_config.benchmark_id, + workload_id=workload_id, + gcs_path=gcs_path, + docker_image=self.task_test_config.docker_image, + accelerator_type=self.task_test_config.accelerator.name, + run_cmds=self.task_test_config.test_script, + num_slices=self.task_test_config.num_slices, + use_vertex_tensorboard=use_vertex_tensorboard, + use_pathways=use_pathways, + ramdisk_directory=ramdisk_directory, + mtc_enabled=mtc_enabled, + xpk_branch=xpk_branch, + ) + + wait_for_workload_start = xpk.wait_for_workload_start.override( + timeout=self.workload_provision_timeout.total_seconds() + )( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + region=self.task_gcp_config.zone[:-2], + cluster_name=self.task_test_config.cluster_name, + ) + + run_interruption_workload = xpk.run_interruption_cmd.override( + owner=self.task_test_config.task_owner + )( + task_id="run_interruption_cmd", + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + region=self.task_gcp_config.zone[:-2], + cluster_name=self.task_test_config.cluster_name, + ) + + run_workload >> wait_for_workload_start >> run_interruption_workload + return group + def post_process(self, result_location: Optional[str] = None) -> DAGNode: """Process metrics and metadata, and insert them into BigQuery tables. diff --git a/xlml/utils/xpk.py b/xlml/utils/xpk.py index 5b287d67..085a90b8 100644 --- a/xlml/utils/xpk.py +++ b/xlml/utils/xpk.py @@ -120,6 +120,12 @@ def run_workload( workload_create_cmd += f" --ramdisk-directory={ramdisk_directory}" 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: @@ -284,6 +290,51 @@ def wait_for_workload_completion( return True +@task.sensor(poke_interval=120, timeout=1200, mode="reschedule") +def run_interruption_cmd( + task_id: str, + project_id: str, + region: str, + cluster_name: str, + workload_id: str, +) -> bool: + """Run command to interrupt pod .""" + core_api = _get_core_api_client(project_id, region, cluster_name) + pods = _list_workload_pods(core_api, workload_id) + + if any(pod.status.phase in ["Pending"] for pod in pods.items): + logging.info("Some of the pods is still pending. Waiting to start") + return False + + try: + for pod in pods.items: + if pod.status.phase == "Failed": + # Don't keep retrying if the pod has failed + raise AirflowFailException(f"Bad pod phase: {pod.status.phase}") + elif pod.status.phase in ["Unknown"]: + raise RuntimeError(f"Bad pod phase: {pod.status.phase}") + finally: + if all(pod.status.phase in ["Running"] for pod in pods.items): + # Pick last one running pod + pod = pods.items[len(pods.items) - 1] + logs = core_api.read_namespaced_pod_log( + name=pod.metadata.name, namespace=pod.metadata.namespace + ) + logging.info(f"Logs for pod {pod.metadata.name}:") + for line in logs.split("\n"): + logging.info(line) + + # TODO --> More sophisticated way to know when the pod start training. + if "completed step:" in logs: + # Here where regex expresion to kill pod will go. First test with killing the pod + result = core_api.delete_namespaced_pod( + name=pod.metadata.name, namespace=pod.metadata.namespace + ) + logging.info("The {pod.metadata.name} pod was successfully deleted.") + return True + return False + + @task(trigger_rule="all_done") def clean_up_workload( workload_id: str,