add a new DAG 'axlearn_emergency_save'

dags/axlearn/axlearn_emergency_checkpoint_save_gcs.py

@@ -0,0 +1,101 @@
+# 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'],
+          )

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

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

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,151 @@ 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.
@@ -343,6 +488,7 @@ def run_model(
           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,
       )
 
       (