add a new DAG 'axlearn_emergency_save' by ernie-chang · Pull Request #186 · CIeNET-International/ml-auto-solutions-2

ernie-chang · 2025-07-29T08:09:25Z

Description

This test verifies the Axlearn Checkpointing saving function. It conducts the following tasks:

Run the training by enabling the checkpointing, until the local and gcs checkpoints have been saved.

Tests

project: cienet-cmcs
name: ernie-axlearn-v5p
zone: us-east5-a
highScaleCheckpointingConfig: enabled
GcsFuseCsiDriver: enabled
Workload Identity Federation: enabled
gcs bucket:

name: cienet-cmcs-axlearn
Hierarchical namespace: enabled
airflow/composer env:

project: cienet-cmcs
name: ernie-ml-solutions
composer version: 2.10.1
airflow version: 2.10.2

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run one-shot tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed.

alfredyu-cienet · 2025-07-29T09:21:41Z

+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities to run workloads with axlearn."""


AXLearn is a proper noun, please be mindful

alfredyu-cienet · 2025-07-29T09:23:35Z

+# 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))


please apply Airflow task group and organize tasks better

why 1 hour? do we have to specify a timeout for this task? do we need to specify timeout for each task?

alfredyu-cienet · 2025-07-29T09:30:20Z

+@task(execution_timeout=timedelta(hours=1))
+def set_up_axlearn_dpd(branch: str):
+  """setup axlearn dependencies."""
+  if branch == LALITAH_BRANCH or branch == SAM_BRANCH:


edge case?

using arbitrary string is not a good design --> caller can specify any string

alfredyu-cienet · 2025-07-29T09:31:22Z

+        f" $HOME/axlearn"
+    )
+
+  # Maybe add these lines


what is this

alfredyu-cienet · 2025-07-29T09:31:58Z

+  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"


why using lkolluru05's fork? this is requested by Google?

alfredyu-cienet · 2025-07-29T09:52:58Z

+
+
+@task
+def generate_workload_id(benchmark_id: str) -> str:


this is a very simple function and it's likely won't fail, do we need it to be an independent task? can't we just integrate it to other method which generates the configuration for AXLearn?

alfredyu-cienet · 2025-07-29T10:02:07Z

+# 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


isn't this supposed to be the doc string of this function?

alfredyu-cienet · 2025-07-29T10:03:39Z

+      f"gcloud container clusters get-credentials {cluster_name} \
+        --region {region} --project {project_id}",


why's this command is considered as activate AXLearn

alfredyu-cienet · 2025-07-29T10:06:39Z

+      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}     "


please format that properly

alfredyu-cienet · 2025-07-29T10:11:12Z

+  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}"
+  )


shorten the code

what exactly are you trying to log? the parameters passed from caller or the actual command that will be executed? ---> data_dir and jax_backend are fixed value

CJ-Tsai-cienet · 2025-07-31T02:33:56Z

+      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,
+        )


gcs_path = gcs_location if gcs_location else name_format.generate_gcs_folder_location( self.task_test_config.gcs_subfolder, self.task_test_config.benchmark_id, )

CJ-Tsai-cienet · 2025-07-31T07:04:15Z

+  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}"


Set up -> Activate

CJ-Tsai-cienet · 2025-07-31T07:17:59Z

+  """Generate a valid workload ID."""
+
+  # Remove all non-alphanumeric characters, and truncate to ensure the result
+  # is less than 40 characters.


the result could be 40 characters

CJ-Tsai-cienet · 2025-07-31T07:28:38Z

+  create_axlearn_conf = [command_string.rstrip("\n")]
+  cmds = [*create_axlearn_conf]


Do we need create_axlearn_conf and command_string?
Could we put the commands in a list directly instead of unpacking from another list?

cmds = [ "cat << 'CONFIG_EOF' >~/axlearn/.axlearn/axlearn.default.config", ... ]

CJ-Tsai-cienet · 2025-07-31T07:45:07Z

+  trace_list = (
+      ("--trace_at_steps=" + ", ".join(map(str, trace_steps)))
+      if trace_steps
+      else " "
+  )


nit: How about moving to L212 and do

if trace_steps: workload_create_cmd += "--trace_at_steps=" + ", ".join(map(str, trace_steps))

CJ-Tsai-cienet · 2025-07-31T07:46:17Z

+      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 "


none of these need fstring

CJ-Tsai-cienet · 2025-07-31T07:47:38Z

+      f"export BASTION_TIER=disabled",
+      f"export DEFAULT_PROJECT_ID=$(gcloud config get project)",


fstring is not required

CJ-Tsai-cienet · 2025-07-31T07:49:58Z

+
+  assert (
+      result.exit_code == 0
+  ), f"Axlearn command failed with code {result.exit_code}"


Run workload axlearn

CJ-Tsai-cienet · 2025-07-31T07:51:56Z

+  setup_jax = None
+  if jax_version:
+    setup_jax = common.set_up_jax_version(jax_version)
+  else:
+    setup_jax = common.set_up_nightly_jax()


setup_jax = common.set_up_jax_version(jax_version) if jax_version else common.set_up_nightly_jax()

Ernie Chang added 4 commits July 29, 2025 08:08

add a new DAG 'axlearn_emergency_save'

c3a60d0

Update format

8b63181

Update format

962844f

Update format

d1c81f1

alfredyu-cienet requested changes Jul 29, 2025

View reviewed changes

Update format

63b1801

CJ-Tsai-cienet reviewed Jul 31, 2025

View reviewed changes

		f"gcloud container clusters get-credentials {cluster_name} \
		--region {region} --project {project_id}",

		create_axlearn_conf = [command_string.rstrip("\n")]
		cmds = [*create_axlearn_conf]

		f"export BASTION_TIER=disabled",
		f"export DEFAULT_PROJECT_ID=$(gcloud config get project)",

Uh oh!

Conversation

ernie-chang commented Jul 29, 2025 • edited by alfredyu-cienet Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ernie-chang commented Jul 29, 2025 •

edited by alfredyu-cienet

Loading