From 71b8d8a5c3e9d19740c46d563f4127425c88fecb Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Tue, 12 Aug 2025 10:03:26 +0800 Subject: [PATCH 01/18] [cienet-private] Enable workflow checks for the primary working branches --- .github/workflows/dag-check.yml | 1 - .github/workflows/pyink-check.yml | 3 +- .github/workflows/pylint-check.yml | 3 +- .github/workflows/require-checklist.yml | 5 ++- .github/workflows/unit-test.yml | 1 - scripts/code-style.sh | 41 +++++++++++++++++++------ 6 files changed, 40 insertions(+), 14 deletions(-) diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml index 96fcb76bf..ac3982d0e 100644 --- a/.github/workflows/dag-check.yml +++ b/.github/workflows/dag-check.yml @@ -3,7 +3,6 @@ name: DAG Check on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/.github/workflows/pyink-check.yml b/.github/workflows/pyink-check.yml index ea7a8f250..c5cb1cd83 100644 --- a/.github/workflows/pyink-check.yml +++ b/.github/workflows/pyink-check.yml @@ -2,11 +2,12 @@ name: Formatter on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: branches: [master] + workflow_dispatch: {} + jobs: format_check: runs-on: ubuntu-latest diff --git a/.github/workflows/pylint-check.yml b/.github/workflows/pylint-check.yml index 5e23d2812..02f4cb31e 100644 --- a/.github/workflows/pylint-check.yml +++ b/.github/workflows/pylint-check.yml @@ -2,12 +2,13 @@ name: Linter on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: branches: [master] + workflow_dispatch: {} + jobs: linting_check: runs-on: ubuntu-latest diff --git a/.github/workflows/require-checklist.yml b/.github/workflows/require-checklist.yml index d15d19d99..4da288575 100644 --- a/.github/workflows/require-checklist.yml +++ b/.github/workflows/require-checklist.yml @@ -2,10 +2,13 @@ name: Require Checklist on: pull_request: types: [opened, edited, synchronize] + + workflow_dispatch: {} + jobs: check_pr_body: runs-on: ubuntu-latest steps: - uses: mheap/require-checklist-action@v2 with: - requireChecklist: false # If this is true and there are no checklists detected, the action will fail \ No newline at end of file + requireChecklist: false # If this is true and there are no checklists detected, the action will fail diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 4f1723cb2..b771ca7e9 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -3,7 +3,6 @@ name: Unit Test on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/scripts/code-style.sh b/scripts/code-style.sh index e02e53c4e..5ade7a620 100755 --- a/scripts/code-style.sh +++ b/scripts/code-style.sh @@ -22,14 +22,37 @@ FOLDERS_TO_FORMAT=("dags" "xlml") echo "[code-style] Running Semgrep static analysis..." semgrep --config scripts/semgrep-rules.yaml --error -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff -done - -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pylint "./$folder" --fail-under=9.6 -done +HEAD_SHA="$(git rev-parse HEAD)" +BASE_BRANCH="dev" + +if ! git rev-parse --verify "$BASE_BRANCH" >/dev/null 2>&1; then + git fetch origin "$BASE_BRANCH":"$BASE_BRANCH" || { + echo "[code-style] base branch '$BASE_BRANCH' not found, skip diff-based check." + exit 0 + } +fi + +CHANGED_PY_FILES="$( + git diff --name-only --diff-filter=ACM "${BASE_BRANCH}" "${HEAD_SHA}" \ + | grep '\.py$' \ + | while read -r f; do + for folder in "${FOLDERS_TO_FORMAT[@]}"; do + if [[ "$f" == "$folder/"* ]]; then + echo "$f" + break + fi + done + done \ + | sort -u +)" + +if [[ -z "${CHANGED_PY_FILES}" ]]; then + echo "[pre-push hook] no changed files detected between ${HEAD_SHA} and ${BASE_BRANCH}" + exit 1 +fi + +pyink ${CHANGED_PY_FILES} --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff + +pylint ${CHANGED_PY_FILES} --fail-under=9.6 --disable=E1123 echo "Successfully clean up all codes." From aaab549b2041cd3c5554459c22f6a9629433a13b Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Tue, 7 Jul 2026 19:40:43 +0000 Subject: [PATCH 02/18] Apply the changes to the affected DAGs --- dags/examples/maxtext_profile_sweep_example_dag.py | 4 ++-- dags/orbax/maxtext_emc_restore_gcs.py | 7 ++++--- dags/orbax/maxtext_emc_restore_local.py | 5 +++-- dags/orbax/maxtext_mtc_restore_local.py | 5 +++-- dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py | 7 ++++--- 5 files changed, 16 insertions(+), 12 deletions(-) diff --git a/dags/examples/maxtext_profile_sweep_example_dag.py b/dags/examples/maxtext_profile_sweep_example_dag.py index d6d6dc1d1..0dc5cf249 100644 --- a/dags/examples/maxtext_profile_sweep_example_dag.py +++ b/dags/examples/maxtext_profile_sweep_example_dag.py @@ -14,7 +14,7 @@ """ An example DAG to extract profile metrics from pretraining mixtral-8x7b model on v6-256. -Profile extraction can be seamlessly integrated with maxtext_sweep_gke_config + run_with_name_gen_and_quarantine. +Profile extraction can be seamlessly integrated with maxtext_sweep_gke_config + (to_name_gen_and_quarantine_task + run). """ import datetime @@ -142,4 +142,4 @@ def dict_to_arg(param_dict): ) for test in maxtext_sweep_gke_test: - test.run_with_name_gen_and_quarantine(quarantine_task_group) + test.to_name_gen_and_quarantine_task(quarantine_task_group).run() diff --git a/dags/orbax/maxtext_emc_restore_gcs.py b/dags/orbax/maxtext_emc_restore_gcs.py index 28420171b..f6aa66d0e 100644 --- a/dags/orbax/maxtext_emc_restore_gcs.py +++ b/dags/orbax/maxtext_emc_restore_gcs.py @@ -128,12 +128,13 @@ run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.DEPP_L, - ).run_with_node_interruption( + ).to_node_interruption_task( + expect_reach_to_step=step_to_interrupt, + last_node=True, + ).run( ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, mtc_enabled=True, skip_post_process=True, - last_node=True, - expect_reach_to_step=step_to_interrupt, max_restart=15, ) diff --git a/dags/orbax/maxtext_emc_restore_local.py b/dags/orbax/maxtext_emc_restore_local.py index 7e3c04b3d..6eccaec94 100644 --- a/dags/orbax/maxtext_emc_restore_local.py +++ b/dags/orbax/maxtext_emc_restore_local.py @@ -126,11 +126,12 @@ run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.DEPP_L, - ).run_with_node_interruption( + ).to_node_interruption_task( + expect_reach_to_step=step_to_interrupt, + ).run( ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, mtc_enabled=True, skip_post_process=True, - expect_reach_to_step=step_to_interrupt, max_restart=15, ) diff --git a/dags/orbax/maxtext_mtc_restore_local.py b/dags/orbax/maxtext_mtc_restore_local.py index 66e8f445f..3b044a182 100644 --- a/dags/orbax/maxtext_mtc_restore_local.py +++ b/dags/orbax/maxtext_mtc_restore_local.py @@ -129,11 +129,12 @@ run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.DEPP_L, - ).run_with_node_interruption( + ).to_node_interruption_task( + expect_reach_to_step=step_to_interrupt, + ).run( ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, mtc_enabled=True, skip_post_process=True, - expect_reach_to_step=step_to_interrupt, max_restart=15, ) diff --git a/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py b/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py index e94cbdbed..23f834827 100644 --- a/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py +++ b/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py @@ -138,13 +138,14 @@ def generate_workload_checkpoints_location(gcs_ckpt_location: str) -> str: run_model_cmds=workload_command, docker_image=image.value, test_owner=test_owner.SHARON_Y, - ).run_with_node_interruption( + ).to_node_interruption_task( + expect_reach_to_step=step_to_interrupt, + check_file_exists=True, + ).run( gcs_location=gcs_location, xpk_branch=MAIN_BRANCH, skip_post_process=True, - expect_reach_to_step=step_to_interrupt, max_restart=15, - check_file_exists=True, ) end_time = validation_util.generate_timestamp.override( From 035b5304b66c7515ec3bdc28400d8e1da251207b Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Tue, 7 Jul 2026 19:52:42 +0000 Subject: [PATCH 03/18] Fix format/linting errors --- dags/orbax/maxtext_emc_restore_gcs.py | 36 ++-- dags/orbax/maxtext_emc_restore_local.py | 34 ++-- dags/orbax/maxtext_mtc_restore_local.py | 34 ++-- ...xt_reg_restore_gcs_with_node_disruption.py | 36 ++-- xlml/apis/task.py | 188 +++++++++++++++++- 5 files changed, 258 insertions(+), 70 deletions(-) diff --git a/dags/orbax/maxtext_emc_restore_gcs.py b/dags/orbax/maxtext_emc_restore_gcs.py index f6aa66d0e..80eb9d308 100644 --- a/dags/orbax/maxtext_emc_restore_gcs.py +++ b/dags/orbax/maxtext_emc_restore_gcs.py @@ -120,22 +120,26 @@ task_id="generate_start_time" )() - 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}-emc", - run_model_cmds=workload_command, - docker_image=image.value, - test_owner=test_owner.DEPP_L, - ).to_node_interruption_task( - expect_reach_to_step=step_to_interrupt, - last_node=True, - ).run( - ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, - mtc_enabled=True, - skip_post_process=True, - max_restart=15, + 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}-emc", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.DEPP_L, + ) + .to_node_interruption_task( + expect_reach_to_step=step_to_interrupt, + last_node=True, + ) + .run( + ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, + mtc_enabled=True, + skip_post_process=True, + max_restart=15, + ) ) end_time = validation_util.generate_timestamp.override( diff --git a/dags/orbax/maxtext_emc_restore_local.py b/dags/orbax/maxtext_emc_restore_local.py index 6eccaec94..4c75b32dc 100644 --- a/dags/orbax/maxtext_emc_restore_local.py +++ b/dags/orbax/maxtext_emc_restore_local.py @@ -118,21 +118,25 @@ task_id="generate_start_time" )() - 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}-emc", - run_model_cmds=workload_command, - docker_image=image.value, - test_owner=test_owner.DEPP_L, - ).to_node_interruption_task( - expect_reach_to_step=step_to_interrupt, - ).run( - ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, - mtc_enabled=True, - skip_post_process=True, - max_restart=15, + 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}-emc", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.DEPP_L, + ) + .to_node_interruption_task( + expect_reach_to_step=step_to_interrupt, + ) + .run( + ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, + mtc_enabled=True, + skip_post_process=True, + max_restart=15, + ) ) end_time = validation_util.generate_timestamp.override( diff --git a/dags/orbax/maxtext_mtc_restore_local.py b/dags/orbax/maxtext_mtc_restore_local.py index 3b044a182..9130625a7 100644 --- a/dags/orbax/maxtext_mtc_restore_local.py +++ b/dags/orbax/maxtext_mtc_restore_local.py @@ -121,21 +121,25 @@ task_id="generate_start_time" )() - 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}-mtc", - run_model_cmds=workload_command, - docker_image=image.value, - test_owner=test_owner.DEPP_L, - ).to_node_interruption_task( - expect_reach_to_step=step_to_interrupt, - ).run( - ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, - mtc_enabled=True, - skip_post_process=True, - max_restart=15, + 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}-mtc", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.DEPP_L, + ) + .to_node_interruption_task( + expect_reach_to_step=step_to_interrupt, + ) + .run( + ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, + mtc_enabled=True, + skip_post_process=True, + max_restart=15, + ) ) end_time = validation_util.generate_timestamp.override( diff --git a/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py b/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py index 23f834827..e30731f47 100644 --- a/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py +++ b/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py @@ -130,22 +130,26 @@ def generate_workload_checkpoints_location(gcs_ckpt_location: str) -> str: task_id="generate_start_time" )() - 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}", - run_model_cmds=workload_command, - docker_image=image.value, - test_owner=test_owner.SHARON_Y, - ).to_node_interruption_task( - expect_reach_to_step=step_to_interrupt, - check_file_exists=True, - ).run( - gcs_location=gcs_location, - xpk_branch=MAIN_BRANCH, - skip_post_process=True, - max_restart=15, + 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}", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.SHARON_Y, + ) + .to_node_interruption_task( + expect_reach_to_step=step_to_interrupt, + check_file_exists=True, + ) + .run( + gcs_location=gcs_location, + xpk_branch=MAIN_BRANCH, + skip_post_process=True, + max_restart=15, + ) ) end_time = validation_util.generate_timestamp.override( diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 9877f2bf3..c8c4cb7ef 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -379,7 +379,38 @@ def run( def run_with_node_interruption( self, - *, + expect_reach_to_step: int, + last_node: bool = False, + check_file_exists: bool = False, + ) -> "XpkNodeInterruptionTask": + return XpkNodeInterruptionTask( + task_test_config=self.task_test_config, + task_gcp_config=self.task_gcp_config, + task_metric_config=self.task_metric_config, + workload_provision_timeout=self.workload_provision_timeout, + expect_reach_to_step=expect_reach_to_step, + last_node=last_node, + check_file_exists=check_file_exists, + ) + + def to_name_gen_and_quarantine_task( + self, + quarantine_task_group, + run_name_env: str = "M_RUN_NAME", + nested_run_name_in_tb_file_location: bool = True, + ) -> "XpkNameGenAndQuarantineTask": + return XpkNameGenAndQuarantineTask( + task_test_config=self.task_test_config, + task_gcp_config=self.task_gcp_config, + task_metric_config=self.task_metric_config, + workload_provision_timeout=self.workload_provision_timeout, + quarantine_task_group=quarantine_task_group, + run_name_env=run_name_env, + nested_run_name_in_tb_file_location=nested_run_name_in_tb_file_location, + ) + + def run_model( + self, gcs_location: Optional[airflow.XComArg] = None, use_vertex_tensorboard: bool = False, expect_reach_to_step: int, @@ -415,9 +446,151 @@ def run_with_node_interruption( check_file_exists: By default, this is False. If set to True, task branch task_path_decider will be performed. Returns: - A task group with the following tasks chained: run_model and - post_process. + 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 = self.launch_workload( + workload_id, + gcs_path, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + max_restart, + ) + 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=gke.zone_to_region(self.task_gcp_config.zone), + 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, + xpk_branch=xpk_branch, + ) + + _ = ( + (workload_id, gcs_path) + >> launch_workload + >> wait_for_workload_completion + >> clean_up_workload + ) + return group, gcs_path + + def launch_workload( + 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, + max_restart: int = 0, + ) -> DAGNode: + """Create the workload and wait for it to provision.""" + with TaskGroup(group_id="launch_workload") 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, + max_restart=max_restart, + ) + 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=gke.zone_to_region(self.task_gcp_config.zone), + cluster_name=self.task_test_config.cluster_name, + ) + _ = run_workload >> wait_for_workload_start + return group + + def post_process(self, result_location: Optional[str] = None) -> DAGNode: + """Process metrics and metadata, and insert them into BigQuery tables. + + Returns: + A DAG node that executes the post process. """ + with TaskGroup(group_id="post_process") as group: + process_id = metric.generate_process_id.override(retries=0)() + post_process_metrics = metric.process_metrics.override(retries=0)( + process_id, + self.task_test_config, + self.task_metric_config, + self.task_gcp_config, + folder_location=result_location, + ) + + if self.task_metric_config and self.task_metric_config.profile: + self.task_metric_config.profile.metrics = ( + metric.xplane_to_metrics.override(retries=0)( + self.task_metric_config.profile.file_location + ) + ) + _ = ( + process_id + >> self.task_metric_config.profile.metrics + >> post_process_metrics + ) + else: + _ = process_id >> post_process_metrics + + return group + + +@dataclasses.dataclass +class XpkNodeInterruptionTask(XpkTask): + """Task for running XPK workloads with node interruption.""" + + expect_reach_to_step: int = 0 + last_node: bool = False + check_file_exists: bool = False + + def run( + 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, + max_restart: int = 0, + ) -> DAGNode: + """Run a test job with node interruption.""" with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: run_model, gcs_path = self.run_model_with_node_interruption( gcs_location=gcs_location, @@ -700,12 +873,11 @@ def run_with_run_name_generation( run_model, gcs_path = self.run_model( use_pathways=use_pathways, xpk_branch=xpk_branch ) - _ = ( - run_name - >> (tb_file_location, profile_file_location) - >> run_model - >> self.post_process(gcs_path) + flow = ( + run_name >> (tb_file_location, profile_file_location) >> run_model ) + if not skip_post_process: + _ = flow >> self.post_process(gcs_path) else: run_model, gcs_path = self.run_model( use_pathways=use_pathways, xpk_branch=xpk_branch From 24bca5de172dd92a289a98123bfd3f85ccc93c42 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Tue, 7 Jul 2026 19:57:40 +0000 Subject: [PATCH 04/18] Fix linting errors --- .../maxtext_profile_sweep_example_dag.py | 69 ++++++++++++------- xlml/apis/task.py | 41 ++++------- 2 files changed, 57 insertions(+), 53 deletions(-) diff --git a/dags/examples/maxtext_profile_sweep_example_dag.py b/dags/examples/maxtext_profile_sweep_example_dag.py index 0dc5cf249..24fc7bc6e 100644 --- a/dags/examples/maxtext_profile_sweep_example_dag.py +++ b/dags/examples/maxtext_profile_sweep_example_dag.py @@ -12,16 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -An example DAG to extract profile metrics from pretraining mixtral-8x7b model on v6-256. -Profile extraction can be seamlessly integrated with maxtext_sweep_gke_config + (to_name_gen_and_quarantine_task + run). +"""An example DAG to extract profile metrics. + +Pretraining mixtral-8x7b model on v6-256. +Profile extraction can be seamlessly integrated with +maxtext_sweep_gke_config + (to_name_gen_and_quarantine_task + run). """ import datetime from airflow import models from airflow.utils.task_group import TaskGroup from dags.common import test_owner -from dags.common.vm_resource import XpkClusters, DockerImage, Project +from dags.common.vm_resource import XpkClusters, Project from dags.multipod.configs import maxtext_sweep_gke_config from xlml.apis import metric_config @@ -35,7 +37,10 @@ def dict_to_arg(param_dict): docker_image = { - "stable": "gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:2025-05-20", + "stable": ( + "gcr.io/tpu-prod-env-multipod/" + "maxtext_jax_stable_stack:2025-05-20" + ), } # https://github.com/AI-Hypercomputer/maxtext/blob/main/benchmarks/maxtext_trillium_model_configs.py @@ -45,9 +50,14 @@ def dict_to_arg(param_dict): "cluster": XpkClusters.TPU_V6E_256_MLPERF_CLUSTER, "train_command": [ f"export BASE_OUTPUT_PATH={BASE_OUTPUT_PATH} && " - "python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml base_output_directory=${BASE_OUTPUT_PATH} model_name=mixtral-8x7b " - # add profiler config: ensure steps > skip_first_n_steps_for_profiler + profiler_steps - "steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3 " + "python3 -m maxtext.trainers.pre_train.train " + "src/maxtext/configs/base.yml " + "base_output_directory=${BASE_OUTPUT_PATH} " + "model_name=mixtral-8x7b " + # add profiler config: ensure steps > + # skip_first_n_steps_for_profiler + profiler_steps + "steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 " + "profiler_steps=3 " + dict_to_arg({ "per_device_batch_size": 12, "ici_fsdp_parallelism": -1, @@ -81,9 +91,14 @@ def dict_to_arg(param_dict): "base_output_directory": "gs://runner-maxtext-logs", "train_command": [ f"export BASE_OUTPUT_PATH={BASE_OUTPUT_PATH} && " - "python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml base_output_directory=${BASE_OUTPUT_PATH} model_name=mixtral-8x7b " - # add profiler config: ensure steps > skip_first_n_steps_for_profiler + profiler_steps - "steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3 " + "python3 -m maxtext.trainers.pre_train.train " + "src/maxtext/configs/base.yml " + "base_output_directory=${BASE_OUTPUT_PATH} " + "model_name=mixtral-8x7b " + # add profiler config: ensure steps > + # skip_first_n_steps_for_profiler + profiler_steps + "steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 " + "profiler_steps=3 " + dict_to_arg({ "per_device_batch_size": 12, "ici_fsdp_parallelism": -1, @@ -120,25 +135,27 @@ def dict_to_arg(param_dict): ) for run_name, test_scripts_details in test_models_tpu.items(): - for image in docker_image.keys(): + for image, img_val in docker_image.items(): # sweep num_slices and other training params # generate run_name and tensorboard/profile location for extraction num_slices = [1] sweep_params = {} - maxtext_sweep_gke_test = maxtext_sweep_gke_config.get_maxtext_sweep_gke_config( - test_owner=test_owner.SHUNING_J, - dataset_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, - composer_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, - dataset_name=metric_config.DatasetOption.XLML_DATASET, - cluster=test_scripts_details["cluster"], - time_out_in_min=test_scripts_details["time_out_in_min"], - base_output_directory=BASE_OUTPUT_PATH, - num_slices=num_slices, - docker_image=docker_image[image], - run_name_prefix=f"maxtext_{image}_{run_name}", - base_run_model_cmds=test_scripts_details["train_command"], - sweep_params=sweep_params, - enable_profile_config=True, # add flag to enable profile extraction + maxtext_sweep_gke_test = ( + maxtext_sweep_gke_config.get_maxtext_sweep_gke_config( + test_owner=test_owner.SHUNING_J, + dataset_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, + composer_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, + dataset_name=metric_config.DatasetOption.XLML_DATASET, + cluster=test_scripts_details["cluster"], + time_out_in_min=test_scripts_details["time_out_in_min"], + base_output_directory=BASE_OUTPUT_PATH, + num_slices=num_slices, + docker_image=img_val, + run_name_prefix=f"maxtext_{image}_{run_name}", + base_run_model_cmds=test_scripts_details["train_command"], + sweep_params=sweep_params, + enable_profile_config=True, # add flag to enable profile extraction + ) ) for test in maxtext_sweep_gke_test: diff --git a/xlml/apis/task.py b/xlml/apis/task.py index c8c4cb7ef..8cb0e628a 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -18,7 +18,7 @@ import dataclasses import datetime import shlex -from typing import Optional, Tuple, Union +from typing import Any, Optional, Tuple, Union import airflow from airflow.models.taskmixin import DAGNode @@ -42,7 +42,7 @@ def run(self) -> DAGNode: Returns: A DAG node that executes this test. """ - ... + pass def run_with_quarantine(self, quarantine_task_group): """Run a test job. If the test job is flaky, wrap it in a special task grop. @@ -67,7 +67,7 @@ def run_queued_resource_test( tpu_name_env_var: bool = False, all_workers: bool = True, skip_post_process: bool = False, - custom_env: dict[str, str] = {}, + custom_env: Optional[dict[str, str]] = None, ): """This is a class to set up tasks for TPU provisioned by Queued Resource. @@ -94,6 +94,9 @@ def run_queued_resource_test( post_process and clean_up. """ + if custom_env is None: + custom_env = {} + with TaskGroup( group_id=task_test_config.benchmark_id, prefix_group_id=True ) as test: @@ -335,7 +338,8 @@ class XpkTask(BaseTask): task_gcp_config: gcp_config.GCPConfig task_metric_config: Optional[metric_config.MetricConfig] = None workload_provision_timeout: datetime.timedelta = datetime.timedelta( - # Set the provision timeout from 300 to 60 minutes for decreasing the duration of failed tasks + # Set the provision timeout from 300 to 60 minutes for decreasing the + # duration of failed tasks minutes=60 ) @@ -624,26 +628,7 @@ def run_model_with_node_interruption( ) -> DAGNode: """Run the TPU/GPU test in `task_test_config` using xpk. - Different behavior for testing node 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. - expect_reach_to_step: The training step at which the node interruption - should be triggered. - use_pathways: Set to True to use the Pathways execution framework. - ramdisk_directory: The directory for enabling emergency checkpointing. - mtc_enabled: Set to True to enable Multi-tier Checkpointing (MTC). - xpk_branch: The specific git branch of the xpk tool to use. - last_node: If True, the interruption will target the last node in the - workload; otherwise, it targets the first node. - max_restart: By default, this is 0. - This will restart the job with flag "--max-restarts" - check_file_exists: By default, this is False. If set to True, - task branch task_path_decider will be performed. - Returns: - A DAG node that executes the model test. + with interruption. """ with TaskGroup(group_id="run_model") as group: workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id) @@ -772,15 +757,17 @@ def launch_workload_with_node_reach_to_step( wait_for_file_to_exist = gcs.wait_for_file_to_exist.override( task_id=task_id_wait_file_exist )( - file_path=f"{gcs_path}/{str(expect_reach_to_step)}/commit_success.txt", + file_path=( + f"{gcs_path}/{str(expect_reach_to_step)}/commit_success.txt" + ), ) task_id_do_nothing = "do_nothing" do_nothing = EmptyOperator(task_id=task_id_do_nothing) @task.branch def task_path_decider(check_file_exists: bool = False) -> str: - """ - Dynamically route the workflow depending on the `check_file_exists`. + """Dynamically route the workflow depending on the + `check_file_exists`. """ if check_file_exists: return f"{group.group_id}.{task_id_wait_file_exist}" From 386c7058e539c9012456e67b737ff8bd768390c7 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Wed, 8 Jul 2026 01:54:45 +0000 Subject: [PATCH 05/18] Fix format/linting errors and apply changes to DAGs that're affected --- ...text_aqtp_version_sweep_gke_example_dag.py | 2 +- .../maxtext_profile_namegen_example_dag.py | 29 +++++++++------- .../maxtext_profile_sweep_example_dag.py | 33 +++++++++---------- .../examples/maxtext_sweep_gke_example_dag.py | 2 +- dags/multipod/maxtext_convergence.py | 24 ++++++++------ xlml/apis/task.py | 4 +-- 6 files changed, 50 insertions(+), 44 deletions(-) diff --git a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py index deaeb3a17..7d66275eb 100644 --- a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py +++ b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py @@ -98,4 +98,4 @@ # Run jobs for test in tests: - test.run_with_run_name_generation() + test.to_name_gen_and_quarantine_task().run() diff --git a/dags/examples/maxtext_profile_namegen_example_dag.py b/dags/examples/maxtext_profile_namegen_example_dag.py index bf108062e..efeec9a1a 100644 --- a/dags/examples/maxtext_profile_namegen_example_dag.py +++ b/dags/examples/maxtext_profile_namegen_example_dag.py @@ -14,7 +14,7 @@ """ An example DAG to extract profile metrics from pretraining mixtral-8x7b model on 1xv4-128. -Profile extraction can be easily integrated with gke_config + run_with_run_name_generation. +Profile extraction can be easily integrated with gke_config + (to_name_gen_and_quarantine_task + run). """ import datetime @@ -88,7 +88,8 @@ ) as dag: for run_name, test_scripts_details in test_models_tpu.items(): for image in docker_image.keys(): - # file_location: pass in base_output_directory, will be altered in `run_with_run_name_generation` + # file_location: pass in base_output_directory, will be altered in + # XpkNameGenAndQuarantineTask.run_with_run_name_generation job_metric_config = metric_config.MetricConfig() # optionally, add tensorboard metrics job_metric_config.tensorboard_summary = metric_config.SummaryConfig( @@ -105,13 +106,17 @@ if "not_add_profile_config" in test_scripts_details: job_metric_config.profile = None - tpu_task = gke_config.get_gke_config( - num_slices=1, - time_out_in_min=test_scripts_details["time_out_in_min"], - test_name=f"maxtext_{image}_{run_name}", - run_model_cmds=test_scripts_details["train_command"], - docker_image=docker_image[image], - test_owner=test_owner.SHUNING_J, - cluster=test_scripts_details["cluster"], - user_specified_job_metric_config=job_metric_config, # customize config - ).run_with_run_name_generation(run_name_env="RUN_NAME") + tpu_task = ( + gke_config.get_gke_config( + num_slices=1, + time_out_in_min=test_scripts_details["time_out_in_min"], + test_name=f"maxtext_{image}_{run_name}", + run_model_cmds=test_scripts_details["train_command"], + docker_image=docker_image[image], + test_owner=test_owner.SHUNING_J, + cluster=test_scripts_details["cluster"], + user_specified_job_metric_config=job_metric_config, # customize config + ) + .to_name_gen_and_quarantine_task(run_name_env="RUN_NAME") + .run() + ) diff --git a/dags/examples/maxtext_profile_sweep_example_dag.py b/dags/examples/maxtext_profile_sweep_example_dag.py index 24fc7bc6e..612531691 100644 --- a/dags/examples/maxtext_profile_sweep_example_dag.py +++ b/dags/examples/maxtext_profile_sweep_example_dag.py @@ -38,8 +38,7 @@ def dict_to_arg(param_dict): docker_image = { "stable": ( - "gcr.io/tpu-prod-env-multipod/" - "maxtext_jax_stable_stack:2025-05-20" + "gcr.io/tpu-prod-env-multipod/" "maxtext_jax_stable_stack:2025-05-20" ), } @@ -140,22 +139,20 @@ def dict_to_arg(param_dict): # generate run_name and tensorboard/profile location for extraction num_slices = [1] sweep_params = {} - maxtext_sweep_gke_test = ( - maxtext_sweep_gke_config.get_maxtext_sweep_gke_config( - test_owner=test_owner.SHUNING_J, - dataset_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, - composer_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, - dataset_name=metric_config.DatasetOption.XLML_DATASET, - cluster=test_scripts_details["cluster"], - time_out_in_min=test_scripts_details["time_out_in_min"], - base_output_directory=BASE_OUTPUT_PATH, - num_slices=num_slices, - docker_image=img_val, - run_name_prefix=f"maxtext_{image}_{run_name}", - base_run_model_cmds=test_scripts_details["train_command"], - sweep_params=sweep_params, - enable_profile_config=True, # add flag to enable profile extraction - ) + maxtext_sweep_gke_test = maxtext_sweep_gke_config.get_maxtext_sweep_gke_config( + test_owner=test_owner.SHUNING_J, + dataset_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, + composer_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, + dataset_name=metric_config.DatasetOption.XLML_DATASET, + cluster=test_scripts_details["cluster"], + time_out_in_min=test_scripts_details["time_out_in_min"], + base_output_directory=BASE_OUTPUT_PATH, + num_slices=num_slices, + docker_image=img_val, + run_name_prefix=f"maxtext_{image}_{run_name}", + base_run_model_cmds=test_scripts_details["train_command"], + sweep_params=sweep_params, + enable_profile_config=True, # add flag to enable profile extraction ) for test in maxtext_sweep_gke_test: diff --git a/dags/examples/maxtext_sweep_gke_example_dag.py b/dags/examples/maxtext_sweep_gke_example_dag.py index 10d76babd..464a26c39 100644 --- a/dags/examples/maxtext_sweep_gke_example_dag.py +++ b/dags/examples/maxtext_sweep_gke_example_dag.py @@ -57,4 +57,4 @@ # Run jobs for test in maxtext_sweep_gke_test: - test.run_with_run_name_generation() + test.to_name_gen_and_quarantine_task().run() diff --git a/dags/multipod/maxtext_convergence.py b/dags/multipod/maxtext_convergence.py index fcb7d9f50..f5e476eaa 100644 --- a/dags/multipod/maxtext_convergence.py +++ b/dags/multipod/maxtext_convergence.py @@ -87,16 +87,20 @@ # The grain dataset takes longer to run, so we give it a longer timeout. The other tests are expected to complete within 5 hours. timeout_in_min = 360 if test_name == "maxtext-convergence-grain" else 300 - test_task = gke_config.get_gke_config( - cluster=XpkClusters.TPU_V6E_256_MLPERF_CLUSTER, - time_out_in_min=timeout_in_min, - test_name=test_name, - run_model_cmds=run_command, - docker_image=DockerImage.MAXTEXT_TPU_JAX_STABLE.value, - test_owner=test_owner.MATT_D, - base_output_directory=base_output_directory, - metric_aggregation_strategy=metric_config.AggregationStrategy.LAST, - ).run_with_run_name_generation() + test_task = ( + gke_config.get_gke_config( + cluster=XpkClusters.TPU_V6E_256_MLPERF_CLUSTER, + time_out_in_min=timeout_in_min, + test_name=test_name, + run_model_cmds=run_command, + docker_image=DockerImage.MAXTEXT_TPU_JAX_STABLE.value, + test_owner=test_owner.MATT_D, + base_output_directory=base_output_directory, + metric_aggregation_strategy=metric_config.AggregationStrategy.LAST, + ) + .to_name_gen_and_quarantine_task() + .run() + ) if test_name not in parallel_test_names: sequential_tests.append(test_task) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 8cb0e628a..be234d34d 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -796,8 +796,8 @@ def run_with_name_gen_and_quarantine( nested_run_name_in_tb_file_location: bool = True, ) -> DAGNode: test_name = self.task_test_config.benchmark_id - if QuarantineTests.is_quarantined(test_name): - with quarantine_task_group: + if QuarantineTests.is_quarantined(test_name) and self.quarantine_task_group: + with self.quarantine_task_group: return self.run_with_run_name_generation( use_pathways, xpk_branch, From 3c280234b506fbba71d4e2eaad62bfa7216d4375 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Wed, 8 Jul 2026 02:09:50 +0000 Subject: [PATCH 06/18] Fix format/linting errors and set the quarantine_task_group arg default to None --- ...text_aqtp_version_sweep_gke_example_dag.py | 8 ++- .../maxtext_profile_namegen_example_dag.py | 54 ++++++++++++------- .../maxtext_profile_sweep_example_dag.py | 4 +- .../examples/maxtext_sweep_gke_example_dag.py | 10 +++- dags/multipod/maxtext_convergence.py | 3 +- xlml/apis/task.py | 2 +- 6 files changed, 55 insertions(+), 26 deletions(-) diff --git a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py index 7d66275eb..495320715 100644 --- a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py +++ b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py @@ -21,7 +21,7 @@ import datetime from airflow import models from dags.common import test_owner -from dags.common.vm_resource import TpuVersion, Zone, Project, XpkClusters, DockerImage +from dags.common.vm_resource import XpkClusters, DockerImage from dags.multipod.configs import maxtext_sweep_gke_config # Set concurrency to number of workers otherwise tasks may time out @@ -54,7 +54,11 @@ for model_size in models: run_cmds = [ "pip show aqtp", - f"bash src/maxtext/configs/tpu/{tpu}/{model_size}.sh EXECUTABLE=train.py OUTPUT_PATH={base_output_directory} PLATFORM=gke", + ( + f"bash src/maxtext/configs/tpu/{tpu}/{model_size}.sh " + f"EXECUTABLE=train.py OUTPUT_PATH={base_output_directory} " + "PLATFORM=gke" + ), ] tests.extend( diff --git a/dags/examples/maxtext_profile_namegen_example_dag.py b/dags/examples/maxtext_profile_namegen_example_dag.py index efeec9a1a..dbc4c81e6 100644 --- a/dags/examples/maxtext_profile_namegen_example_dag.py +++ b/dags/examples/maxtext_profile_namegen_example_dag.py @@ -12,15 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -An example DAG to extract profile metrics from pretraining mixtral-8x7b model on 1xv4-128. +"""An example DAG to extract profile metrics. + +Pretraining mixtral-8x7b model on 1xv4-128. Profile extraction can be easily integrated with gke_config + (to_name_gen_and_quarantine_task + run). """ import datetime from airflow import models from dags.common import test_owner -from dags.common.vm_resource import XpkClusters, DockerImage +from dags.common.vm_resource import XpkClusters from dags.multipod.configs import gke_config from xlml.apis import metric_config @@ -28,36 +29,50 @@ BASE_OUTPUT_PATH = "gs://runner-maxtext-logs" docker_image = { - "stable": "gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:2025-05-20", + "stable": ( + "gcr.io/tpu-prod-env-multipod/" "maxtext_jax_stable_stack:2025-05-20" + ), } base_command = ( f"export BASE_OUTPUT_PATH={BASE_OUTPUT_PATH} && " - + "python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml base_output_directory=gs://runner-maxtext-logs run_name=${RUN_NAME} model_name=mixtral-8x7b tokenizer_path=assets/tokenizer.mistral-v1 dataset_path=gs://maxtext-dataset per_device_batch_size=4 enable_checkpointing=false ici_fsdp_parallelism=-1 max_target_length=1024 async_checkpointing=false attention=flash dtype=bfloat16 weight_dtype=bfloat16" + "python3 -m maxtext.trainers.pre_train.train " + "src/maxtext/configs/base.yml " + "base_output_directory=gs://runner-maxtext-logs " + "run_name=${RUN_NAME} model_name=mixtral-8x7b " + "tokenizer_path=assets/tokenizer.mistral-v1 " + "dataset_path=gs://maxtext-dataset per_device_batch_size=4 " + "enable_checkpointing=false ici_fsdp_parallelism=-1 " + "max_target_length=1024 async_checkpointing=false " + "attention=flash dtype=bfloat16 weight_dtype=bfloat16" ) test_models_tpu = { # use: upload single profile from the first host, extract profile - # add profiler config: ensure steps > skip_first_n_steps_for_profiler + profiler_steps + # add profiler config: ensure steps > + # skip_first_n_steps_for_profiler + profiler_steps "mixtral-8x7b_pretraining-megablox_config-true_upload-one": { "cluster": XpkClusters.TPU_V4_128_CLUSTER, "time_out_in_min": 60, "train_command": [ - base_command - + " steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3", + base_command + " steps=10 profiler=xplane " + "skip_first_n_steps_for_profiler=5 profiler_steps=3", ], }, # use: upload profiles from all hosts, extract one of the profiles - # add profiler config: ensure steps > skip_first_n_steps_for_profiler + profiler_steps + # add profiler config: ensure steps > + # skip_first_n_steps_for_profiler + profiler_steps "mixtral-8x7b_pretraining-megablox_config-true_upload-all": { "cluster": XpkClusters.TPU_V4_128_CLUSTER, "time_out_in_min": 60, "train_command": [ - base_command - + " steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3 upload_all_profiler_results=True", + base_command + " steps=10 profiler=xplane " + "skip_first_n_steps_for_profiler=5 profiler_steps=3 " + "upload_all_profiler_results=True", ], }, - # testing: handle edge case, attempt to extract, find no match, proceed to post_process without error + # testing: handle edge case, attempt to extract, find no match, + # proceed to post_process without error "testing_config-true_upload-none": { "cluster": XpkClusters.TPU_V4_128_CLUSTER, "time_out_in_min": 60, @@ -65,13 +80,14 @@ base_command + " steps=10", ], }, - # testing: not generate profile location, not extract profile in post_process + # testing: not generate profile location, not extract profile in + # post_process "testing_config-false_upload-one": { "cluster": XpkClusters.TPU_V4_128_CLUSTER, "time_out_in_min": 60, "train_command": [ - base_command - + " steps=10 profiler=xplane skip_first_n_steps_for_profiler=5 profiler_steps=3", + base_command + " steps=10 profiler=xplane " + "skip_first_n_steps_for_profiler=5 profiler_steps=3", ], "not_add_profile_config": True, }, @@ -87,7 +103,7 @@ concurrency=2, ) as dag: for run_name, test_scripts_details in test_models_tpu.items(): - for image in docker_image.keys(): + for image, img_val in docker_image.items(): # file_location: pass in base_output_directory, will be altered in # XpkNameGenAndQuarantineTask.run_with_run_name_generation job_metric_config = metric_config.MetricConfig() @@ -112,10 +128,12 @@ time_out_in_min=test_scripts_details["time_out_in_min"], test_name=f"maxtext_{image}_{run_name}", run_model_cmds=test_scripts_details["train_command"], - docker_image=docker_image[image], + docker_image=img_val, test_owner=test_owner.SHUNING_J, cluster=test_scripts_details["cluster"], - user_specified_job_metric_config=job_metric_config, # customize config + user_specified_job_metric_config=( + job_metric_config + ), # customize config ) .to_name_gen_and_quarantine_task(run_name_env="RUN_NAME") .run() diff --git a/dags/examples/maxtext_profile_sweep_example_dag.py b/dags/examples/maxtext_profile_sweep_example_dag.py index 612531691..3b0257852 100644 --- a/dags/examples/maxtext_profile_sweep_example_dag.py +++ b/dags/examples/maxtext_profile_sweep_example_dag.py @@ -24,7 +24,7 @@ from airflow.utils.task_group import TaskGroup from dags.common import test_owner from dags.common.vm_resource import XpkClusters, Project -from dags.multipod.configs import maxtext_sweep_gke_config +from dags.multipod.configs import maxtext_sweep_gke_config as sweep_config from xlml.apis import metric_config SCHEDULED_TIME = None @@ -139,7 +139,7 @@ def dict_to_arg(param_dict): # generate run_name and tensorboard/profile location for extraction num_slices = [1] sweep_params = {} - maxtext_sweep_gke_test = maxtext_sweep_gke_config.get_maxtext_sweep_gke_config( + maxtext_sweep_gke_test = sweep_config.get_maxtext_sweep_gke_config( test_owner=test_owner.SHUNING_J, dataset_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project=Project.CLOUD_ML_AUTO_SOLUTIONS.value, diff --git a/dags/examples/maxtext_sweep_gke_example_dag.py b/dags/examples/maxtext_sweep_gke_example_dag.py index 464a26c39..7437bc7f8 100644 --- a/dags/examples/maxtext_sweep_gke_example_dag.py +++ b/dags/examples/maxtext_sweep_gke_example_dag.py @@ -21,7 +21,7 @@ import datetime from airflow import models from dags.common import test_owner -from dags.common.vm_resource import TpuVersion, Zone, Project, XpkClusters, DockerImage +from dags.common.vm_resource import XpkClusters, DockerImage from dags.multipod.configs import maxtext_sweep_gke_config # Set concurrency to number of workers otherwise tasks may time out @@ -37,7 +37,13 @@ # MaxText set up and run commands base_output_directory = "gs://maxtext-experiments-multipod" base_run_model_cmds = [ - f"python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml base_output_directory={base_output_directory} dataset_path=gs://max-datasets-rogue enable_checkpointing=false global_parameter_scale=16 steps=10", + ( + "python3 -m maxtext.trainers.pre_train.train " + "src/maxtext/configs/base.yml " + f"base_output_directory={base_output_directory} " + "dataset_path=gs://max-datasets-rogue enable_checkpointing=false " + "global_parameter_scale=16 steps=10" + ), ] # Get list of MaxText GKE XPK jobs diff --git a/dags/multipod/maxtext_convergence.py b/dags/multipod/maxtext_convergence.py index f5e476eaa..681ffaa79 100644 --- a/dags/multipod/maxtext_convergence.py +++ b/dags/multipod/maxtext_convergence.py @@ -84,7 +84,8 @@ sequential_tests = [] for test_name, run_command in convergence_tests.items(): - # The grain dataset takes longer to run, so we give it a longer timeout. The other tests are expected to complete within 5 hours. + # The grain dataset takes longer to run, so we give it a longer timeout. + # The other tests are expected to complete within 5 hours. timeout_in_min = 360 if test_name == "maxtext-convergence-grain" else 300 test_task = ( diff --git a/xlml/apis/task.py b/xlml/apis/task.py index be234d34d..8e21e7833 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -399,7 +399,7 @@ def run_with_node_interruption( def to_name_gen_and_quarantine_task( self, - quarantine_task_group, + quarantine_task_group: Any = None, run_name_env: str = "M_RUN_NAME", nested_run_name_in_tb_file_location: bool = True, ) -> "XpkNameGenAndQuarantineTask": From aa4cfb97259e6a510d698c5af6af18337a1b796c Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Wed, 8 Jul 2026 08:34:01 +0000 Subject: [PATCH 07/18] Reconstruct the XpkTask --- ...text_aqtp_version_sweep_gke_example_dag.py | 2 +- .../maxtext_profile_namegen_example_dag.py | 29 ++-- .../maxtext_profile_sweep_example_dag.py | 3 +- .../examples/maxtext_sweep_gke_example_dag.py | 2 +- dags/multipod/configs/gke_config.py | 130 +++++++++++++++++- .../configs/maxtext_sweep_gke_config.py | 6 +- dags/multipod/maxtext_convergence.py | 24 ++-- dags/orbax/maxtext_emc_restore_gcs.py | 35 ++--- dags/orbax/maxtext_emc_restore_local.py | 33 ++--- dags/orbax/maxtext_mtc_restore_local.py | 33 ++--- ...xt_reg_restore_gcs_with_node_disruption.py | 35 ++--- xlml/apis/task.py | 70 +--------- 12 files changed, 220 insertions(+), 182 deletions(-) diff --git a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py index 495320715..aa484a44c 100644 --- a/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py +++ b/dags/examples/maxtext_aqtp_version_sweep_gke_example_dag.py @@ -102,4 +102,4 @@ # Run jobs for test in tests: - test.to_name_gen_and_quarantine_task().run() + test.run() diff --git a/dags/examples/maxtext_profile_namegen_example_dag.py b/dags/examples/maxtext_profile_namegen_example_dag.py index dbc4c81e6..115a4fd17 100644 --- a/dags/examples/maxtext_profile_namegen_example_dag.py +++ b/dags/examples/maxtext_profile_namegen_example_dag.py @@ -122,19 +122,16 @@ if "not_add_profile_config" in test_scripts_details: job_metric_config.profile = None - tpu_task = ( - gke_config.get_gke_config( - num_slices=1, - time_out_in_min=test_scripts_details["time_out_in_min"], - test_name=f"maxtext_{image}_{run_name}", - run_model_cmds=test_scripts_details["train_command"], - docker_image=img_val, - test_owner=test_owner.SHUNING_J, - cluster=test_scripts_details["cluster"], - user_specified_job_metric_config=( - job_metric_config - ), # customize config - ) - .to_name_gen_and_quarantine_task(run_name_env="RUN_NAME") - .run() - ) + tpu_task = gke_config.get_gke_config_with_name_gen_and_quarantine( + num_slices=1, + time_out_in_min=test_scripts_details["time_out_in_min"], + test_name=f"maxtext_{image}_{run_name}", + run_model_cmds=test_scripts_details["train_command"], + docker_image=img_val, + test_owner=test_owner.SHUNING_J, + cluster=test_scripts_details["cluster"], + user_specified_job_metric_config=( + job_metric_config + ), # customize config + run_name_env="RUN_NAME", + ).run() diff --git a/dags/examples/maxtext_profile_sweep_example_dag.py b/dags/examples/maxtext_profile_sweep_example_dag.py index 3b0257852..28e5ff05b 100644 --- a/dags/examples/maxtext_profile_sweep_example_dag.py +++ b/dags/examples/maxtext_profile_sweep_example_dag.py @@ -153,7 +153,8 @@ def dict_to_arg(param_dict): base_run_model_cmds=test_scripts_details["train_command"], sweep_params=sweep_params, enable_profile_config=True, # add flag to enable profile extraction + quarantine_task_group=quarantine_task_group, ) for test in maxtext_sweep_gke_test: - test.to_name_gen_and_quarantine_task(quarantine_task_group).run() + test.run() diff --git a/dags/examples/maxtext_sweep_gke_example_dag.py b/dags/examples/maxtext_sweep_gke_example_dag.py index 7437bc7f8..2d11d9e79 100644 --- a/dags/examples/maxtext_sweep_gke_example_dag.py +++ b/dags/examples/maxtext_sweep_gke_example_dag.py @@ -63,4 +63,4 @@ # Run jobs for test in maxtext_sweep_gke_test: - test.to_name_gen_and_quarantine_task().run() + test.run() diff --git a/dags/multipod/configs/gke_config.py b/dags/multipod/configs/gke_config.py index da705b630..bfe1db888 100644 --- a/dags/multipod/configs/gke_config.py +++ b/dags/multipod/configs/gke_config.py @@ -19,7 +19,7 @@ from xlml.apis.xpk_cluster_config import XpkClusterConfig from dags import gcs_bucket from dags.common.vm_resource import TpuVersion, Project, XpkClusters, GpuVersion, CpuVersion -from typing import Iterable +from typing import Any, Iterable import datetime @@ -81,6 +81,134 @@ def get_gke_config( ) +def get_gke_config_with_interrupt( + time_out_in_min: int, + test_name: str, + docker_image: str, + test_owner: str, + run_model_cmds: Iterable[str], + expect_reach_to_step: int, + cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, + num_slices: int = 1, + dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + base_output_directory: str = None, + metric_aggregation_strategy: metric_config.AggregationStrategy = None, + user_specified_job_metric_config: metric_config.MetricConfig = None, + last_node: bool = False, + check_file_exists: bool = False, +) -> task.XpkNodeInterruptionTask: + job_gcp_config = gcp_config.GCPConfig( + project_name=cluster.project, + zone=cluster.zone, + dataset_name=dataset_name, + dataset_project=dataset_project, + composer_project=composer_project, + ) + + job_test_config = test_config.TpuGkeTest( + test_config.Tpu( + version=cluster.device_version, + cores=cluster.core_count, + ), + test_name=test_name, + run_model_cmds=run_model_cmds, + set_up_cmds=None, + timeout=datetime.timedelta(minutes=time_out_in_min), + task_owner=test_owner, + num_slices=num_slices, + cluster_name=cluster.name, + docker_image=docker_image, + ) + job_metric_config = user_specified_job_metric_config + if job_metric_config is None: + job_metric_config = ( + metric_config.MetricConfig( + tensorboard_summary=metric_config.SummaryConfig( + file_location=base_output_directory, + aggregation_strategy=metric_aggregation_strategy, + use_regex_file_location=True, + ), + ) + if base_output_directory and metric_aggregation_strategy + else None + ) + + return task.XpkNodeInterruptionTask( + task_test_config=job_test_config, + task_gcp_config=job_gcp_config, + task_metric_config=job_metric_config, + expect_reach_to_step=expect_reach_to_step, + last_node=last_node, + check_file_exists=check_file_exists, + ) + + +def get_gke_config_with_name_gen_and_quarantine( + time_out_in_min: int, + test_name: str, + docker_image: str, + test_owner: str, + run_model_cmds: Iterable[str], + cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, + num_slices: int = 1, + dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + base_output_directory: str = None, + metric_aggregation_strategy: metric_config.AggregationStrategy = None, + user_specified_job_metric_config: metric_config.MetricConfig = None, + quarantine_task_group: Any = None, + run_name_env: str = "M_RUN_NAME", + nested_run_name_in_tb_file_location: bool = True, +) -> task.XpkNameGenAndQuarantineTask: + job_gcp_config = gcp_config.GCPConfig( + project_name=cluster.project, + zone=cluster.zone, + dataset_name=dataset_name, + dataset_project=dataset_project, + composer_project=composer_project, + ) + + job_test_config = test_config.TpuGkeTest( + test_config.Tpu( + version=cluster.device_version, + cores=cluster.core_count, + ), + test_name=test_name, + run_model_cmds=run_model_cmds, + set_up_cmds=None, + timeout=datetime.timedelta(minutes=time_out_in_min), + task_owner=test_owner, + num_slices=num_slices, + cluster_name=cluster.name, + docker_image=docker_image, + ) + job_metric_config = user_specified_job_metric_config + if job_metric_config is None: + job_metric_config = ( + metric_config.MetricConfig( + tensorboard_summary=metric_config.SummaryConfig( + file_location=base_output_directory, + aggregation_strategy=metric_aggregation_strategy, + use_regex_file_location=True, + ), + ) + if base_output_directory and metric_aggregation_strategy + else None + ) + + return task.XpkNameGenAndQuarantineTask( + task_test_config=job_test_config, + task_gcp_config=job_gcp_config, + task_metric_config=job_metric_config, + quarantine_task_group=quarantine_task_group, + run_name_env=run_name_env, + nested_run_name_in_tb_file_location=nested_run_name_in_tb_file_location, + ) + + def get_gke_maxtext_nightly_config( time_out_in_min: int, test_name: str, diff --git a/dags/multipod/configs/maxtext_sweep_gke_config.py b/dags/multipod/configs/maxtext_sweep_gke_config.py index 8cdeb69a0..582593826 100644 --- a/dags/multipod/configs/maxtext_sweep_gke_config.py +++ b/dags/multipod/configs/maxtext_sweep_gke_config.py @@ -54,7 +54,8 @@ def get_maxtext_sweep_gke_config( dataset_project: str = None, composer_project: str = None, enable_profile_config: bool = False, -) -> List[task.XpkTask]: + quarantine_task_group: Any = None, +) -> List[task.XpkNameGenAndQuarantineTask]: if not dataset_project: dataset_project = cluster.project if not composer_project: @@ -122,10 +123,11 @@ def get_maxtext_sweep_gke_config( file_location=base_output_directory, ) - xpk_task = task.XpkTask( + xpk_task = task.XpkNameGenAndQuarantineTask( task_test_config=job_test_config, task_gcp_config=job_gcp_config, task_metric_config=job_metric_config, + quarantine_task_group=quarantine_task_group, ) xpk_task_list.append(xpk_task) diff --git a/dags/multipod/maxtext_convergence.py b/dags/multipod/maxtext_convergence.py index 681ffaa79..5340bda6a 100644 --- a/dags/multipod/maxtext_convergence.py +++ b/dags/multipod/maxtext_convergence.py @@ -88,20 +88,16 @@ # The other tests are expected to complete within 5 hours. timeout_in_min = 360 if test_name == "maxtext-convergence-grain" else 300 - test_task = ( - gke_config.get_gke_config( - cluster=XpkClusters.TPU_V6E_256_MLPERF_CLUSTER, - time_out_in_min=timeout_in_min, - test_name=test_name, - run_model_cmds=run_command, - docker_image=DockerImage.MAXTEXT_TPU_JAX_STABLE.value, - test_owner=test_owner.MATT_D, - base_output_directory=base_output_directory, - metric_aggregation_strategy=metric_config.AggregationStrategy.LAST, - ) - .to_name_gen_and_quarantine_task() - .run() - ) + test_task = gke_config.get_gke_config_with_name_gen_and_quarantine( + cluster=XpkClusters.TPU_V6E_256_MLPERF_CLUSTER, + time_out_in_min=timeout_in_min, + test_name=test_name, + run_model_cmds=run_command, + docker_image=DockerImage.MAXTEXT_TPU_JAX_STABLE.value, + test_owner=test_owner.MATT_D, + base_output_directory=base_output_directory, + metric_aggregation_strategy=metric_config.AggregationStrategy.LAST, + ).run() if test_name not in parallel_test_names: sequential_tests.append(test_task) diff --git a/dags/orbax/maxtext_emc_restore_gcs.py b/dags/orbax/maxtext_emc_restore_gcs.py index 80eb9d308..1edaadcca 100644 --- a/dags/orbax/maxtext_emc_restore_gcs.py +++ b/dags/orbax/maxtext_emc_restore_gcs.py @@ -120,26 +120,21 @@ task_id="generate_start_time" )() - 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}-emc", - run_model_cmds=workload_command, - docker_image=image.value, - test_owner=test_owner.DEPP_L, - ) - .to_node_interruption_task( - expect_reach_to_step=step_to_interrupt, - last_node=True, - ) - .run( - ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, - mtc_enabled=True, - skip_post_process=True, - max_restart=15, - ) + maxtext_chkpt_run_test = gke_config.get_gke_config_with_interrupt( + num_slices=slice_num, + cluster=test_config.cluster, + time_out_in_min=60, + test_name=f"{test_config.short_id}-emc", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.DEPP_L, + expect_reach_to_step=step_to_interrupt, + last_node=True, + ).run( + ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, + mtc_enabled=True, + skip_post_process=True, + max_restart=15, ) end_time = validation_util.generate_timestamp.override( diff --git a/dags/orbax/maxtext_emc_restore_local.py b/dags/orbax/maxtext_emc_restore_local.py index 4c75b32dc..a1372fac9 100644 --- a/dags/orbax/maxtext_emc_restore_local.py +++ b/dags/orbax/maxtext_emc_restore_local.py @@ -118,25 +118,20 @@ task_id="generate_start_time" )() - 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}-emc", - run_model_cmds=workload_command, - docker_image=image.value, - test_owner=test_owner.DEPP_L, - ) - .to_node_interruption_task( - expect_reach_to_step=step_to_interrupt, - ) - .run( - ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, - mtc_enabled=True, - skip_post_process=True, - max_restart=15, - ) + maxtext_chkpt_run_test = gke_config.get_gke_config_with_interrupt( + num_slices=slice_num, + cluster=test_config.cluster, + time_out_in_min=60, + test_name=f"{test_config.short_id}-emc", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.DEPP_L, + expect_reach_to_step=step_to_interrupt, + ).run( + ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, + mtc_enabled=True, + skip_post_process=True, + max_restart=15, ) end_time = validation_util.generate_timestamp.override( diff --git a/dags/orbax/maxtext_mtc_restore_local.py b/dags/orbax/maxtext_mtc_restore_local.py index 9130625a7..5548b6f80 100644 --- a/dags/orbax/maxtext_mtc_restore_local.py +++ b/dags/orbax/maxtext_mtc_restore_local.py @@ -121,25 +121,20 @@ task_id="generate_start_time" )() - 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}-mtc", - run_model_cmds=workload_command, - docker_image=image.value, - test_owner=test_owner.DEPP_L, - ) - .to_node_interruption_task( - expect_reach_to_step=step_to_interrupt, - ) - .run( - ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, - mtc_enabled=True, - skip_post_process=True, - max_restart=15, - ) + maxtext_chkpt_run_test = gke_config.get_gke_config_with_interrupt( + num_slices=slice_num, + cluster=test_config.cluster, + time_out_in_min=60, + test_name=f"{test_config.short_id}-mtc", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.DEPP_L, + expect_reach_to_step=step_to_interrupt, + ).run( + ramdisk_directory=test_config_util.DEFAULT_RAM_DISK, + mtc_enabled=True, + skip_post_process=True, + max_restart=15, ) end_time = validation_util.generate_timestamp.override( diff --git a/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py b/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py index e30731f47..5e9ef528d 100644 --- a/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py +++ b/dags/orbax/maxtext_reg_restore_gcs_with_node_disruption.py @@ -130,26 +130,21 @@ def generate_workload_checkpoints_location(gcs_ckpt_location: str) -> str: task_id="generate_start_time" )() - 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}", - run_model_cmds=workload_command, - docker_image=image.value, - test_owner=test_owner.SHARON_Y, - ) - .to_node_interruption_task( - expect_reach_to_step=step_to_interrupt, - check_file_exists=True, - ) - .run( - gcs_location=gcs_location, - xpk_branch=MAIN_BRANCH, - skip_post_process=True, - max_restart=15, - ) + maxtext_chkpt_run_test = gke_config.get_gke_config_with_interrupt( + num_slices=slice_num, + cluster=test_config.cluster, + time_out_in_min=60, + test_name=f"{test_config.short_id}", + run_model_cmds=workload_command, + docker_image=image.value, + test_owner=test_owner.SHARON_Y, + expect_reach_to_step=step_to_interrupt, + check_file_exists=True, + ).run( + gcs_location=gcs_location, + xpk_branch=MAIN_BRANCH, + skip_post_process=True, + max_restart=15, ) end_time = validation_util.generate_timestamp.override( diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 8e21e7833..7e9b56db7 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -381,38 +381,6 @@ def run( return group - def run_with_node_interruption( - self, - expect_reach_to_step: int, - last_node: bool = False, - check_file_exists: bool = False, - ) -> "XpkNodeInterruptionTask": - return XpkNodeInterruptionTask( - task_test_config=self.task_test_config, - task_gcp_config=self.task_gcp_config, - task_metric_config=self.task_metric_config, - workload_provision_timeout=self.workload_provision_timeout, - expect_reach_to_step=expect_reach_to_step, - last_node=last_node, - check_file_exists=check_file_exists, - ) - - def to_name_gen_and_quarantine_task( - self, - quarantine_task_group: Any = None, - run_name_env: str = "M_RUN_NAME", - nested_run_name_in_tb_file_location: bool = True, - ) -> "XpkNameGenAndQuarantineTask": - return XpkNameGenAndQuarantineTask( - task_test_config=self.task_test_config, - task_gcp_config=self.task_gcp_config, - task_metric_config=self.task_metric_config, - workload_provision_timeout=self.workload_provision_timeout, - quarantine_task_group=quarantine_task_group, - run_name_env=run_name_env, - nested_run_name_in_tb_file_location=nested_run_name_in_tb_file_location, - ) - def run_model( self, gcs_location: Optional[airflow.XComArg] = None, @@ -582,50 +550,16 @@ class XpkNodeInterruptionTask(XpkTask): last_node: bool = False check_file_exists: bool = False - def run( - 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, - max_restart: int = 0, - ) -> DAGNode: - """Run a test job with node interruption.""" - with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: - run_model, gcs_path = self.run_model_with_node_interruption( - gcs_location=gcs_location, - use_vertex_tensorboard=use_vertex_tensorboard, - expect_reach_to_step=expect_reach_to_step, - use_pathways=use_pathways, - ramdisk_directory=ramdisk_directory, - mtc_enabled=mtc_enabled, - xpk_branch=xpk_branch, - last_node=last_node, - max_restart=max_restart, - check_file_exists=check_file_exists, - ) - if not skip_post_process: - _ = run_model >> self.post_process(gcs_path) - return group - - def run_model_with_node_interruption( + def run_model( self, - *, gcs_location: Optional[airflow.XComArg] = None, use_vertex_tensorboard: bool = False, - expect_reach_to_step: int, use_pathways: bool = False, ramdisk_directory: str = "", mtc_enabled: bool = False, xpk_branch: str = xpk.MAIN_BRANCH, - last_node: bool = False, max_restart: int = 0, - check_file_exists: bool = False, - ) -> DAGNode: + ) -> Tuple[DAGNode, str]: """Run the TPU/GPU test in `task_test_config` using xpk. with interruption. From fded7822c84cffe8fa96a489f5e4cb1fd741c687 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Fri, 10 Jul 2026 02:41:59 +0000 Subject: [PATCH 08/18] Reuse and override functions --- xlml/apis/task.py | 470 ++++++++++++++-------------------------------- 1 file changed, 144 insertions(+), 326 deletions(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 7e9b56db7..5ba3ad814 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -385,7 +385,6 @@ def run_model( self, gcs_location: Optional[airflow.XComArg] = None, use_vertex_tensorboard: bool = False, - expect_reach_to_step: int, use_pathways: bool = False, skip_post_process: bool = False, ramdisk_directory: str = "", @@ -429,6 +428,11 @@ def run_model( self.task_test_config.gcs_subfolder, self.task_test_config.benchmark_id, ) + + dummy_op_for_teardown = EmptyOperator( + task_id="dummy_op_for_teardown" + ).as_setup() + launch_workload = self.launch_workload( workload_id, gcs_path, @@ -456,9 +460,13 @@ def run_model( xpk_branch=xpk_branch, ) + intermediary_flow = self.intermediary_flow( + launch_workload, workload_id, gcs_path, xpk_branch + ) + _ = ( (workload_id, gcs_path) - >> launch_workload + >> intermediary_flow >> wait_for_workload_completion >> clean_up_workload ) @@ -541,6 +549,15 @@ def post_process(self, result_location: Optional[str] = None) -> DAGNode: return group + def intermediary_flow( + self, + launch_workload: DAGNode, + workload_id: str, + gcs_path: str, + xpk_branch: str, + ) -> DAGNode: + return launch_workload + @dataclasses.dataclass class XpkNodeInterruptionTask(XpkTask): @@ -550,131 +567,53 @@ class XpkNodeInterruptionTask(XpkTask): last_node: bool = False check_file_exists: bool = False - def run_model( + def run( 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, max_restart: int = 0, - ) -> Tuple[DAGNode, str]: - """Run the TPU/GPU test in `task_test_config` using xpk. - - with interruption. - """ - 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, - ) - - dummy_op_for_teardown = EmptyOperator( - task_id="dummy_op_for_teardown" - ).as_setup() - - launch_workload_and_wait_for_reach_step = ( - self.launch_workload_with_node_reach_to_step( - workload_id, - gcs_path, - expect_reach_to_step, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - max_restart, - check_file_exists, - ) - ) - - run_node_interruption = xpk.delete_node.override( - owner=self.task_test_config.task_owner, trigger_rule="none_failed" - )( - project=self.task_gcp_config.project_name, - zone=self.task_gcp_config.zone, - cluster_name=self.task_test_config.cluster_name, - workload_id=workload_id, - dry_run=False, - last_node=last_node, - ) - - 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=gke.zone_to_region(self.task_gcp_config.zone), - 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, - xpk_branch=xpk_branch, - ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True) + ) -> DAGNode: + """Run a test job with node interruption.""" + return super().run( + gcs_location=gcs_location, + use_vertex_tensorboard=use_vertex_tensorboard, + use_pathways=use_pathways, + skip_post_process=skip_post_process, + ramdisk_directory=ramdisk_directory, + mtc_enabled=mtc_enabled, + xpk_branch=xpk_branch, + max_restart=max_restart, + ) - _ = ( - (workload_id, gcs_path) - >> dummy_op_for_teardown - >> launch_workload_and_wait_for_reach_step - >> run_node_interruption - >> wait_for_workload_completion - >> clean_up_workload - ) - return group, gcs_path + def intermediary_flow( + self, + launch_workload: DAGNode, + workload_id: str, + gcs_path: str, + xpk_branch: str, + ) -> DAGNode: + """Inject interruption steps after launch_workload.""" + return self.wait_for_workload_reach_step_and_interruption( + launch_workload, workload_id, gcs_path, xpk_branch + ) - def launch_workload_with_node_reach_to_step( + def wait_for_workload_reach_step_and_interruption( self, + launch_workload: DAGNode, workload_id: str, gcs_path: str, - expect_reach_to_step: int, - use_vertex_tensorboard: bool, - use_pathways: bool = False, - ramdisk_directory: str = "", - mtc_enabled: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - max_restart: int = 0, - check_file_exists: bool = False, + xpk_branch: str, ) -> DAGNode: - """Create the workload and wait for it to provision.""" - with TaskGroup(group_id="launch_workload_with_node_reach_to_step") 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, - max_restart=max_restart, - ) - 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=gke.zone_to_region(self.task_gcp_config.zone), - cluster_name=self.task_test_config.cluster_name, - ) + """Wait for workload to reach specific step and trigger node deletion.""" + with TaskGroup( + group_id="wait_for_workload_reach_step_and_interruption" + ) as group: wait_for_workload_to_reach_step = ( xpk.wait_for_workload_reach_step.override( task_id="wait_for_workload_reach_step" @@ -683,7 +622,7 @@ def launch_workload_with_node_reach_to_step( project_id=self.task_gcp_config.project_name, region=gke.zone_to_region(self.task_gcp_config.zone), cluster_name=self.task_test_config.cluster_name, - expect_reach_to_step=str(expect_reach_to_step), + expect_reach_to_step=str(self.expect_reach_to_step), ) ) @@ -692,7 +631,7 @@ def launch_workload_with_node_reach_to_step( task_id=task_id_wait_file_exist )( file_path=( - f"{gcs_path}/{str(expect_reach_to_step)}/commit_success.txt" + f"{gcs_path}/{str(self.expect_reach_to_step)}/commit_success.txt" ), ) task_id_do_nothing = "do_nothing" @@ -709,254 +648,133 @@ def task_path_decider(check_file_exists: bool = False) -> str: # Conditional checks: depending on the `check_file_exists` argument # specified by the upper-level caller. - maybe_check_file_exists = task_path_decider(check_file_exists) + maybe_check_file_exists = task_path_decider(self.check_file_exists) + + run_node_interruption = xpk.delete_node.override( + owner=self.task_test_config.task_owner, trigger_rule="none_failed" + )( + project=self.task_gcp_config.project_name, + zone=self.task_gcp_config.zone, + cluster_name=self.task_test_config.cluster_name, + workload_id=workload_id, + dry_run=False, + last_node=self.last_node, + ) _ = ( - run_workload - >> wait_for_workload_start + launch_workload >> wait_for_workload_to_reach_step >> maybe_check_file_exists ) _ = maybe_check_file_exists >> [wait_for_file_to_exist, do_nothing] + _ = [wait_for_file_to_exist, do_nothing] >> run_node_interruption - return group - - def run_with_name_gen_and_quarantine( - self, - quarantine_task_group, - use_pathways: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - run_name_env: str = "M_RUN_NAME", - nested_run_name_in_tb_file_location: bool = True, - ) -> DAGNode: - test_name = self.task_test_config.benchmark_id - if QuarantineTests.is_quarantined(test_name) and self.quarantine_task_group: - with self.quarantine_task_group: - return self.run_with_run_name_generation( - use_pathways, - xpk_branch, - run_name_env, - nested_run_name_in_tb_file_location, - ) - else: - return self.run_with_run_name_generation( - use_pathways, - xpk_branch, - run_name_env, - nested_run_name_in_tb_file_location, - ) - - def run_with_run_name_generation( - self, - use_pathways: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - run_name_env: str = "M_RUN_NAME", - nested_run_name_in_tb_file_location: bool = True, - ) -> DAGNode: - """Generate a unique run name, tensorboard file location, - and profile file location (if metric config has profile), - then run a test job within a docker image. - - Returns: - A task group with the following tasks chained: generate_run_name, - generate_tb_file_location, generate_profile_file_location (optional), - run provision, run_model, post_process. - """ - with TaskGroup( - group_id=self.task_test_config.benchmark_id, prefix_group_id=True - ) as group: - run_name = name_format.generate_run_name( - self.task_test_config.benchmark_id - ) - tb_file_location = name_format.generate_tb_file_location( - run_name, - self.task_metric_config.tensorboard_summary.file_location, - nested_run_name_in_tb_file_location, - ) + return group - # Set run_name in run_model_cmds - new_run_model_cmds = [f"export {run_name_env}={run_name}"] - for cmd in self.task_test_config.run_model_cmds: - new_run_model_cmds.append(cmd) - self.task_test_config.run_model_cmds = new_run_model_cmds +@dataclasses.dataclass +class XpkNameGenAndQuarantineTask(XpkTask): + """Task for running XPK workloads with name generation and quarantine.""" - # Update tensorboard file location - self.task_metric_config.tensorboard_summary.file_location = ( - tb_file_location - ) + quarantine_task_group: Any = None + run_name_env: str = "M_RUN_NAME" + nested_run_name_in_tb_file_location: bool = True - # Update profile file location - if self.task_metric_config.profile: - profile_file_location = name_format.generate_profile_file_location( - run_name, self.task_metric_config.profile.file_location - ) - self.task_metric_config.profile.file_location = profile_file_location - run_model, gcs_path = self.run_model( - use_pathways=use_pathways, xpk_branch=xpk_branch - ) - flow = ( - run_name >> (tb_file_location, profile_file_location) >> run_model - ) - if not skip_post_process: - _ = flow >> self.post_process(gcs_path) - else: - run_model, gcs_path = self.run_model( - use_pathways=use_pathways, xpk_branch=xpk_branch - ) - _ = ( - run_name - >> tb_file_location - >> run_model - >> self.post_process(gcs_path) - ) - return group + # Temporary variables to pass tasks from run() to intermediary_flow() + _run_name_task: Any = dataclasses.field(default=None, init=False, repr=False) + _tb_file_location_task: Any = dataclasses.field(default=None, init=False, repr=False) + _profile_file_location_task: Any = dataclasses.field(default=None, init=False, repr=False) - def run_model( + def run( 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, max_restart: int = 0, ) -> DAGNode: - """Run the TPU/GPU test in `task_test_config` using xpk. - - 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, - ) + """Run a test job with name generation and quarantine.""" + test_name = self.task_test_config.benchmark_id - dummy_op_for_teardown = EmptyOperator( - task_id="dummy_op_for_teardown" - ).as_setup() + # 1. Generate run name and file location tasks first (they are Airflow task objects) + self._run_name_task = name_format.generate_run_name( + self.task_test_config.benchmark_id + ) + self._tb_file_location_task = name_format.generate_tb_file_location( + self._run_name_task, + self.task_metric_config.tensorboard_summary.file_location, + self.nested_run_name_in_tb_file_location, + ) - launch_workload = self.launch_workload( - workload_id, - gcs_path, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - max_restart, - ) - 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=gke.zone_to_region(self.task_gcp_config.zone), - cluster_name=self.task_test_config.cluster_name, - ) + # Set run_name in run_model_cmds + new_run_model_cmds = [f"export {self.run_name_env}={self._run_name_task}"] + for cmd in self.task_test_config.run_model_cmds: + new_run_model_cmds.append(cmd) + self.task_test_config.run_model_cmds = new_run_model_cmds - 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, - xpk_branch=xpk_branch, - ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True) + # Update tensorboard file location + self.task_metric_config.tensorboard_summary.file_location = ( + self._tb_file_location_task + ) - _ = ( - (workload_id, gcs_path) - >> dummy_op_for_teardown - >> launch_workload - >> wait_for_workload_completion - >> clean_up_workload + # Update profile file location if profile config exists + if self.task_metric_config.profile: + self._profile_file_location_task = name_format.generate_profile_file_location( + self._run_name_task, self.task_metric_config.profile.file_location ) - return group, gcs_path + self.task_metric_config.profile.file_location = self._profile_file_location_task - def launch_workload( - 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, - max_restart: int = 0, - ) -> DAGNode: - """Create the workload and wait for it to provision.""" - with TaskGroup(group_id="launch_workload") 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, + # 2. Run under quarantine task group if applicable + if QuarantineTests.is_quarantined(test_name) and self.quarantine_task_group: + with self.quarantine_task_group: + return super().run( + gcs_location=gcs_location, + use_vertex_tensorboard=use_vertex_tensorboard, + use_pathways=use_pathways, + skip_post_process=skip_post_process, + ramdisk_directory=ramdisk_directory, + mtc_enabled=mtc_enabled, + xpk_branch=xpk_branch, + max_restart=max_restart, + ) + else: + return super().run( + gcs_location=gcs_location, use_vertex_tensorboard=use_vertex_tensorboard, use_pathways=use_pathways, + skip_post_process=skip_post_process, ramdisk_directory=ramdisk_directory, mtc_enabled=mtc_enabled, xpk_branch=xpk_branch, max_restart=max_restart, ) - 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=gke.zone_to_region(self.task_gcp_config.zone), - cluster_name=self.task_test_config.cluster_name, - ) - _ = run_workload >> wait_for_workload_start - return group - def post_process(self, result_location: Optional[str] = None) -> DAGNode: - """Process metrics and metadata, and insert them into BigQuery tables. - - Returns: - A DAG node that executes the post process. - """ - with TaskGroup(group_id="post_process") as group: - process_id = metric.generate_process_id.override(retries=0)() - post_process_metrics = metric.process_metrics.override(retries=0)( - process_id, - self.task_test_config, - self.task_metric_config, - self.task_gcp_config, - folder_location=result_location, + def intermediary_flow( + self, + launch_workload: DAGNode, + workload_id: str, + gcs_path: str, + xpk_branch: str, + ) -> DAGNode: + """Inject name generation tasks before launch_workload.""" + if self.task_metric_config.profile: + flow = ( + self._run_name_task + >> (self._tb_file_location_task, self._profile_file_location_task) + >> launch_workload ) + else: + flow = ( + self._run_name_task + >> self._tb_file_location_task + >> launch_workload + ) + return flow - if self.task_metric_config and self.task_metric_config.profile: - self.task_metric_config.profile.metrics = ( - metric.xplane_to_metrics.override(retries=0)( - self.task_metric_config.profile.file_location - ) - ) - _ = ( - process_id - >> self.task_metric_config.profile.metrics - >> post_process_metrics - ) - else: - _ = process_id >> post_process_metrics - return group @dataclasses.dataclass From 985fb02073ce5c699fb7ec3fed3eb917c6df9a0a Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Fri, 10 Jul 2026 03:58:53 +0000 Subject: [PATCH 09/18] Correct the rebase mistake on fixing DAG status changes --- xlml/apis/task.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 5ba3ad814..d69fba5d3 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -458,12 +458,13 @@ def run_model( zone=self.task_gcp_config.zone, cluster_name=self.task_test_config.cluster_name, xpk_branch=xpk_branch, - ) + ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True) intermediary_flow = self.intermediary_flow( launch_workload, workload_id, gcs_path, xpk_branch ) + _ = dummy_op_for_teardown >> launch_workload _ = ( (workload_id, gcs_path) >> intermediary_flow From 3ba8dc0baea133660b5981bc114dd64dabf37826 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Fri, 10 Jul 2026 08:49:45 +0000 Subject: [PATCH 10/18] fix: use kwargs to avoid parameter misalignments in XpkTask --- xlml/apis/task.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index d69fba5d3..ad1df14b2 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -368,13 +368,14 @@ def run( """ with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: run_model, gcs_path = self.run_model( - gcs_location, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - max_restart, + gcs_location=gcs_location, + use_vertex_tensorboard=use_vertex_tensorboard, + use_pathways=use_pathways, + skip_post_process=skip_post_process, + ramdisk_directory=ramdisk_directory, + mtc_enabled=mtc_enabled, + xpk_branch=xpk_branch, + max_restart=max_restart, ) if not skip_post_process: _ = run_model >> self.post_process(gcs_path) From df05adddd7abbb0b935fc30fd1852a0a2c1ea5d9 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Fri, 10 Jul 2026 10:24:43 +0000 Subject: [PATCH 11/18] Fix pre-commit errors --- xlml/apis/task.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index ad1df14b2..ec9c0c3bc 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -673,6 +673,7 @@ def task_path_decider(check_file_exists: bool = False) -> str: return group + @dataclasses.dataclass class XpkNameGenAndQuarantineTask(XpkTask): """Task for running XPK workloads with name generation and quarantine.""" @@ -683,8 +684,12 @@ class XpkNameGenAndQuarantineTask(XpkTask): # Temporary variables to pass tasks from run() to intermediary_flow() _run_name_task: Any = dataclasses.field(default=None, init=False, repr=False) - _tb_file_location_task: Any = dataclasses.field(default=None, init=False, repr=False) - _profile_file_location_task: Any = dataclasses.field(default=None, init=False, repr=False) + _tb_file_location_task: Any = dataclasses.field( + default=None, init=False, repr=False + ) + _profile_file_location_task: Any = dataclasses.field( + default=None, init=False, repr=False + ) def run( self, @@ -724,10 +729,14 @@ def run( # Update profile file location if profile config exists if self.task_metric_config.profile: - self._profile_file_location_task = name_format.generate_profile_file_location( - self._run_name_task, self.task_metric_config.profile.file_location + self._profile_file_location_task = ( + name_format.generate_profile_file_location( + self._run_name_task, self.task_metric_config.profile.file_location + ) + ) + self.task_metric_config.profile.file_location = ( + self._profile_file_location_task ) - self.task_metric_config.profile.file_location = self._profile_file_location_task # 2. Run under quarantine task group if applicable if QuarantineTests.is_quarantined(test_name) and self.quarantine_task_group: @@ -770,15 +779,11 @@ def intermediary_flow( ) else: flow = ( - self._run_name_task - >> self._tb_file_location_task - >> launch_workload + self._run_name_task >> self._tb_file_location_task >> launch_workload ) return flow - - @dataclasses.dataclass class GpuCreateResourceTask(BaseTask): """This is a class to set up tasks for GPU. From 7565253afab158ff27598c5577664c9c9fee7d83 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Fri, 10 Jul 2026 11:04:19 +0000 Subject: [PATCH 12/18] Fix pylint errors --- .../configs/maxtext_sweep_gke_config.py | 13 +++++++++---- xlml/apis/task.py | 17 +++++++++-------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/dags/multipod/configs/maxtext_sweep_gke_config.py b/dags/multipod/configs/maxtext_sweep_gke_config.py index 582593826..61182aad2 100644 --- a/dags/multipod/configs/maxtext_sweep_gke_config.py +++ b/dags/multipod/configs/maxtext_sweep_gke_config.py @@ -16,7 +16,7 @@ import datetime from xlml.apis import gcp_config, metric_config, task, test_config -from dags.common.vm_resource import TpuVersion, XpkClusterConfig +from dags.common.vm_resource import XpkClusterConfig import itertools from typing import List, Iterable, Dict, Any @@ -49,8 +49,12 @@ def get_maxtext_sweep_gke_config( docker_image: str, base_output_directory: str, base_run_model_cmds: Iterable[str], - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.BENCHMARK_DATASET, - metric_aggregation_strategy: metric_config.AggregationStrategy = metric_config.AggregationStrategy.MEDIAN, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.BENCHMARK_DATASET + ), + metric_aggregation_strategy: metric_config.AggregationStrategy = ( + metric_config.AggregationStrategy.MEDIAN + ), dataset_project: str = None, composer_project: str = None, enable_profile_config: bool = False, @@ -77,7 +81,8 @@ def get_maxtext_sweep_gke_config( for param, values in sweep_params.items(): sweep_params_list.append([(param, val) for val in values]) - # Generate all combinations of sweep param configurations and create a XpkTask for each one + # Generate all combinations of sweep param configurations + # and create a XpkTask for each one xpk_task_list = [] for idx, config in enumerate(itertools.product(*sweep_params_list)): config_dict = {key: value for (key, value) in config} diff --git a/xlml/apis/task.py b/xlml/apis/task.py index ec9c0c3bc..9b72d0fd0 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -387,13 +387,13 @@ def run_model( gcs_location: Optional[airflow.XComArg] = None, use_vertex_tensorboard: bool = False, use_pathways: bool = False, - skip_post_process: bool = False, + skip_post_process: bool = False, # pylint: disable=unused-argument ramdisk_directory: str = "", mtc_enabled: bool = False, xpk_branch: str = xpk.MAIN_BRANCH, - last_node: bool = False, + last_node: bool = False, # pylint: disable=unused-argument max_restart: int = 0, - check_file_exists: bool = False, + check_file_exists: bool = False, # pylint: disable=unused-argument ) -> DAGNode: """Run a test job within a docker image. @@ -554,9 +554,9 @@ def post_process(self, result_location: Optional[str] = None) -> DAGNode: def intermediary_flow( self, launch_workload: DAGNode, - workload_id: str, - gcs_path: str, - xpk_branch: str, + workload_id: str, # pylint: disable=unused-argument + gcs_path: str, # pylint: disable=unused-argument + xpk_branch: str, # pylint: disable=unused-argument ) -> DAGNode: return launch_workload @@ -610,7 +610,7 @@ def wait_for_workload_reach_step_and_interruption( launch_workload: DAGNode, workload_id: str, gcs_path: str, - xpk_branch: str, + xpk_branch: str, # pylint: disable=unused-argument ) -> DAGNode: """Wait for workload to reach specific step and trigger node deletion.""" with TaskGroup( @@ -706,7 +706,8 @@ def run( """Run a test job with name generation and quarantine.""" test_name = self.task_test_config.benchmark_id - # 1. Generate run name and file location tasks first (they are Airflow task objects) + # 1. Generate run name and file location tasks first + # (they are Airflow task objects) self._run_name_task = name_format.generate_run_name( self.task_test_config.benchmark_id ) From e96aa550cf579d964bd21ab4ca9fbf9112c5b334 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Fri, 10 Jul 2026 11:20:35 +0000 Subject: [PATCH 13/18] Fix linting errors --- .../maxtext_profile_namegen_example_dag.py | 3 +- dags/multipod/configs/gke_config.py | 67 +++++++++++++------ 2 files changed, 49 insertions(+), 21 deletions(-) diff --git a/dags/examples/maxtext_profile_namegen_example_dag.py b/dags/examples/maxtext_profile_namegen_example_dag.py index 115a4fd17..71f3ac1e4 100644 --- a/dags/examples/maxtext_profile_namegen_example_dag.py +++ b/dags/examples/maxtext_profile_namegen_example_dag.py @@ -15,7 +15,8 @@ """An example DAG to extract profile metrics. Pretraining mixtral-8x7b model on 1xv4-128. -Profile extraction can be easily integrated with gke_config + (to_name_gen_and_quarantine_task + run). +Profile extraction can be easily integrated with gke_config ++ (to_name_gen_and_quarantine_task + run). """ import datetime diff --git a/dags/multipod/configs/gke_config.py b/dags/multipod/configs/gke_config.py index bfe1db888..1876856fb 100644 --- a/dags/multipod/configs/gke_config.py +++ b/dags/multipod/configs/gke_config.py @@ -14,13 +14,13 @@ """Utilities to construct configs for maxtext DAG on GKE.""" -from dags.common import test_owner +import datetime +from typing import Any, Iterable + +from dags import gcs_bucket +from dags.common.vm_resource import Project, XpkClusters from xlml.apis import gcp_config, metric_config, task, test_config from xlml.apis.xpk_cluster_config import XpkClusterConfig -from dags import gcs_bucket -from dags.common.vm_resource import TpuVersion, Project, XpkClusters, GpuVersion, CpuVersion -from typing import Any, Iterable -import datetime def get_gke_config( @@ -31,7 +31,9 @@ def get_gke_config( run_model_cmds: Iterable[str], cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, base_output_directory: str = None, @@ -90,7 +92,9 @@ def get_gke_config_with_interrupt( expect_reach_to_step: int, cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, base_output_directory: str = None, @@ -153,7 +157,9 @@ def get_gke_config_with_name_gen_and_quarantine( run_model_cmds: Iterable[str], cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, base_output_directory: str = None, @@ -216,7 +222,9 @@ def get_gke_maxtext_nightly_config( test_owner: str, cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, ) -> task.XpkTask: @@ -234,18 +242,25 @@ def get_gke_maxtext_nightly_config( base_output_directory = ( f"{gcs_bucket.BASE_OUTPUT_DIR}/maxtext/nightly/automated/{current_date}" ) - run_name = f"{num_slices}slice-V{cluster.device_version.value}_{cluster.core_count}-maxtext-nightly-{current_datetime}" + run_name = ( + f"{num_slices}slice-V{cluster.device_version.value}_{cluster.core_count}" + f"-maxtext-nightly-{current_datetime}" + ) run_model_cmds = ( "bash src/dependencies/scripts/preflight.sh PLATFORM=GKE", ( "JAX_PLATFORM_NAME=TPU XLA_FLAGS='--xla_dump_to=/tmp/xla_dump/'" " ENABLE_PJRT_COMPATIBILITY=true" - f" python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name={run_name}" - f" base_output_directory={base_output_directory}" + " python3 -m maxtext.trainers.pre_train.train" + " src/maxtext/configs/base.yml" + f" run_name={run_name} base_output_directory={base_output_directory}" " dataset_path=gs://max-datasets-rogue dataset_type=synthetic" - " model_name=llama3-8b per_device_batch_size=12 reuse_example_batch=1 metrics_file='metrics.txt'" - " steps=50 enable_checkpointing=false profiler=xplane upload_all_profiler_results=true skip_first_n_steps_for_profiler=10 profiler_steps=10 gcs_metrics=true" + " model_name=llama3-8b per_device_batch_size=12 reuse_example_batch=1" + " metrics_file='metrics.txt' steps=50 enable_checkpointing=false" + " profiler=xplane upload_all_profiler_results=true" + " skip_first_n_steps_for_profiler=10 profiler_steps=10" + " gcs_metrics=true" ), ) @@ -316,7 +331,9 @@ def get_gke_gpt3_6b_nightly_config( test_owner: str, cluster: XpkClusterConfig = XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, ) -> task.XpkTask: @@ -334,18 +351,26 @@ def get_gke_gpt3_6b_nightly_config( base_output_directory = ( f"{gcs_bucket.BASE_OUTPUT_DIR}/maxtext/nightly/automated/{current_date}" ) - run_name = f"{num_slices}slice-V{cluster.device_version.value}_{cluster.core_count}-gpt3-6b-nightly-{current_datetime}" + run_name = ( + f"{num_slices}slice-V{cluster.device_version.value}_{cluster.core_count}" + f"-gpt3-6b-nightly-{current_datetime}" + ) run_model_cmds = ( "bash src/dependencies/scripts/preflight.sh PLATFORM=GKE", ( "JAX_PLATFORM_NAME=TPU XLA_FLAGS='--xla_dump_to=/tmp/xla_dump/'" " ENABLE_PJRT_COMPATIBILITY=true" - f" python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name={run_name} model_name=gpt3-6b" + " python3 -m maxtext.trainers.pre_train.train" + " src/maxtext/configs/base.yml" + f" run_name={run_name} model_name=gpt3-6b" f" base_output_directory={base_output_directory}" " dataset_path=gs://max-datasets-rogue dataset_type=synthetic" - " per_device_batch_size=12 reuse_example_batch=1 global_parameter_scale=1 metrics_file='metrics.txt'" - " steps=50 enable_checkpointing=false profiler=xplane upload_all_profiler_results=true skip_first_n_steps_for_profiler=10 profiler_steps=10 gcs_metrics=true" + " per_device_batch_size=12 reuse_example_batch=1" + " global_parameter_scale=1 metrics_file='metrics.txt' steps=50" + " enable_checkpointing=false profiler=xplane" + " upload_all_profiler_results=true skip_first_n_steps_for_profiler=10" + " profiler_steps=10 gcs_metrics=true" ), ) @@ -379,7 +404,9 @@ def get_maxtext_cpu_end_to_end_gke_config( cluster: XpkClusterConfig = XpkClusters.CPU_N2_STANDARD_64_CLUSTER, machine_count: int = 1, num_slices: int = 1, - dataset_name: metric_config.DatasetOption = metric_config.DatasetOption.XLML_DATASET, + dataset_name: metric_config.DatasetOption = ( + metric_config.DatasetOption.XLML_DATASET + ), dataset_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, composer_project: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value, base_output_directory: str = None, From 0cea1b95fa890f91b5255a55515d007c6cddb101 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Wed, 15 Jul 2026 15:02:30 +0800 Subject: [PATCH 14/18] Refactor the XpkClass --- xlml/apis/task.py | 336 ++++++++++++++++++++++++---------------------- 1 file changed, 173 insertions(+), 163 deletions(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 9b72d0fd0..f3c44f6c4 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -367,7 +367,7 @@ def run( post_process. """ with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: - run_model, gcs_path = self.run_model( + run_model, gcs_path = self._run_model( gcs_location=gcs_location, use_vertex_tensorboard=use_vertex_tensorboard, use_pathways=use_pathways, @@ -378,22 +378,19 @@ def run( max_restart=max_restart, ) if not skip_post_process: - _ = run_model >> self.post_process(gcs_path) + _ = run_model >> self._post_process(gcs_path) return group - def run_model( + def _run_model( self, gcs_location: Optional[airflow.XComArg] = None, use_vertex_tensorboard: bool = False, use_pathways: bool = False, - skip_post_process: bool = False, # pylint: disable=unused-argument ramdisk_directory: str = "", mtc_enabled: bool = False, xpk_branch: str = xpk.MAIN_BRANCH, - last_node: bool = False, # pylint: disable=unused-argument max_restart: int = 0, - check_file_exists: bool = False, # pylint: disable=unused-argument ) -> DAGNode: """Run a test job within a docker image. @@ -404,24 +401,18 @@ def run_model( gcs_location: GCS path for all artifacts of the test. use_vertex_tensorboard: Set to True to view workload data on Vertex AI Tensorboard. - expect_reach_to_step: The training step at which the node interruption - should be triggered. use_pathways: Set to True to use the Pathways execution framework. - skip_post_process: If True, the post processing step will be skipped. ramdisk_directory: The directory for enabling emergency checkpointing. mtc_enabled: Set to True to enable Multi-tier Checkpointing (MTC). xpk_branch: The specific git branch of the xpk tool to use. - last_node: If True, the interruption will target the last node in the - workload; otherwise, it targets the first node. max_restart: By default, this is 0. This will restart the job with flag "--max-restarts" - check_file_exists: By default, this is False. If set to True, - task branch task_path_decider will be performed. 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: @@ -434,7 +425,7 @@ def run_model( task_id="dummy_op_for_teardown" ).as_setup() - launch_workload = self.launch_workload( + launch_workload = self._launch_workload( workload_id, gcs_path, use_vertex_tensorboard, @@ -444,37 +435,46 @@ def run_model( xpk_branch, max_restart, ) - 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=gke.zone_to_region(self.task_gcp_config.zone), - cluster_name=self.task_test_config.cluster_name, + + wait_for_workload_completion = self._get_wait_for_workload_completion( + workload_id ) - 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, - xpk_branch=xpk_branch, + clean_up_workload = self._get_clean_up_workload( + workload_id, xpk_branch ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True) - intermediary_flow = self.intermediary_flow( - launch_workload, workload_id, gcs_path, xpk_branch - ) - - _ = dummy_op_for_teardown >> launch_workload _ = ( (workload_id, gcs_path) - >> intermediary_flow + >> dummy_op_for_teardown + >> launch_workload >> wait_for_workload_completion >> clean_up_workload ) return group, gcs_path - def launch_workload( + def _get_wait_for_workload_completion(self, workload_id: str) -> DAGNode: + return 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=gke.zone_to_region(self.task_gcp_config.zone), + cluster_name=self.task_test_config.cluster_name, + ) + + def _get_clean_up_workload( + self, workload_id: str, xpk_branch: str + ) -> DAGNode: + return 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, + xpk_branch=xpk_branch, + ) + + def _launch_workload( self, workload_id: str, gcs_path: str, @@ -519,7 +519,7 @@ def launch_workload( _ = run_workload >> wait_for_workload_start return group - def post_process(self, result_location: Optional[str] = None) -> DAGNode: + def _post_process(self, result_location: Optional[str] = None) -> DAGNode: """Process metrics and metadata, and insert them into BigQuery tables. Returns: @@ -551,15 +551,6 @@ def post_process(self, result_location: Optional[str] = None) -> DAGNode: return group - def intermediary_flow( - self, - launch_workload: DAGNode, - workload_id: str, # pylint: disable=unused-argument - gcs_path: str, # pylint: disable=unused-argument - xpk_branch: str, # pylint: disable=unused-argument - ) -> DAGNode: - return launch_workload - @dataclasses.dataclass class XpkNodeInterruptionTask(XpkTask): @@ -569,48 +560,87 @@ class XpkNodeInterruptionTask(XpkTask): last_node: bool = False check_file_exists: bool = False - def run( + def _run_model( 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, max_restart: int = 0, ) -> DAGNode: - """Run a test job with node interruption.""" - return super().run( - gcs_location=gcs_location, - use_vertex_tensorboard=use_vertex_tensorboard, - use_pathways=use_pathways, - skip_post_process=skip_post_process, - ramdisk_directory=ramdisk_directory, - mtc_enabled=mtc_enabled, - xpk_branch=xpk_branch, - max_restart=max_restart, - ) + """Run a test job within a docker image. - def intermediary_flow( - self, - launch_workload: DAGNode, - workload_id: str, - gcs_path: str, - xpk_branch: str, - ) -> DAGNode: - """Inject interruption steps after launch_workload.""" - return self.wait_for_workload_reach_step_and_interruption( - launch_workload, workload_id, gcs_path, xpk_branch - ) + Will run a workload with an injected interruption of a GKE node. + Then is expected to automatically restart and continuing running. + + 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. + use_pathways: Set to True to use the Pathways execution framework. + ramdisk_directory: The directory for enabling emergency checkpointing. + mtc_enabled: Set to True to enable Multi-tier Checkpointing (MTC). + xpk_branch: The specific git branch of the xpk tool to use. + max_restart: By default, this is 0. + This will restart the job with flag "--max-restarts" + 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, + ) + + dummy_op_for_teardown = EmptyOperator( + task_id="dummy_op_for_teardown" + ).as_setup() + + launch_workload = self._launch_workload( + workload_id, + gcs_path, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + max_restart, + ) + + wait_for_workload_reach_step_and_interruption = ( + self.wait_for_workload_reach_step_and_interruption( + workload_id, gcs_path + ) + ) + + wait_for_workload_completion = self._get_wait_for_workload_completion( + workload_id + ) + + clean_up_workload = self._get_clean_up_workload( + workload_id, xpk_branch + ).as_teardown(setups=dummy_op_for_teardown, on_failure_fail_dagrun=True) + + _ = ( + dummy_op_for_teardown + >> launch_workload + >> wait_for_workload_reach_step_and_interruption + >> wait_for_workload_completion + >> clean_up_workload + ) + return group, gcs_path def wait_for_workload_reach_step_and_interruption( self, - launch_workload: DAGNode, workload_id: str, gcs_path: str, - xpk_branch: str, # pylint: disable=unused-argument ) -> DAGNode: """Wait for workload to reach specific step and trigger node deletion.""" with TaskGroup( @@ -663,11 +693,7 @@ def task_path_decider(check_file_exists: bool = False) -> str: last_node=self.last_node, ) - _ = ( - launch_workload - >> wait_for_workload_to_reach_step - >> maybe_check_file_exists - ) + _ = wait_for_workload_to_reach_step >> maybe_check_file_exists _ = maybe_check_file_exists >> [wait_for_file_to_exist, do_nothing] _ = [wait_for_file_to_exist, do_nothing] >> run_node_interruption @@ -682,107 +708,91 @@ class XpkNameGenAndQuarantineTask(XpkTask): run_name_env: str = "M_RUN_NAME" nested_run_name_in_tb_file_location: bool = True - # Temporary variables to pass tasks from run() to intermediary_flow() - _run_name_task: Any = dataclasses.field(default=None, init=False, repr=False) - _tb_file_location_task: Any = dataclasses.field( - default=None, init=False, repr=False - ) - _profile_file_location_task: Any = dataclasses.field( - default=None, init=False, repr=False - ) - def run( 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, - max_restart: int = 0, ) -> DAGNode: - """Run a test job with name generation and quarantine.""" - test_name = self.task_test_config.benchmark_id - - # 1. Generate run name and file location tasks first - # (they are Airflow task objects) - self._run_name_task = name_format.generate_run_name( - self.task_test_config.benchmark_id - ) - self._tb_file_location_task = name_format.generate_tb_file_location( - self._run_name_task, - self.task_metric_config.tensorboard_summary.file_location, - self.nested_run_name_in_tb_file_location, - ) - - # Set run_name in run_model_cmds - new_run_model_cmds = [f"export {self.run_name_env}={self._run_name_task}"] - for cmd in self.task_test_config.run_model_cmds: - new_run_model_cmds.append(cmd) - self.task_test_config.run_model_cmds = new_run_model_cmds - - # Update tensorboard file location - self.task_metric_config.tensorboard_summary.file_location = ( - self._tb_file_location_task - ) - - # Update profile file location if profile config exists - if self.task_metric_config.profile: - self._profile_file_location_task = ( - name_format.generate_profile_file_location( - self._run_name_task, self.task_metric_config.profile.file_location - ) - ) - self.task_metric_config.profile.file_location = ( - self._profile_file_location_task - ) + """Generate a unique run name, tensorboard file location, + and profile file location (if metric config has profile), + then run a test job within a docker image. - # 2. Run under quarantine task group if applicable - if QuarantineTests.is_quarantined(test_name) and self.quarantine_task_group: + Returns: + A task group with the following tasks chained: generate_run_name, + generate_tb_file_location, generate_profile_file_location (optional), + run provision, run_model, post_process. + """ + test_name = self.task_test_config.benchmark_id + if QuarantineTests.is_quarantined(test_name): with self.quarantine_task_group: - return super().run( - gcs_location=gcs_location, - use_vertex_tensorboard=use_vertex_tensorboard, - use_pathways=use_pathways, - skip_post_process=skip_post_process, - ramdisk_directory=ramdisk_directory, - mtc_enabled=mtc_enabled, - xpk_branch=xpk_branch, - max_restart=max_restart, + return self._run_with_run_name_generation( + use_pathways, + xpk_branch, + self.run_name_env, + self.nested_run_name_in_tb_file_location, ) else: - return super().run( - gcs_location=gcs_location, - use_vertex_tensorboard=use_vertex_tensorboard, - use_pathways=use_pathways, - skip_post_process=skip_post_process, - ramdisk_directory=ramdisk_directory, - mtc_enabled=mtc_enabled, - xpk_branch=xpk_branch, - max_restart=max_restart, + return self._run_with_run_name_generation( + use_pathways, + xpk_branch, + self.run_name_env, + self.nested_run_name_in_tb_file_location, ) - def intermediary_flow( - self, - launch_workload: DAGNode, - workload_id: str, - gcs_path: str, - xpk_branch: str, + def _run_with_run_name_generation( + self, use_pathways: bool = False, xpk_branch: str = xpk.MAIN_BRANCH ) -> DAGNode: - """Inject name generation tasks before launch_workload.""" - if self.task_metric_config.profile: - flow = ( - self._run_name_task - >> (self._tb_file_location_task, self._profile_file_location_task) - >> launch_workload + with TaskGroup( + group_id=self.task_test_config.benchmark_id, prefix_group_id=True + ) as group: + run_name = name_format.generate_run_name( + self.task_test_config.benchmark_id ) - else: - flow = ( - self._run_name_task >> self._tb_file_location_task >> launch_workload + tb_file_location = name_format.generate_tb_file_location( + run_name, + self.task_metric_config.tensorboard_summary.file_location, + self.nested_run_name_in_tb_file_location, + ) + + # Set run_name in run_model_cmds + new_run_model_cmds = [f"export {self.run_name_env}={run_name}"] + for cmd in self.task_test_config.run_model_cmds: + new_run_model_cmds.append(cmd) + self.task_test_config.run_model_cmds = new_run_model_cmds + + # Update tensorboard file location + self.task_metric_config.tensorboard_summary.file_location = ( + tb_file_location ) - return flow + + # Update profile file location + if self.task_metric_config.profile: + profile_file_location = name_format.generate_profile_file_location( + run_name, self.task_metric_config.profile.file_location + ) + self.task_metric_config.profile.file_location = profile_file_location + + run_model, gcs_path = self._run_model( + use_pathways=use_pathways, xpk_branch=xpk_branch + ) + _ = ( + run_name + >> (tb_file_location, profile_file_location) + >> run_model + >> self._post_process(gcs_path) + ) + else: + run_model, gcs_path = self._run_model( + use_pathways=use_pathways, xpk_branch=xpk_branch + ) + _ = ( + run_name + >> tb_file_location + >> run_model + >> self._post_process(gcs_path) + ) + + return group @dataclasses.dataclass From d2ef52c2ddaa7647cc0d3cb0652f5683c0cda126 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Wed, 15 Jul 2026 15:32:07 +0800 Subject: [PATCH 15/18] Fix DAG failures --- xlml/apis/task.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index f3c44f6c4..9e1b927a8 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -728,15 +728,11 @@ def run( return self._run_with_run_name_generation( use_pathways, xpk_branch, - self.run_name_env, - self.nested_run_name_in_tb_file_location, ) else: return self._run_with_run_name_generation( use_pathways, xpk_branch, - self.run_name_env, - self.nested_run_name_in_tb_file_location, ) def _run_with_run_name_generation( From 9c14fad07520d5b7693bc5e1beb9741a114350b3 Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Wed, 15 Jul 2026 15:51:17 +0800 Subject: [PATCH 16/18] Fix DAG failures --- xlml/apis/task.py | 1 - 1 file changed, 1 deletion(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 9e1b927a8..9d68dd865 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -371,7 +371,6 @@ def run( gcs_location=gcs_location, use_vertex_tensorboard=use_vertex_tensorboard, use_pathways=use_pathways, - skip_post_process=skip_post_process, ramdisk_directory=ramdisk_directory, mtc_enabled=mtc_enabled, xpk_branch=xpk_branch, From e7bc686ec244147565019bd935be8979e146cd7b Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Wed, 15 Jul 2026 16:00:19 +0800 Subject: [PATCH 17/18] Fix DAG failures --- dags/multipod/maxtext_e2e_tpu_post_training.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/multipod/maxtext_e2e_tpu_post_training.py b/dags/multipod/maxtext_e2e_tpu_post_training.py index 352504448..659604fba 100644 --- a/dags/multipod/maxtext_e2e_tpu_post_training.py +++ b/dags/multipod/maxtext_e2e_tpu_post_training.py @@ -123,7 +123,7 @@ def get_workload_name(model, mode, length=6): run_model_cmds=training_cmd, docker_image="{{ params.docker_image }}", test_owner=test_owner.SURBHI_J, - ).run_model( + ).run( use_pathways=True, ) From 9fe57d74dc728ac0a1d809b503d43444b43e62da Mon Sep 17 00:00:00 2001 From: Ryan Liu Date: Wed, 15 Jul 2026 16:30:35 +0800 Subject: [PATCH 18/18] Rename the function --- xlml/apis/task.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 9d68dd865..8b875834a 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -614,7 +614,7 @@ def _run_model( ) wait_for_workload_reach_step_and_interruption = ( - self.wait_for_workload_reach_step_and_interruption( + self._wait_for_workload_reach_step_and_interruption( workload_id, gcs_path ) ) @@ -636,7 +636,7 @@ def _run_model( ) return group, gcs_path - def wait_for_workload_reach_step_and_interruption( + def _wait_for_workload_reach_step_and_interruption( self, workload_id: str, gcs_path: str,