From c3a60d0f9a5322a76467f7427fdb353d549a1dcb Mon Sep 17 00:00:00 2001 From: Ernie Chang Date: Tue, 29 Jul 2025 08:08:21 +0000 Subject: [PATCH 1/5] add a new DAG 'axlearn_emergency_save' --- .../axlearn_emergency_checkpoint_save_gcs.py | 102 +++ dags/axlearn/configs/axlearn_config.py | 96 +++ dags/common/test_owner.py | 1 + xlml/apis/task.py | 434 ++++++++++++- xlml/utils/axlearn.py | 589 ++++++++++++++++++ 5 files changed, 1220 insertions(+), 2 deletions(-) create mode 100644 dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py create mode 100644 dags/axlearn/configs/axlearn_config.py create mode 100644 xlml/utils/axlearn.py diff --git a/dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py b/dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py new file mode 100644 index 000000000..176b38214 --- /dev/null +++ b/dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py @@ -0,0 +1,102 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A DAG to run all supported ML models with the latest JAX/FLAX version.""" + +import datetime +from airflow import models +from dags import composer_env +from dags.multipod.configs.common import SetupMode +from dags.common import test_owner +from dags.common.vm_resource import TpuVersion, Zone, RuntimeVersion, DockerImage +from dags.axlearn.configs import axlearn_config as config +from airflow.utils.task_group import TaskGroup +from datetime import timedelta + + +# Run once a day at 6 pm UTC (11 am PST) +SCHEDULED_TIME = '0 18 * * *' if composer_env.is_prod_env() else None + +v5p_conf = { + 'cluster_name': 'ernie-axlearn-v5p', + 'tpu_version': TpuVersion.V5P, + 'tpu_cores': 128, + 'tpu_zone': Zone.US_EAST5_A.value, + 'is_tpu_reserved': False, + 'project_name': 'cienet-cmcs', + 'network': 'ernie-net', + 'subnetwork': 'ernie-subnet', + 'module': 'text.gpt.c4_trainer', + 'model_config': 'fuji-7B-v3-flash-orbax', + 'num_replica': 2, + 'runtime_version': RuntimeVersion.V2_ALPHA_TPUV5.value, +} + +common = { + 'time_out_in_min': 180, + 'task_owner': test_owner.ERNIE_C, +} + +default_args = { + 'retries': 3, + 'retry_delay': timedelta(minutes=25), +} + +with models.DAG( + dag_id='axlearn_chk_save_gcs', + schedule=SCHEDULED_TIME, + tags=[ + 'multipod_team', + 'tpu', + 'axlearn', + ], + catchup=False, + default_args=default_args, +) as dag: + with TaskGroup( + group_id='axlearn_tpu_training', prefix_group_id=False + ) as axlearn_training: + docker_images = [( + SetupMode.NIGHTLY, + DockerImage.AXLEARN_TPU_JAX_STABLE_STACK, + )] + + test_configs = { + 'v5p-128': [2], + } + for mode, image in docker_images: + for accelerator, slices in test_configs.items(): + for slice_num in slices: + # AXLearn head against JAX head + # Runs Fuji training on v5p-128 in the provided GCP Project + jax_main_fuji_v5p_8 = config.get_axlearn_tpu_config( + cluster_name=v5p_conf['cluster_name'], + docker_image=image.value, + tpu_version=v5p_conf['tpu_version'], + tpu_cores=v5p_conf['tpu_cores'], + tpu_zone=v5p_conf['tpu_zone'], + runtime_version=v5p_conf['runtime_version'], + project_name=v5p_conf['project_name'], + network=v5p_conf['network'], + subnetwork=v5p_conf['subnetwork'], + is_tpu_reserved=v5p_conf['is_tpu_reserved'], + num_replica=v5p_conf['num_replica'], + model_config=v5p_conf['model_config'], + time_out_in_min=common['time_out_in_min'], + task_owner=common['task_owner'], + ).run( + module=v5p_conf['module'], + model_config=v5p_conf['model_config'], + ) + diff --git a/dags/axlearn/configs/axlearn_config.py b/dags/axlearn/configs/axlearn_config.py new file mode 100644 index 000000000..c9d86dcb9 --- /dev/null +++ b/dags/axlearn/configs/axlearn_config.py @@ -0,0 +1,96 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Utilities to construct configs for solutionsteam_jax_bite DAG.""" + + +import datetime +from typing import Tuple, Optional +from dags.common import test_owner +from xlml.apis import gcp_config, metric_config, task, test_config +from dags import gcs_bucket +from dags.sparsity_diffusion_devx.configs import common +from dags.common.vm_resource import TpuVersion, Project +from airflow.models.taskmixin import DAGNode + + +GCS_SUBFOLDER_PREFIX = test_owner.Team.SPARSITY_DIFFUSION_DEVX.value + + +def set_up_axlearn(pinned_version, jax_version) -> Tuple[str]: + reset_version = "" + if pinned_version: + reset_version = f"cd axlearn && git reset --hard {pinned_version} && cd .." + + setup_jax = None + if jax_version: + setup_jax = common.set_up_jax_version(jax_version) + else: + setup_jax = common.set_up_nightly_jax() + + return ( + common.UPGRADE_PIP, + common.UPGRADE_SETUPTOOLS, + common.UPGRADE_PACKAGING, + "git clone https://github.com/lkolluru05/axlearn", + reset_version, + "python -m pip install ./axlearn[core,dev]", + "python -m pip install ./axlearn[gcp]", + *setup_jax, + ) + + +def get_axlearn_tpu_config( + cluster_name: str, + docker_image: str, + tpu_version: TpuVersion, + tpu_cores: int, + tpu_zone: str, + runtime_version: str, + model_config: str, + time_out_in_min: int, + task_owner: str, + num_replica: int, + is_tpu_reserved: bool = False, + project_name: Optional[Project] = Project.CLOUD_ML_AUTO_SOLUTIONS.value, + network: str = "default", + subnetwork: str = "default", +) -> task.AxlearnTask: + """Setup the axlearn tpu env config.""" + job_gcp_config = gcp_config.GCPConfig( + project_name=project_name, + zone=tpu_zone, + dataset_name=metric_config.DatasetOption.XLML_DATASET, + ) + test_name = f"axlearn_tpu_training_{model_config}" + + job_test_config = test_config.TpuGkeTest( + test_config.Tpu( + version=tpu_version, + cores=tpu_cores, + ), + test_name=test_name, + run_model_cmds=None, + set_up_cmds=None, + timeout=datetime.timedelta(minutes=time_out_in_min), + task_owner=test_owner, + num_slices=num_replica, + cluster_name=cluster_name, + docker_image=docker_image, + ) + + return task.AxlearnTask( + task_test_config=job_test_config, + task_gcp_config=job_gcp_config, + ) diff --git a/dags/common/test_owner.py b/dags/common/test_owner.py index 4789e4a4c..0d543ac8d 100644 --- a/dags/common/test_owner.py +++ b/dags/common/test_owner.py @@ -55,6 +55,7 @@ class Team(enum.Enum): # Multi-tier Checkpointing ABHINAV_S = "ABHINAV S." XUEFENG_G = "XUEFENG G." +ERNIE_C = "ERNIE C." # MLCompass ORTI_B = "Orti B." diff --git a/xlml/apis/task.py b/xlml/apis/task.py index fceddcf63..730aa4c62 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -24,7 +24,7 @@ from airflow.models.taskmixin import DAGNode from airflow.utils.task_group import TaskGroup from xlml.apis import gcp_config, metric_config, test_config -from xlml.utils import gpu, metric, name_format, ssh, tpu, xpk, gke +from xlml.utils import gpu, metric, name_format, ssh, tpu, xpk, axlearn, gke class BaseTask(abc.ABC): @@ -145,6 +145,152 @@ def run_queued_resource_test( return test +@dataclasses.dataclass +class AxlearnTask(BaseTask): + """This is a class to set up tasks for TPU/GPU axlearn. + + Attributes: + task_test_config: Test configs to run on this TPU/GPU. + task_gcp_config: Runtime TPU/GPU creation parameters. + task_metric_config: Metric configs to process metrics. + workload_provision_timeout: Time allowed for provisioning a workload. + """ + + task_test_config: Union[ + test_config.TpuGkeTest, test_config.GpuXpkTest, test_config.CpuGkeTest + ] + task_gcp_config: gcp_config.GCPConfig + task_metric_config: Optional[metric_config.MetricConfig] = None + workload_provision_timeout: datetime.timedelta = datetime.timedelta( + minutes=300 + ) + + def run( + self, + *, + gcs_location: Optional[airflow.XComArg] = None, + module: Optional[str] = None, + model_config: Optional[str] = None, + ) -> DAGNode: + """Run a test job within a docker image. + + Attributes: + gcs_location: GCS path for all artifacts of the test. + module: Set run module. + model_config: Set model config. + use_vertex_tensorboard: Set to True to view workload data on + Vertex AI Tensorboard. + + Returns: + A task group with the following tasks chained: run_model and + post_process. + """ + with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: + self.run_model( + gcs_location, + axlearn_branch=axlearn.LALITAH_BRANCH, + module=module, + model_config=model_config, + ) + return group + + def run_model( + self, + gcs_location: Optional[airflow.XComArg] = None, + module: Optional[str] = None, + model_config: Optional[str] = None, + axlearn_branch: str = "", + ) -> DAGNode: + """Run the TPU/GPU test in `get_bite_tpu_config` using axlearn. + + Attributes: + gcs_location: GCS path for all artifacts of the test. + use_vertex_tensorboard: Set to True to view workload data on + Vertex AI Tensorboard. + + Returns: + A DAG node that executes the model test. + """ + with TaskGroup(group_id="run_model") as group: + workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id) + if gcs_location: + gcs_path = gcs_location + else: + gcs_path = name_format.generate_gcs_folder_location( + self.task_test_config.gcs_subfolder, + self.task_test_config.benchmark_id, + ) + launch_workload = self.launch_workload( + workload_id, + gcs_path, + axlearn_branch, + module, + model_config, + ) + + + ((workload_id, gcs_path) >> launch_workload) + return group, gcs_path + + def launch_workload( + self, + workload_id: str, + gcs_path: str, + axlearn_branch: str = "", + module: str = None, + model_config: str = None, + ) -> DAGNode: + """Create the workload and wait for it to provision.""" + with TaskGroup(group_id="launch_workload") as group: + setup_axlearn_dpd = axlearn.set_up_axlearn_dpd(axlearn_branch) + create_conf_axlearn = axlearn.create_conf_axlearn( + cluster_name=self.task_test_config.cluster_name, + project_id=self.task_gcp_config.project_name, + zone=self.task_gcp_config.zone, + ) + activate_axlearn = axlearn.activate_axlearn( + cluster_name=self.task_test_config.cluster_name, + project_id=self.task_gcp_config.project_name, + region=self.task_gcp_config.zone[:-2], + ) + run_workload = axlearn.run_workload_axlearn( + 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, + run_name=self.task_test_config.test_name, + benchmark_id=self.task_test_config.benchmark_id, + workload_id=workload_id, + gcs_path=gcs_path, + accelerator_type=self.task_test_config.accelerator.name, + run_cmds="", + module=module, + model_config=model_config, + trainer_dir=f"gs://{self.task_gcp_config.project_name}-axlearn/{self.task_test_config.test_name}-nr-{self.task_test_config.num_slices}", + num_replicas=self.task_test_config.num_slices, + axlearn_branch=axlearn_branch, + trace_steps=[40, 90, 140, 190, 240] + ) + + wait_for_workload_start = xpk.wait_for_workload_start.override( + timeout=self.workload_provision_timeout.total_seconds() + )( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + region=self.task_gcp_config.zone[:-2], + cluster_name=self.task_test_config.cluster_name, + ) + + ( + setup_axlearn_dpd + >> create_conf_axlearn + >> activate_axlearn + >> run_workload + >> wait_for_workload_start + ) + return group + + @dataclasses.dataclass class XpkTask(BaseTask): """This is a class to set up tasks for TPU/GPU provisioned by XPK tool. @@ -201,6 +347,44 @@ def run( return group + def run_with_workload_id( + 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, + workload_id: str = None, + ) -> DAGNode: + """Run a test job within a docker image with specific workload id. + + Attributes: + gcs_location: GCS path for all artifacts of the test. + use_vertex_tensorboard: Set to True to view workload data on + Vertex AI Tensorboard. + + Returns: + A task group with the following tasks chained: run_model and + post_process. + """ + with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: + run_model, gcs_path = self.run_model_with_workload_id( + gcs_location, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + workload_id, + ) + if not skip_post_process: + run_model >> self.post_process(gcs_path) + + return group + def run_with_name_gen_and_quarantine( self, quarantine_task_group, @@ -292,6 +476,43 @@ def run_with_run_name_generation( ) return group + def run_with_node_interruption( + self, + *, + gcs_location: Optional[airflow.XComArg] = None, + use_vertex_tensorboard: bool = False, + use_pathways: bool = False, + skip_post_process: bool = False, + ramdisk_directory: str = "", + mtc_enabled: bool = False, + xpk_branch: str = xpk.MAIN_BRANCH, + last_node: bool = False, + ) -> DAGNode: + """Run a test job within a docker image. + + Attributes: + gcs_location: GCS path for all artifacts of the test. + use_vertex_tensorboard: Set to True to view workload data on + Vertex AI Tensorboard. + + Returns: + A task group with the following tasks chained: run_model and + post_process. + """ + with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: + run_model, gcs_path = self.run_model_with_node_interruption( + gcs_location, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + last_node, + ) + if not skip_post_process: + run_model >> self.post_process(gcs_path) + return group + def run_model( self, gcs_location: Optional[airflow.XComArg] = None, @@ -338,6 +559,142 @@ def run_model( 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 run_model_with_node_interruption( + self, + gcs_location: Optional[airflow.XComArg] = None, + use_vertex_tensorboard: bool = False, + use_pathways: bool = False, + ramdisk_directory: str = "", + mtc_enabled: bool = False, + xpk_branch: str = xpk.MAIN_BRANCH, + last_node: bool = False, + ) -> DAGNode: + """Run the TPU/GPU test in `task_test_config` using xpk. + Different behaviour 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. + + Returns: + A DAG node that executes the model test. + """ + with TaskGroup(group_id="run_model") as group: + workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id) + if gcs_location: + gcs_path = gcs_location + else: + gcs_path = name_format.generate_gcs_folder_location( + self.task_test_config.gcs_subfolder, + self.task_test_config.benchmark_id, + ) + + launch_workload_with_interruption = ( + self.launch_workload_with_node_interruption( + workload_id, + gcs_path, + use_vertex_tensorboard, + use_pathways, + ramdisk_directory, + mtc_enabled, + xpk_branch, + 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=self.task_gcp_config.zone[:-2], + cluster_name=self.task_test_config.cluster_name, + ) + + clean_up_workload = xpk.clean_up_workload( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + zone=self.task_gcp_config.zone, + cluster_name=self.task_test_config.cluster_name, + xpk_branch=xpk_branch, + ) + + ( + (workload_id, gcs_path) + >> launch_workload_with_interruption + >> wait_for_workload_completion + >> clean_up_workload + ) + return group, gcs_path + + def run_model_with_workload_id( + self, + gcs_location: Optional[airflow.XComArg] = None, + use_vertex_tensorboard: bool = False, + use_pathways: bool = False, + ramdisk_directory: str = "", + mtc_enabled: bool = False, + xpk_branch: str = xpk.MAIN_BRANCH, + workload_id: str = None, + ) -> DAGNode: + """Run the TPU/GPU test in `task_test_config` using xpk + with specific workload id. + + 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: + if workload_id is None: + 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, + ) + wait_for_workload_completion = xpk.wait_for_workload_completion.override( + timeout=int(self.task_test_config.timeout.total_seconds()), + )( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + region=self.task_gcp_config.zone[:-2], + cluster_name=self.task_test_config.cluster_name, + ) + clean_up_workload = xpk.clean_up_workload( workload_id=workload_id, project_id=self.task_gcp_config.project_name, @@ -396,6 +753,74 @@ def launch_workload( run_workload >> wait_for_workload_start return group + def launch_workload_with_node_interruption( + self, + workload_id: str, + gcs_path: str, + use_vertex_tensorboard: bool, + use_pathways: bool = False, + ramdisk_directory: str = "", + mtc_enabled: bool = False, + xpk_branch: str = xpk.MAIN_BRANCH, + last_node: bool = False, + ) -> DAGNode: + """Create the workload and wait for it to provision.""" + with TaskGroup(group_id="launch_workload_with_interruption") as group: + run_workload = xpk.run_workload.override( + owner=self.task_test_config.task_owner + )( + task_id="run_workload", + cluster_project=self.task_gcp_config.project_name, + zone=self.task_gcp_config.zone, + cluster_name=self.task_test_config.cluster_name, + benchmark_id=self.task_test_config.benchmark_id, + workload_id=workload_id, + gcs_path=gcs_path, + docker_image=self.task_test_config.docker_image, + accelerator_type=self.task_test_config.accelerator.name, + run_cmds=self.task_test_config.test_script, + num_slices=self.task_test_config.num_slices, + use_vertex_tensorboard=use_vertex_tensorboard, + use_pathways=use_pathways, + ramdisk_directory=ramdisk_directory, + mtc_enabled=mtc_enabled, + xpk_branch=xpk_branch, + ) + wait_for_workload_start = xpk.wait_for_workload_start.override( + timeout=self.workload_provision_timeout.total_seconds() + )( + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + region=self.task_gcp_config.zone[:-2], + cluster_name=self.task_test_config.cluster_name, + ) + wait_to_reach_step_to_interrupt = xpk.wait_for_reach_step_to_interrupt( + task_id="wait_to_reach_step_to_interrupt", + workload_id=workload_id, + project_id=self.task_gcp_config.project_name, + region=self.task_gcp_config.zone[:-2], + cluster_name=self.task_test_config.cluster_name, + step_to_interrupt="40", + ) + run_node_interruption = xpk.delete_node.override( + owner=self.task_test_config.task_owner + )( + 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, + ) + + ( + run_workload + >> wait_for_workload_start + >> wait_to_reach_step_to_interrupt + >> run_node_interruption + ) + return group + def post_process(self, result_location: Optional[str] = None) -> DAGNode: """Process metrics and metadata, and insert them into BigQuery tables. @@ -530,7 +955,12 @@ def run_with_existing_instance(self) -> DAGNode: def provision_via_existing_instance( self, - ) -> Tuple[DAGNode, airflow.XComArg, airflow.XComArg, airflow.XComArg,]: + ) -> Tuple[ + DAGNode, + airflow.XComArg, + airflow.XComArg, + airflow.XComArg, + ]: """Provision an existing GPU accelerator. Returns: diff --git a/xlml/utils/axlearn.py b/xlml/utils/axlearn.py new file mode 100644 index 000000000..bbf2d85a8 --- /dev/null +++ b/xlml/utils/axlearn.py @@ -0,0 +1,589 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 # +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Utilities to run workloads with xpk (https://github.com/AI-Hypercomputer/xpk).""" + +import uuid +from absl import logging +from airflow.decorators import task +from airflow.hooks.subprocess import SubprocessHook +import re +from datetime import timedelta + +LALITAH_BRANCH = "lkolluru-orbax-fuji-v2" +SAM_BRANCH = "orbax-fuji-v2" + + +# This function do some hacks to get Axlearn working with Airlfow +# One of them is deleting some unuseful packages in [dev] dependencies. We only need to run axlearn CLI +@task(execution_timeout=timedelta(hours=1)) +def set_up_axlearn_dpd(branch: str): + if branch == LALITAH_BRANCH or branch == SAM_BRANCH: + logging.info(f"Using custom branch ==> {branch}") + clone_branch = ( + f"git clone --branch {branch} https://github.com/lkolluru05/axlearn.git" + f" $HOME/axlearn" + ) + + # Maybe add these lines + install_python3_cmd = [ + "rm -rf ~/.pyenv", + "rm -rf ~/my_venv", + "curl https://pyenv.run | bash", + f"echo 'export PYENV_ROOT=\"$HOME/.pyenv\"' >> ~/.bashrc ", + f"echo 'export PYENV_ROOT=\"$HOME/.pyenv\"' >> ~/.profile ", + f"echo '[[ -d $PYENV_ROOT/bin ]] && export PATH=\"$PYENV_ROOT/bin:$PATH\"' >> ~/.bashrc ", + f"echo '[[ -d $PYENV_ROOT/bin ]] && export PATH=\"$PYENV_ROOT/bin:$PATH\"' >> ~/.profile", + f"echo 'eval \"$(pyenv init -)\"' >> ~/.bashrc ", + f"echo 'eval \"$(pyenv init -)\"' >> ~/.profile", + f"source ~/.bashrc ", + f"source ~/.profile", + f"pyenv install 3.10.12 && pyenv global 3.10.12", + "python -m venv ~/my_venv", + f"source ~/my_venv/bin/activate" + ] + + # TODO: Need to think a better way to do this. + hack_dpndcies_cmd = [ + "sed -i '/^dev = \[/,/^\]/ { /^dev = \[/b; /\[core\]/ { s/^\s*#\s*//; b; }; /\[orbax\]/ { s/^\s*#\s*//; b; }; /^\s*\"/ { s/^\s*#\?\s*//; s/^/#/; }; }' pyproject.toml" + ] + + cmds = [ + "set -xue", + "rm -rf $HOME/axlearn", + clone_branch, + *install_python3_cmd, + "python --version", + f"cd ~/axlearn/ ", + *hack_dpndcies_cmd, + f"pip install -e '.[core,gcp,dev]'", + "pip list", + "pyenv rehash", + "which axlearn", + ] + hook = SubprocessHook() + result = hook.run_command(["bash", "-c", ";".join(cmds)]) + + assert ( + result.exit_code == 0 + ), f"Set up axlearn dependencies command failed with code {result.exit_code}" + + +@task +def activate_axlearn( + cluster_name: str, + project_id: str, + region: str, +): + """Activate axlearn.""" + + cmds = [ + "set -xue", + "cat ~/.bashrc", + "cd ~/axlearn", + "source ~/my_venv/bin/activate", + "python --version", + "which axlearn", + "axlearn gcp config activate", + f"gcloud container clusters get-credentials {cluster_name} --region {region} --project {project_id}", + ] + hook = SubprocessHook() + result = hook.run_command(["bash", "-c", ";".join(cmds)]) + assert ( + result.exit_code == 0 + ), f"Set up axlearn dependencies command failed with code {result.exit_code}" + + +@task +def generate_workload_id(benchmark_id: str) -> str: + """Generate a valid workload ID.""" + + # Remove all non-alphanumeric characters, and truncate to ensure the result + # is less than 40 characters. + short_benchmark = re.sub(r"[^a-zA-Z0-9-]+", "", benchmark_id)[:32] + short_id = str(uuid.uuid4())[:8] + return f"{short_benchmark}{short_id}" + + +@task(execution_timeout=timedelta(hours=1)) +def create_conf_axlearn( + cluster_name: str, + project_id: str, + zone: str, +): + """create axlearn config file.""" + command_string = f"cat << 'CONFIG_EOF' > ~/axlearn/.axlearn/axlearn.default.config\n [gcp]\n_active = \"{project_id}:{zone}\"\n\n[gcp.\"{project_id}:{zone}\"]\nproject = \"{project_id}\"\nregion = \"{zone[:-2]}\"\nzone = \"{zone}\"\ngke_cluster = \"{cluster_name}\"\ncluster = \"{cluster_name}\"\nlabels = \"tpu-v5p\"\ndocker_repo = \"us-docker.pkg.dev/{project_id}/axlearn\"\ndefault_dockerfile = \"Dockerfile\"\npermanent_bucket = \"{project_id}-axlearn\"\nprivate_bucket = \"{project_id}-axlearn\"\nttl_bucket = \"{project_id}-axlearn\"\nCONFIG_EOF\n" + # command_string = """cat << 'CONFIG_EOF' > ~/axlearn/.axlearn/axlearn.default.config + # ["gcp.cienet-cmcs:us-east5-a"] + # project = "cienet-cmcs" + # region = "us-east5" + # zone = "us-east5-a" + # gke_cluster = "camiloquinones-axlearn" + # cluster = "camiloquinones-axlearn" + # labels = "tpu-v5p" + # docker_repo = "us-docker.pkg.dev/cienet-cmcs/axlearn" + # default_dockerfile = "Dockerfile" + # permanent_bucket = "cienet-cmcs-axlearn" + # private_bucket = "cienet-cmcs-axlearn" + # ttl_bucket = "cienet-cmcs-axlearn"CONFIG_EOF""" + create_axlearn_conf = [command_string.rstrip('\n')] + cmds = [ + *create_axlearn_conf + ] + hook = SubprocessHook() + result = hook.run_command( + ["bash", "-c",";".join(cmds)] + ) + + assert ( + result.exit_code == 0 + ), f"Config Axlearn file command failed with code {result.exit_code}" + + +@task(execution_timeout=timedelta(hours=1)) +def run_workload_axlearn( + cluster_project: str, + zone: str, + cluster_name: str, + run_name: str, + accelerator_type: str = "", + module: str = "", + model_config: str = "", + trainer_dir: str = "", + num_replicas: int = 1, + trace_steps: list[str] = None, +): + """Run workload through axlearn tool.""" + + trace_list = ( + ("--trace_at_steps=" + ", ".join(map(str, trace_steps))) + if trace_steps + else " " + ) + export_var = [ + f"export CLUSTER={cluster_name}", + f"export NAME={run_name}", + f"export BASTION_TIER=disabled", + f"export DEFAULT_PROJECT_ID=$(gcloud config get project)", + "export PROJECT_ID=${PROJECT_ID:-$DEFAULT_PROJECT_ID}", + f"export INSTANCE_TYPE={accelerator_type}", + f"export NUM_REPLICAS={num_replicas}", + f"export MODULE={module}", + f"export MODEL_CONFIG={model_config}", + f"export TRAIN_DIR={trainer_dir}" + ] + logging.info(f" Cluster: {cluster_name} -- num-replicas={num_replicas} --run_name={run_name} --project={cluster_project} --zone={zone} --instance-type={accelerator_type} --module={module} --config={model_config} --trainer_dir={trainer_dir} --data_dir=gs://axlearn-public/tensorflow_datasets --jax_backend=tpu --mesh_selector={accelerator_type} --initialization_timeout=1200 Trace: {trace_list}") + workload_create_cmd = ( + f"axlearn gcp launch run --cluster=$CLUSTER --runner_name gke_tpu_single " + f" --name=$NAME --instance_type=$INSTANCE_TYPE --num_replicas=$NUM_REPLICAS --bundler_spec=allow_dirty=True " + f"--bundler_type=artifactregistry --bundler_spec=image=tpu --bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu " + f"-- \"ulimit -n 1048576; ulimit -c 0; python3 -c 'import jax; jax.devices()'; python3 -m axlearn.common.launch_trainer_main\" " + f"--module=$MODULE --config=$MODEL_CONFIG --trainer_dir=$TRAIN_DIR " + f"--data_dir=gs://axlearn-public/tensorflow_datasets --jax_backend=tpu --mesh_selector=$INSTANCE_TYPE --initialization_timeout=1200 {trace_list} " + ) + + cmds = [ + "set -xue", + "source ~/my_venv/bin/activate", + "cd ~/axlearn", + "axlearn gcp config activate", + f"gcloud container clusters get-credentials {cluster_name} --region {zone[:-2]} --project {cluster_project}", + *export_var, + workload_create_cmd, + ] + + hook = SubprocessHook() + result = hook.run_command(["bash", "-c", ";".join(cmds)]) + + assert ( + result.exit_code == 0 + ), f"Axlearn command failed with code {result.exit_code}" + + +# +# +# def _get_core_api_client( +# project_id: str, region: str, cluster_name: str +# ) -> k8s_client.CoreV1Api: +# """Create a core API client for the given cluster.""" +# client = gke.get_authenticated_client(project_id, region, cluster_name) +# +# # Initilize the client +# core_api = k8s_client.CoreV1Api(client) +# logging.info("Successful initilize k8s client from cluster response.") +# return core_api +# +# +# def _list_workload_pods( +# core_api: k8s_client.CoreV1Api, workload_id: str +# ) -> k8s_client.V1PodList: +# """List all pods for the given workload.""" +# logging.info(f"Getting pods for workload_id: {workload_id}") +# pods = core_api.list_namespaced_pod( +# label_selector=f"jobset.sigs.k8s.io/jobset-name={workload_id}", +# namespace="default", +# ) +# return pods +# +# +# def _get_batch_api_client( +# project_id: str, region: str, cluster_name: str +# ) -> k8s_client.BatchV1Api: +# """Create a batch API client for the given cluster.""" +# client = gke.get_authenticated_client(project_id, region, cluster_name) +# +# # Initilize the client +# bat h_api = k8s_client.BatchV1Api(client) +# logging.info( +# "Successful initilize k8s batch api client from cluster response." +# ) +# return batch_api +# +# +# def _get_workload_job( +# batch_api: k8s_client.BatchV1Api, workload_id: str +# ) -> k8s_client.V1Job: +# """Get the job for a given workload.""" +# logging.info(f"Getting job for workload_id: {workload_id}") +# job = batch_api.list_namespaced_job( +# label_selector=f"jobset.sigs.k8s.io/jobset-name={workload_id}", +# namespace="default", +# ) +# if len(jobs.items) == 0: +# logging.info(f"Getting job for workload_id: {workload_id}") +# return None +# +# if len(jobs.items) > 1: +# logging.info(f"Got more than one job for workload_id: {workload_id}") +# f r i, job in enumerate(jobs.items): +# logging.info(f"Job {i=}") +# logging.info(f"{job}") + +# return jobs.items[0] +# +# +# def extract_numbers(pod_name: str) -> Tuple[int, int]: +# """Extract slice and pod numbers from pod name.""" +# m tch = re.search(r"slice-job-(\d+)-(\d+)-", pod_name) +# if match: +# return int(match.group(1)), int(match.group(2)) +# return (0, 0) +# +# +# def _find_target_pod_node( +# project_id: str, +# region: str, +# cluster_name: str, +# workload_id: str, +# last_node: bool = False, +# ) -> str: +# """find the node name for the workload.""" +# core_api = _get_core_api_client(project_id, region, cluster_name) +# pods = _list_workload_pods(core_api, workload_id) +# pod_node_pairs = [] +# pattern = re.compile(r".*slice-job-(\d+)-(\d+)-\w+") +# +# for pod in pods.items: +# if pod.status.phase == "Running" and pod.spec.node_name: +# if pattern.match(pod.metadata.name): +# pod_node_pairs.append((pod.metadata.name, pod.spec.node_name)) + +# # Find the pod with the highest slice and pod numbers. +# +# # Sort by slice number, then by pod number, and get the last (highest) one +# sorted_pairs = sorted(pod_node_pairs, key=lambda x: extract_numbers(x[0])) +# target_pod, target_node = sorted_pairs[0] +# if last_node: +# target_pod, target_node = sorted_pairs[-1] +# +# logging.info("Identified Pod for node deletion:") +# logging.info(f" Pod Name: {target_pod}") +# logging.info(f" Node Name: {target_node}") +# logging.info("-" * 72) + +# return target_node +# +# +# @task.sensor(poke_interval=60, timeout=600, mode="reschedule") +# def wait_for_workload_start( +# workload_id: str, project_id: str, region: str, cluster_name: str +# ) -> bool: +# """Check if the workload has started.""" +# core_api = _get_core_api_client(project_id, region, cluster_name) +# pods = _list_workload_pods(core_api, workload_id) +# print(f"Found {len(pods.items)} pods for workload {workload_id}") +# return len(pods.items) > 0 +# +# +# @ta k.sensor(poke_interval=120, timeout=3600, mode="reschedule") +# def wait_for_reach_step_to_interrupt( +# task_id: str, +# project_id: str, +# region: str, +# cluster_name: str, +# workload_id: str, +# step_to_interrupt: str, +# ) -> bool: +# """ +# Watch any given training pod, check the given step is already reach before +# deleting a node +# """ +# core_api = _get_core_api_client(project_id, region, cluster_name) +# pods = _list_workload_pods(core_api, workload_id) +# +# i any(pod.status.phase in ["Pending"] for pod in pods.items): +# logging.info("Some of the pods is still pending. Waiting to start") +# return False +# +# try: +# for pod in pods.items: +# i pod.status.phase == "Failed": +# # Don't keep retrying if the pod has failed +# raise AirflowFailException(f"Bad pod phase: {pod.status.phase}") +# elif pod.status.phase in ["Unknown"]: +# raise RuntimeError(f"Bad pod phase: {pod.status.phase}") +# finally: +# i all(pod.status.phase in ["Running"] for pod in pods.items): +# +# # Pick last one running pod +# pod = pods.items[len(pods.items) - 1] +# logs = core_api.read_namespaced_pod_log( +# name=pod.metadata.name, namespace=pod.metadata.namespace +# ) +# +# # Just to debug TO BE DELETED +# logging.info(f"Logs for pod {pod.metadata.name}:") +# for line in logs.split("\n"): +# logging.info(line) +# if f"completed step: {step_to_interrupt}" in logs: +# # Here we return true because we are sure the step "step_to_interrupt" is already save +# logging.info("The step to be interrupt is {step_to_interrupt}") +# return True +# return False +# +# +# @task.sensor(poke_interval=60, timeout=600, mode="reschedule") +# def wait_for_workload_completion( +# workload_id: str, project_id: str, region: str, cluster_name: str +# ) -> bool: +# """Check the workload status.""" +# core_api = _get_core_api_client(project_id, region, cluster_name) +# pods = _list_workload_pods(core_api, workload_id) +# +# if not pods.items: +# logging.info(f"No pods found for workload selector: {workload_id}.") +# +# # Pathways jobs delete all pods on failure so we must also check if the job +# # is complete +# batch_api = _get_batch_api_client(project_id, region, cluster_name) +# job = _get_workload_job(batch_api, workload_id) +# if job is None: +# logging.info( +# f"No pods or jobs were found for workload selector: {workload_id}" +# ) +# return False +# +# if any(condition.type == "Failed" for condition in job.status.conditions): +# # Don't keep retrying if the job has failed +# raise AirflowFailException('Job has condition type: "Failed"') +# +# if any(condition.type == "Complete" for condition in job.status.conditions): +# logging.info( +# "No pods found but job is complete for workload selector:" +# f" {workload_id}" +# ) +# return True + +# return False +# +# if any(pod.status.phase in ["Pending", "Running"] for pod in pods.items): +# logging.info("At least one pod has yet to complete.") +# return False +# +# try: +# for pod in pods.items: +# if pod.status.phase == "Failed": +# # Don't keep retrying if the pod has failed +# raise AirflowFailException(f"Bad pod phase: {pod.status.phase}") +# elif pod.status.phase in ["Unknown"]: +# raise RuntimeError(f"Bad pod phase: {pod.status.phase}") +# finally: +# # TODO(jonbolin): log printing for GPUs, which have multiple containers +# if len(pod.spec.containers) == 1: +# # Print the logs of the last pod checked - either the first failed pod or +# # the last successful one. +# logs = core_api.read_namespaced_pod_log( +# name=pod.metadata.name, namespace=pod.metadata.namespace +# ) +# logging.info(f"Logs for pod {pod.metadata.name}:") +# for line in logs.split("\n"): +# logging.info(line) +# url = LOGGING_URL_FORMAT.format( +# project=project_id, +# region=region, +# cluster=cluster_name, +# workload_id=workload_id, +# ) +# logging.info(f"Link to logs: {url}") +# +# logging.info("All pod(s) phase are succeeded.") +# return True +# +# +# @ta k(trigger_rule="all_done") +# def clean_up_workload( +# workload_id: str, +# project_id: str, +# zone: str, +# cluster_name: str, +# xpk_branch: str = MAIN_BRANCH, +# ) -> bool: +# """Delete workload.""" +# with tempfile.TemporaryDirectory() as tmpdir: +# workload_delete_cmd = ( +# f"python {tmpdir}/xpk/xpk.py workload delete" +# f" --cluster={cluster_name} --workload={workload_id}" +# f" --project={project_id} --zone={zone}" +# ) +# +# cmds = get_xpk_setup_cmd(tmpdir, xpk_branch) +# cmds.append(workload_delete_cmd) +# hook = SubprocessHook() +# result = hook.run_command( +# ["bash", "-c", ";".join(cmds)], +# env={**os.environ, "KUBECONFIG": os.path.join(tmpdir, "xpk.conf")}, +# ) +# assert ( +# result.exit_code == 0 +# ), f"XPK clean-up failed with code {result.exit_code}" +# +# +# @task +# def delete_node( +# cluster_name: str, +# workload_id: str, +# zone: str, +# project: str, +# dry_run: bool = False, +# last_node: bool = False, +# ) -> None: +# """Delete node.""" +# nod _name = _find_target_pod_node( +# project, +# zone[:-2], +# cluster_name, +# workload_id, +# last_node, +# ) +# # Delete the specified compute instance. +# if dry_run: +# logging.info( +# f"DRY RUN: Would delete node: {node_name}" +# f"in zone: {zone} (project: {project})" +# ) +# return +# +# logging.info(f"Proceeding to delete node: {node_name}") +# try: +# # Initialize the Compute Engine client +# instances_client = compute_v1.InstancesClient() +# +# # Delete the instance +# operation = instances_client.delete( +# project=project, zone=zone, instance=node_name +# ) +# +# logging.info(f"Deletion operation started for node: {node_name}") +# logging.info(f"Operation: {operation.name}") +# logging.info(f"Deletion command executed for node: {node_name}") +# except Exception as e: +# logging.info(f"Error deleting node {node_name}: {e}", file=sys.stderr) +# sys.exit(1) +# +# +# @task +# def simple_sleep(sleep_seconds: int) -> None: +# """ +# A simple task that pauses execution for a specified number of seconds +# using time.sleep(). +# +# Note: This task occupies a worker slot for the entire sleep duration. +# It is not a sensor and does not use the 'poke' or 'reschedule' mechanism. +# +# Args: +# sleep_seconds: The number of seconds the task should sleep. +# """ +# if sleep_seconds < 0: +# logging.warning( +# f"Requested sleep time is negative: {sleep_seconds}. Skipping sleep." +# ) +# return # Or raise an error, depending on desired behavior +# logging.info( +# f"Simple Sleep Task: Starting sleep for {sleep_seconds} seconds." +# ) +# # --- The sleep happens here --- +# time.sleep(sleep_seconds) +# # ----------------------------- +# logging.info( +# f"Simple Sleep Task: Finished sleeping after {sleep_seconds} seconds." +# ) +# return +# +# +# @task.sensor(poke_interval=120, timeout=1200, mode="reschedule") +# def wait_for_training_step_complete( +# project_id: str, +# region: str, +# cluster_name: str, +# workload_id: str, +# step: str, +# polling_time: int = 20, +# ) -> bool: +# """Restore workload from specific step and calculate elapsed time.""" +# start_time = time.monotonic() # Record the start time + +# core_api = _get_core_api_client(project_id, region, cluster_name) +# +# try: +# w ile True: +# pods = _list_workload_pods(core_api, workload_id) +# if any(pod.status.phase in ["Failed", "Unknown"] for pod in pods.items): +# elapsed_time = time.monotonic() - start_time +# logging.error( +# f"Bad pod phase. Sensor failed after {elapsed_time:.2f} seconds." +# ) +# raise AirflowFailException("Bad pod phase") +# +# if all(pod.status.phase in ["Running"] for pod in pods.items): +# # Pick last one running pod +# pod = pods.items[len(pods.items) - 1] +# logs = core_api.read_namespaced_pod_log( +# name=pod.metadata.name, namespace=pod.metadata.namespace +# ) +# logging.info(f"Logs for pod {pod.metadata.name}:") +# for line in logs.split("\n"): +# logging.info(line) +# +# if f"completed step: {step}" in logs: +# elapsed_time = time.monotonic() - start_time +# logging.info( +# f"Stop training at step {step}. Sensor completed successfully in {elapsed_time:.2f} seconds." +# ) +# return True +# time.sleep(polling_time) +# except Exception as e: +# elapsed_time = time.monotonic() - start_time +# logging.error( +# f"An unexpected error occurred: {e}. Sensor failed after {elapsed_time:.2f} seconds." +# ) +# raise From 8b63181c2a6e727b9e8fe0881fa814d76db76ea1 Mon Sep 17 00:00:00 2001 From: Ernie Chang Date: Tue, 29 Jul 2025 08:29:16 +0000 Subject: [PATCH 2/5] Update format --- .../axlearn_emergency_checkpoint_save_gcs.py | 1 - xlml/apis/task.py | 20 +++---- xlml/utils/axlearn.py | 57 ++++++++++++------- 3 files changed, 44 insertions(+), 34 deletions(-) diff --git a/dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py b/dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py index 176b38214..85abd3f96 100644 --- a/dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py +++ b/dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py @@ -99,4 +99,3 @@ module=v5p_conf['module'], model_config=v5p_conf['model_config'], ) - diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 730aa4c62..245d0d419 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -228,7 +228,6 @@ def run_model( model_config, ) - ((workload_id, gcs_path) >> launch_workload) return group, gcs_path @@ -269,7 +268,7 @@ def launch_workload( trainer_dir=f"gs://{self.task_gcp_config.project_name}-axlearn/{self.task_test_config.test_name}-nr-{self.task_test_config.num_slices}", num_replicas=self.task_test_config.num_slices, axlearn_branch=axlearn_branch, - trace_steps=[40, 90, 140, 190, 240] + trace_steps=[40, 90, 140, 190, 240], ) wait_for_workload_start = xpk.wait_for_workload_start.override( @@ -282,11 +281,11 @@ def launch_workload( ) ( - setup_axlearn_dpd - >> create_conf_axlearn - >> activate_axlearn - >> run_workload - >> wait_for_workload_start + setup_axlearn_dpd + >> create_conf_axlearn + >> activate_axlearn + >> run_workload + >> wait_for_workload_start ) return group @@ -955,12 +954,7 @@ def run_with_existing_instance(self) -> DAGNode: def provision_via_existing_instance( self, - ) -> Tuple[ - DAGNode, - airflow.XComArg, - airflow.XComArg, - airflow.XComArg, - ]: + ) -> Tuple[DAGNode, airflow.XComArg, airflow.XComArg, airflow.XComArg,]: """Provision an existing GPU accelerator. Returns: diff --git a/xlml/utils/axlearn.py b/xlml/utils/axlearn.py index bbf2d85a8..cae3e56c1 100644 --- a/xlml/utils/axlearn.py +++ b/xlml/utils/axlearn.py @@ -13,12 +13,12 @@ """Utilities to run workloads with xpk (https://github.com/AI-Hypercomputer/xpk).""" +from datetime import timedelta +import re import uuid from absl import logging from airflow.decorators import task from airflow.hooks.subprocess import SubprocessHook -import re -from datetime import timedelta LALITAH_BRANCH = "lkolluru-orbax-fuji-v2" SAM_BRANCH = "orbax-fuji-v2" @@ -28,6 +28,7 @@ # One of them is deleting some unuseful packages in [dev] dependencies. We only need to run axlearn CLI @task(execution_timeout=timedelta(hours=1)) def set_up_axlearn_dpd(branch: str): + """setup axlearn dependencies.""" if branch == LALITAH_BRANCH or branch == SAM_BRANCH: logging.info(f"Using custom branch ==> {branch}") clone_branch = ( @@ -50,7 +51,7 @@ def set_up_axlearn_dpd(branch: str): f"source ~/.profile", f"pyenv install 3.10.12 && pyenv global 3.10.12", "python -m venv ~/my_venv", - f"source ~/my_venv/bin/activate" + f"source ~/my_venv/bin/activate", ] # TODO: Need to think a better way to do this. @@ -95,7 +96,8 @@ def activate_axlearn( "python --version", "which axlearn", "axlearn gcp config activate", - f"gcloud container clusters get-credentials {cluster_name} --region {region} --project {project_id}", + f"gcloud container clusters get-credentials {cluster_name} \ + --region {region} --project {project_id}", ] hook = SubprocessHook() result = hook.run_command(["bash", "-c", ";".join(cmds)]) @@ -122,7 +124,7 @@ def create_conf_axlearn( zone: str, ): """create axlearn config file.""" - command_string = f"cat << 'CONFIG_EOF' > ~/axlearn/.axlearn/axlearn.default.config\n [gcp]\n_active = \"{project_id}:{zone}\"\n\n[gcp.\"{project_id}:{zone}\"]\nproject = \"{project_id}\"\nregion = \"{zone[:-2]}\"\nzone = \"{zone}\"\ngke_cluster = \"{cluster_name}\"\ncluster = \"{cluster_name}\"\nlabels = \"tpu-v5p\"\ndocker_repo = \"us-docker.pkg.dev/{project_id}/axlearn\"\ndefault_dockerfile = \"Dockerfile\"\npermanent_bucket = \"{project_id}-axlearn\"\nprivate_bucket = \"{project_id}-axlearn\"\nttl_bucket = \"{project_id}-axlearn\"\nCONFIG_EOF\n" + command_string = f'cat << \'CONFIG_EOF\' > ~/axlearn/.axlearn/axlearn.default.config\n [gcp]\n_active = "{project_id}:{zone}"\n\n[gcp."{project_id}:{zone}"]\nproject = "{project_id}"\nregion = "{zone[:-2]}"\nzone = "{zone}"\ngke_cluster = "{cluster_name}"\ncluster = "{cluster_name}"\nlabels = "tpu-v5p"\ndocker_repo = "us-docker.pkg.dev/{project_id}/axlearn"\ndefault_dockerfile = "Dockerfile"\npermanent_bucket = "{project_id}-axlearn"\nprivate_bucket = "{project_id}-axlearn"\nttl_bucket = "{project_id}-axlearn"\nCONFIG_EOF\n' # command_string = """cat << 'CONFIG_EOF' > ~/axlearn/.axlearn/axlearn.default.config # ["gcp.cienet-cmcs:us-east5-a"] # project = "cienet-cmcs" @@ -136,14 +138,10 @@ def create_conf_axlearn( # permanent_bucket = "cienet-cmcs-axlearn" # private_bucket = "cienet-cmcs-axlearn" # ttl_bucket = "cienet-cmcs-axlearn"CONFIG_EOF""" - create_axlearn_conf = [command_string.rstrip('\n')] - cmds = [ - *create_axlearn_conf - ] + create_axlearn_conf = [command_string.rstrip("\n")] + cmds = [*create_axlearn_conf] hook = SubprocessHook() - result = hook.run_command( - ["bash", "-c",";".join(cmds)] - ) + result = hook.run_command(["bash", "-c", ";".join(cmds)]) assert ( result.exit_code == 0 @@ -180,16 +178,34 @@ def run_workload_axlearn( f"export NUM_REPLICAS={num_replicas}", f"export MODULE={module}", f"export MODEL_CONFIG={model_config}", - f"export TRAIN_DIR={trainer_dir}" + f"export TRAIN_DIR={trainer_dir}", ] - logging.info(f" Cluster: {cluster_name} -- num-replicas={num_replicas} --run_name={run_name} --project={cluster_project} --zone={zone} --instance-type={accelerator_type} --module={module} --config={model_config} --trainer_dir={trainer_dir} --data_dir=gs://axlearn-public/tensorflow_datasets --jax_backend=tpu --mesh_selector={accelerator_type} --initialization_timeout=1200 Trace: {trace_list}") + logging.info(f" Cluster: {cluster_name} \ + -- num-replicas={num_replicas} \ + --run_name={run_name} \ + --project={cluster_project} \ + --zone={zone} \ + --instance-type={accelerator_type} \ + --module={module} \ + --config={model_config} \ + --trainer_dir={trainer_dir} \ + --data_dir=gs://axlearn-public/tensorflow_datasets \ + --jax_backend=tpu \ + --mesh_selector={accelerator_type} \ + --initialization_timeout=1200 Trace: {trace_list}") workload_create_cmd = ( - f"axlearn gcp launch run --cluster=$CLUSTER --runner_name gke_tpu_single " - f" --name=$NAME --instance_type=$INSTANCE_TYPE --num_replicas=$NUM_REPLICAS --bundler_spec=allow_dirty=True " - f"--bundler_type=artifactregistry --bundler_spec=image=tpu --bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu " + f"axlearn gcp launch run --cluster=$CLUSTER " + f"--runner_name gke_tpu_single " + f" --name=$NAME --instance_type=$INSTANCE_TYPE " + f"--num_replicas=$NUM_REPLICAS --bundler_spec=allow_dirty=True " + f"--bundler_type=artifactregistry --bundler_spec=image=tpu " + f"--bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu " f"-- \"ulimit -n 1048576; ulimit -c 0; python3 -c 'import jax; jax.devices()'; python3 -m axlearn.common.launch_trainer_main\" " - f"--module=$MODULE --config=$MODEL_CONFIG --trainer_dir=$TRAIN_DIR " - f"--data_dir=gs://axlearn-public/tensorflow_datasets --jax_backend=tpu --mesh_selector=$INSTANCE_TYPE --initialization_timeout=1200 {trace_list} " + f"--module=$MODULE --config=$MODEL_CONFIG " + f"--trainer_dir=$TRAIN_DIR " + f"--data_dir=gs://axlearn-public/tensorflow_datasets " + f"--jax_backend=tpu --mesh_selector=$INSTANCE_TYPE " + f"--initialization_timeout=1200 {trace_list} " ) cmds = [ @@ -197,7 +213,8 @@ def run_workload_axlearn( "source ~/my_venv/bin/activate", "cd ~/axlearn", "axlearn gcp config activate", - f"gcloud container clusters get-credentials {cluster_name} --region {zone[:-2]} --project {cluster_project}", + f"gcloud container clusters get-credentials {cluster_name} \ + --region {zone[:-2]} --project {cluster_project}", *export_var, workload_create_cmd, ] From 962844ffc8f846b33d993f9fc78ae0f78bc25741 Mon Sep 17 00:00:00 2001 From: Ernie Chang Date: Tue, 29 Jul 2025 08:39:08 +0000 Subject: [PATCH 3/5] Update format --- xlml/apis/task.py | 278 ------------------------------------------ xlml/utils/axlearn.py | 28 +++-- 2 files changed, 15 insertions(+), 291 deletions(-) diff --git a/xlml/apis/task.py b/xlml/apis/task.py index 245d0d419..db81137b9 100644 --- a/xlml/apis/task.py +++ b/xlml/apis/task.py @@ -346,44 +346,6 @@ def run( return group - def run_with_workload_id( - 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, - workload_id: str = None, - ) -> DAGNode: - """Run a test job within a docker image with specific workload id. - - Attributes: - gcs_location: GCS path for all artifacts of the test. - use_vertex_tensorboard: Set to True to view workload data on - Vertex AI Tensorboard. - - Returns: - A task group with the following tasks chained: run_model and - post_process. - """ - with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: - run_model, gcs_path = self.run_model_with_workload_id( - gcs_location, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - workload_id, - ) - if not skip_post_process: - run_model >> self.post_process(gcs_path) - - return group - def run_with_name_gen_and_quarantine( self, quarantine_task_group, @@ -475,43 +437,6 @@ def run_with_run_name_generation( ) return group - def run_with_node_interruption( - self, - *, - gcs_location: Optional[airflow.XComArg] = None, - use_vertex_tensorboard: bool = False, - use_pathways: bool = False, - skip_post_process: bool = False, - ramdisk_directory: str = "", - mtc_enabled: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - last_node: bool = False, - ) -> DAGNode: - """Run a test job within a docker image. - - Attributes: - gcs_location: GCS path for all artifacts of the test. - use_vertex_tensorboard: Set to True to view workload data on - Vertex AI Tensorboard. - - Returns: - A task group with the following tasks chained: run_model and - post_process. - """ - with TaskGroup(group_id=self.task_test_config.benchmark_id) as group: - run_model, gcs_path = self.run_model_with_node_interruption( - gcs_location, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - last_node, - ) - if not skip_post_process: - run_model >> self.post_process(gcs_path) - return group - def run_model( self, gcs_location: Optional[airflow.XComArg] = None, @@ -574,141 +499,6 @@ def run_model( ) return group, gcs_path - def run_model_with_node_interruption( - self, - gcs_location: Optional[airflow.XComArg] = None, - use_vertex_tensorboard: bool = False, - use_pathways: bool = False, - ramdisk_directory: str = "", - mtc_enabled: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - last_node: bool = False, - ) -> DAGNode: - """Run the TPU/GPU test in `task_test_config` using xpk. - Different behaviour 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. - - Returns: - A DAG node that executes the model test. - """ - with TaskGroup(group_id="run_model") as group: - workload_id = xpk.generate_workload_id(self.task_test_config.benchmark_id) - if gcs_location: - gcs_path = gcs_location - else: - gcs_path = name_format.generate_gcs_folder_location( - self.task_test_config.gcs_subfolder, - self.task_test_config.benchmark_id, - ) - - launch_workload_with_interruption = ( - self.launch_workload_with_node_interruption( - workload_id, - gcs_path, - use_vertex_tensorboard, - use_pathways, - ramdisk_directory, - mtc_enabled, - xpk_branch, - 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=self.task_gcp_config.zone[:-2], - cluster_name=self.task_test_config.cluster_name, - ) - - clean_up_workload = xpk.clean_up_workload( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - zone=self.task_gcp_config.zone, - cluster_name=self.task_test_config.cluster_name, - xpk_branch=xpk_branch, - ) - - ( - (workload_id, gcs_path) - >> launch_workload_with_interruption - >> wait_for_workload_completion - >> clean_up_workload - ) - return group, gcs_path - - def run_model_with_workload_id( - self, - gcs_location: Optional[airflow.XComArg] = None, - use_vertex_tensorboard: bool = False, - use_pathways: bool = False, - ramdisk_directory: str = "", - mtc_enabled: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - workload_id: str = None, - ) -> DAGNode: - """Run the TPU/GPU test in `task_test_config` using xpk - with specific workload id. - - 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: - if workload_id is None: - 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, - ) - wait_for_workload_completion = xpk.wait_for_workload_completion.override( - timeout=int(self.task_test_config.timeout.total_seconds()), - )( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - region=self.task_gcp_config.zone[:-2], - cluster_name=self.task_test_config.cluster_name, - ) - - clean_up_workload = xpk.clean_up_workload( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - zone=self.task_gcp_config.zone, - cluster_name=self.task_test_config.cluster_name, - ) - - ( - (workload_id, gcs_path) - >> launch_workload - >> wait_for_workload_completion - >> clean_up_workload - ) - return group, gcs_path - def launch_workload( self, workload_id: str, @@ -752,74 +542,6 @@ def launch_workload( run_workload >> wait_for_workload_start return group - def launch_workload_with_node_interruption( - self, - workload_id: str, - gcs_path: str, - use_vertex_tensorboard: bool, - use_pathways: bool = False, - ramdisk_directory: str = "", - mtc_enabled: bool = False, - xpk_branch: str = xpk.MAIN_BRANCH, - last_node: bool = False, - ) -> DAGNode: - """Create the workload and wait for it to provision.""" - with TaskGroup(group_id="launch_workload_with_interruption") as group: - run_workload = xpk.run_workload.override( - owner=self.task_test_config.task_owner - )( - task_id="run_workload", - cluster_project=self.task_gcp_config.project_name, - zone=self.task_gcp_config.zone, - cluster_name=self.task_test_config.cluster_name, - benchmark_id=self.task_test_config.benchmark_id, - workload_id=workload_id, - gcs_path=gcs_path, - docker_image=self.task_test_config.docker_image, - accelerator_type=self.task_test_config.accelerator.name, - run_cmds=self.task_test_config.test_script, - num_slices=self.task_test_config.num_slices, - use_vertex_tensorboard=use_vertex_tensorboard, - use_pathways=use_pathways, - ramdisk_directory=ramdisk_directory, - mtc_enabled=mtc_enabled, - xpk_branch=xpk_branch, - ) - wait_for_workload_start = xpk.wait_for_workload_start.override( - timeout=self.workload_provision_timeout.total_seconds() - )( - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - region=self.task_gcp_config.zone[:-2], - cluster_name=self.task_test_config.cluster_name, - ) - wait_to_reach_step_to_interrupt = xpk.wait_for_reach_step_to_interrupt( - task_id="wait_to_reach_step_to_interrupt", - workload_id=workload_id, - project_id=self.task_gcp_config.project_name, - region=self.task_gcp_config.zone[:-2], - cluster_name=self.task_test_config.cluster_name, - step_to_interrupt="40", - ) - run_node_interruption = xpk.delete_node.override( - owner=self.task_test_config.task_owner - )( - 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, - ) - - ( - run_workload - >> wait_for_workload_start - >> wait_to_reach_step_to_interrupt - >> run_node_interruption - ) - return group - def post_process(self, result_location: Optional[str] = None) -> DAGNode: """Process metrics and metadata, and insert them into BigQuery tables. diff --git a/xlml/utils/axlearn.py b/xlml/utils/axlearn.py index cae3e56c1..f72591eb4 100644 --- a/xlml/utils/axlearn.py +++ b/xlml/utils/axlearn.py @@ -180,19 +180,21 @@ def run_workload_axlearn( f"export MODEL_CONFIG={model_config}", f"export TRAIN_DIR={trainer_dir}", ] - logging.info(f" Cluster: {cluster_name} \ - -- num-replicas={num_replicas} \ - --run_name={run_name} \ - --project={cluster_project} \ - --zone={zone} \ - --instance-type={accelerator_type} \ - --module={module} \ - --config={model_config} \ - --trainer_dir={trainer_dir} \ - --data_dir=gs://axlearn-public/tensorflow_datasets \ - --jax_backend=tpu \ - --mesh_selector={accelerator_type} \ - --initialization_timeout=1200 Trace: {trace_list}") + logging.info( + f" Cluster: {cluster_name} \ + -- num-replicas={num_replicas} \ + --run_name={run_name} \ + --project={cluster_project} \ + --zone={zone} \ + --instance-type={accelerator_type} \ + --module={module} \ + --config={model_config} \ + --trainer_dir={trainer_dir} \ + --data_dir=gs://axlearn-public/tensorflow_datasets \ + --jax_backend=tpu \ + --mesh_selector={accelerator_type} \ + --initialization_timeout=1200 Trace: {trace_list}" + ) workload_create_cmd = ( f"axlearn gcp launch run --cluster=$CLUSTER " f"--runner_name gke_tpu_single " From d1c81f10a15847c3c0ecf0cbc928afb37d29bfa4 Mon Sep 17 00:00:00 2001 From: Ernie Chang Date: Tue, 29 Jul 2025 08:52:29 +0000 Subject: [PATCH 4/5] Update format --- xlml/utils/axlearn.py | 53 ++++++++++++++++++++++++------------------- 1 file changed, 30 insertions(+), 23 deletions(-) diff --git a/xlml/utils/axlearn.py b/xlml/utils/axlearn.py index f72591eb4..93b09d4f6 100644 --- a/xlml/utils/axlearn.py +++ b/xlml/utils/axlearn.py @@ -11,7 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Utilities to run workloads with xpk (https://github.com/AI-Hypercomputer/xpk).""" +"""Utilities to run workloads with axlearn.""" from datetime import timedelta import re @@ -25,7 +25,8 @@ # This function do some hacks to get Axlearn working with Airlfow -# One of them is deleting some unuseful packages in [dev] dependencies. We only need to run axlearn CLI +# One of them is deleting some unuseful packages in [dev] dependencies. +# We only need to run axlearn CLI @task(execution_timeout=timedelta(hours=1)) def set_up_axlearn_dpd(branch: str): """setup axlearn dependencies.""" @@ -125,7 +126,8 @@ def create_conf_axlearn( ): """create axlearn config file.""" command_string = f'cat << \'CONFIG_EOF\' > ~/axlearn/.axlearn/axlearn.default.config\n [gcp]\n_active = "{project_id}:{zone}"\n\n[gcp."{project_id}:{zone}"]\nproject = "{project_id}"\nregion = "{zone[:-2]}"\nzone = "{zone}"\ngke_cluster = "{cluster_name}"\ncluster = "{cluster_name}"\nlabels = "tpu-v5p"\ndocker_repo = "us-docker.pkg.dev/{project_id}/axlearn"\ndefault_dockerfile = "Dockerfile"\npermanent_bucket = "{project_id}-axlearn"\nprivate_bucket = "{project_id}-axlearn"\nttl_bucket = "{project_id}-axlearn"\nCONFIG_EOF\n' - # command_string = """cat << 'CONFIG_EOF' > ~/axlearn/.axlearn/axlearn.default.config + # command_string = """cat << 'CONFIG_EOF' > + # ~/axlearn/.axlearn/axlearn.default.config # ["gcp.cienet-cmcs:us-east5-a"] # project = "cienet-cmcs" # region = "us-east5" @@ -181,19 +183,19 @@ def run_workload_axlearn( f"export TRAIN_DIR={trainer_dir}", ] logging.info( - f" Cluster: {cluster_name} \ - -- num-replicas={num_replicas} \ - --run_name={run_name} \ - --project={cluster_project} \ - --zone={zone} \ - --instance-type={accelerator_type} \ - --module={module} \ - --config={model_config} \ - --trainer_dir={trainer_dir} \ - --data_dir=gs://axlearn-public/tensorflow_datasets \ - --jax_backend=tpu \ - --mesh_selector={accelerator_type} \ - --initialization_timeout=1200 Trace: {trace_list}" + f" Cluster: {cluster_name} \ + -- num-replicas={num_replicas} \ + --run_name={run_name} \ + --project={cluster_project} \ + --zone={zone} \ + --instance-type={accelerator_type} \ + --module={module} \ + --config={model_config} \ + --trainer_dir={trainer_dir} \ + --data_dir=gs://axlearn-public/tensorflow_datasets \ + --jax_backend=tpu \ + --mesh_selector={accelerator_type} \ + --initialization_timeout=1200 Trace: {trace_list}" ) workload_create_cmd = ( f"axlearn gcp launch run --cluster=$CLUSTER " @@ -385,7 +387,8 @@ def run_workload_axlearn( # for line in logs.split("\n"): # logging.info(line) # if f"completed step: {step_to_interrupt}" in logs: -# # Here we return true because we are sure the step "step_to_interrupt" is already save +# # Here we return true because we are sure the step +# "step_to_interrupt" is already save # logging.info("The step to be interrupt is {step_to_interrupt}") # return True # return False @@ -402,8 +405,8 @@ def run_workload_axlearn( # if not pods.items: # logging.info(f"No pods found for workload selector: {workload_id}.") # -# # Pathways jobs delete all pods on failure so we must also check if the job -# # is complete +# # Pathways jobs delete all pods on failure so we must also check +# if the job is complete # batch_api = _get_batch_api_client(project_id, region, cluster_name) # job = _get_workload_job(batch_api, workload_id) # if job is None: @@ -416,7 +419,8 @@ def run_workload_axlearn( # # Don't keep retrying if the job has failed # raise AirflowFailException('Job has condition type: "Failed"') # -# if any(condition.type == "Complete" for condition in job.status.conditions): +# if any(condition.type == "Complete" +# for condition in job.status.conditions): # logging.info( # "No pods found but job is complete for workload selector:" # f" {workload_id}" @@ -439,7 +443,8 @@ def run_workload_axlearn( # finally: # # TODO(jonbolin): log printing for GPUs, which have multiple containers # if len(pod.spec.containers) == 1: -# # Print the logs of the last pod checked - either the first failed pod or +# # Print the logs of the last pod checked +# - either the first failed pod or # # the last successful one. # logs = core_api.read_namespaced_pod_log( # name=pod.metadata.name, namespace=pod.metadata.namespace @@ -596,13 +601,15 @@ def run_workload_axlearn( # if f"completed step: {step}" in logs: # elapsed_time = time.monotonic() - start_time # logging.info( -# f"Stop training at step {step}. Sensor completed successfully in {elapsed_time:.2f} seconds." +# f"Stop training at step {step}. +# Sensor completed successfully in {elapsed_time:.2f} seconds." # ) # return True # time.sleep(polling_time) # except Exception as e: # elapsed_time = time.monotonic() - start_time # logging.error( -# f"An unexpected error occurred: {e}. Sensor failed after {elapsed_time:.2f} seconds." +# f"An unexpected error occurred: {e}. +# Sensor failed after {elapsed_time:.2f} seconds." # ) # raise From 63b1801984f8fc30f3a0ec85afbb4734719c7869 Mon Sep 17 00:00:00 2001 From: Ernie Chang Date: Wed, 30 Jul 2025 08:04:36 +0000 Subject: [PATCH 5/5] Update format --- xlml/utils/axlearn.py | 68 +++++++++++++++++++++---------------------- 1 file changed, 33 insertions(+), 35 deletions(-) diff --git a/xlml/utils/axlearn.py b/xlml/utils/axlearn.py index 93b09d4f6..cbef71fd5 100644 --- a/xlml/utils/axlearn.py +++ b/xlml/utils/axlearn.py @@ -11,7 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Utilities to run workloads with axlearn.""" +"""Utilities to run workloads with AXLearn.""" from datetime import timedelta import re @@ -27,15 +27,13 @@ # This function do some hacks to get Axlearn working with Airlfow # One of them is deleting some unuseful packages in [dev] dependencies. # We only need to run axlearn CLI -@task(execution_timeout=timedelta(hours=1)) +@task def set_up_axlearn_dpd(branch: str): """setup axlearn dependencies.""" - if branch == LALITAH_BRANCH or branch == SAM_BRANCH: - logging.info(f"Using custom branch ==> {branch}") - clone_branch = ( - f"git clone --branch {branch} https://github.com/lkolluru05/axlearn.git" - f" $HOME/axlearn" - ) + logging.info(f"Using custom branch ==> {branch}") + clone_branch = ( + f"git clone --branch {branch} https://github.com/lkolluru05/axlearn.git $HOME/axlearn" + ) # Maybe add these lines install_python3_cmd = [ @@ -182,34 +180,34 @@ def run_workload_axlearn( f"export MODEL_CONFIG={model_config}", f"export TRAIN_DIR={trainer_dir}", ] - logging.info( - f" Cluster: {cluster_name} \ - -- num-replicas={num_replicas} \ - --run_name={run_name} \ - --project={cluster_project} \ - --zone={zone} \ - --instance-type={accelerator_type} \ - --module={module} \ - --config={model_config} \ - --trainer_dir={trainer_dir} \ - --data_dir=gs://axlearn-public/tensorflow_datasets \ - --jax_backend=tpu \ - --mesh_selector={accelerator_type} \ - --initialization_timeout=1200 Trace: {trace_list}" - ) + # logging.info( + # f" Cluster: {cluster_name} \ + # -- num-replicas={num_replicas} \ + # --run_name={run_name} \ + # --project={cluster_project} \ + # --zone={zone} \ + # --instance-type={accelerator_type} \ + # --module={module} \ + # --config={model_config} \ + # --trainer_dir={trainer_dir} \ + # --data_dir=gs://axlearn-public/tensorflow_datasets \ + # --jax_backend=tpu \ + # --mesh_selector={accelerator_type} \ + # --initialization_timeout=1200 Trace: {trace_list}" + # ) workload_create_cmd = ( - f"axlearn gcp launch run --cluster=$CLUSTER " - f"--runner_name gke_tpu_single " - f" --name=$NAME --instance_type=$INSTANCE_TYPE " - f"--num_replicas=$NUM_REPLICAS --bundler_spec=allow_dirty=True " - f"--bundler_type=artifactregistry --bundler_spec=image=tpu " - f"--bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu " - f"-- \"ulimit -n 1048576; ulimit -c 0; python3 -c 'import jax; jax.devices()'; python3 -m axlearn.common.launch_trainer_main\" " - f"--module=$MODULE --config=$MODEL_CONFIG " - f"--trainer_dir=$TRAIN_DIR " - f"--data_dir=gs://axlearn-public/tensorflow_datasets " - f"--jax_backend=tpu --mesh_selector=$INSTANCE_TYPE " - f"--initialization_timeout=1200 {trace_list} " + f"axlearn gcp launch run --cluster=$CLUSTER " + f"--runner_name gke_tpu_single " + f" --name=$NAME --instance_type=$INSTANCE_TYPE " + f"--num_replicas=$NUM_REPLICAS --bundler_spec=allow_dirty=True " + f"--bundler_type=artifactregistry --bundler_spec=image=tpu " + f"--bundler_spec=dockerfile=Dockerfile --bundler_spec=target=tpu " + f"-- \"ulimit -n 1048576; ulimit -c 0; python3 -c 'import jax; jax.devices()'; python3 -m axlearn.common.launch_trainer_main\" " + f"--module=$MODULE --config=$MODEL_CONFIG " + f"--trainer_dir=$TRAIN_DIR " + f"--data_dir=gs://axlearn-public/tensorflow_datasets " + f"--jax_backend=tpu --mesh_selector=$INSTANCE_TYPE " + f"--initialization_timeout=1200 {trace_list} " ) cmds = [