From 4a5fd5dc48d541fc68c3d6efe657f85ed1e411ac Mon Sep 17 00:00:00 2001
From: Andrew S <andrewsal@google.com>
Date: Tue, 18 Mar 2025 11:54:55 -0700
Subject: [PATCH] Updated TPU unittests to cover more TPU types / JAX versions

- Expanded test to v5e, v5p and JAX 0.5.1 and nightly
- Allow for different projects/networks for test VMs
- Merges tests into existing project_bite_tpu_e2e DAG to minimize number of DAG entries
---
 .../sparsity_diffusion_devx/configs/common.py |  14 +++
 .../configs/project_bite_config.py            |  67 ++++++++---
 .../project_bite_tpu_e2e.py                   | 109 ++++++++++++------
 3 files changed, 139 insertions(+), 51 deletions(-)

diff --git a/dags/sparsity_diffusion_devx/configs/common.py b/dags/sparsity_diffusion_devx/configs/common.py
index 0f273c101..c9bc36523 100644
--- a/dags/sparsity_diffusion_devx/configs/common.py
+++ b/dags/sparsity_diffusion_devx/configs/common.py
@@ -32,3 +32,17 @@ def set_up_nightly_jax() -> Tuple[str]:
       ),
       "pip install git+https://github.com/google/jax",
   )
+
+
+def set_up_jax_version(version) -> Tuple[str]:
+  return (
+      (
+          f"pip install jax[tpu]=={version}  -f "
+          "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
+      ),
+      (
+          f"pip install --pre jaxlib=={version} -f"
+          " https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html"
+      ),
+      f"pip install jax=={version}",
+  )
diff --git a/dags/sparsity_diffusion_devx/configs/project_bite_config.py b/dags/sparsity_diffusion_devx/configs/project_bite_config.py
index 0573b8ddb..9c2072f2d 100644
--- a/dags/sparsity_diffusion_devx/configs/project_bite_config.py
+++ b/dags/sparsity_diffusion_devx/configs/project_bite_config.py
@@ -94,18 +94,21 @@ def get_bite_tpu_config(
   )
 
 
-def get_bite_tpu_unittests_config(
-    tpu_version: TpuVersion,
-    tpu_cores: int,
-    tpu_zone: str,
-    runtime_version: str,
-    time_out_in_min: int,
-    task_owner: str,
-    is_tpu_reserved: bool = False,
-    pinned_version: Optional[str] = None,
-):
-  unittest_setupcmds = (
-      # create configuration files needed
+def dockerfile_build_cmd(jax_version):
+  # Generate pip commands to install certain version of JAX/libTPU e.g.
+  # pip install --pre jaxlib==0.5.1  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+  # pip install jax[tpu]==0.5.1  -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+  # pip install jax==0.5.1
+  if jax_version:
+    pip_tpu_jax_install = "\n".join(
+        ["RUN " + x for x in common.set_up_jax_version(jax_version)]
+    )
+  else:
+    pip_tpu_jax_install = "\n".join(
+        ["RUN " + x for x in common.set_up_nightly_jax()]
+    )
+
+  return (
       """cat > Dockerfile_CI <<EOF
 FROM python:3.10-slim
 WORKDIR /workspace
@@ -118,12 +121,33 @@ def get_bite_tpu_unittests_config(
 RUN pip install -e '.[core,dev,gcp]'
 RUN pip install grain
 RUN pip install google-cloud-aiplatform
-RUN pip install -U --pre libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-RUN pip install --pre -U jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
-RUN pip install git+https://github.com/google/jax
+"""
+      + pip_tpu_jax_install
+      + """
 RUN pip freeze
 EOF
-""",
+"""
+  )
+
+
+def get_bite_tpu_unittests_config(
+    tpu_version: TpuVersion,
+    tpu_cores: int,
+    tpu_zone: str,
+    runtime_version: str,
+    time_out_in_min: int,
+    task_owner: str,
+    network: str = "default",
+    subnetwork: str = "default",
+    is_tpu_reserved: bool = False,
+    jax_version: Optional[str] = None,
+    test_suffix: Optional[str] = None,
+    project_name: Optional[Project] = Project.CLOUD_ML_AUTO_SOLUTIONS.value,
+):
+
+  unittest_setupcmds = (
+      # create configuration files needed
+      dockerfile_build_cmd(jax_version),
       # create script to run the tests inside of the container
       # incluedes a basic sanity check python script which prints out TPU env info for reference
       """cat > run_tpu_tests.sh <<EOF
@@ -144,19 +168,26 @@ def get_bite_tpu_unittests_config(
       "sudo docker run --network=host --privileged --ulimit memlock=-1:-1  ml-auto-solutions/tpu_unittests  /bin/bash -c '/workspace/run_tpu_tests.sh'",
   )
   job_gcp_config = gcp_config.GCPConfig(
-      project_name=Project.CLOUD_ML_AUTO_SOLUTIONS.value,
+      project_name=project_name,
       zone=tpu_zone,
       dataset_name=metric_config.DatasetOption.XLML_DATASET,
   )
 
+  if test_suffix:
+    test_name = f"bite_unittests_{test_suffix}"
+  else:
+    test_name = "bite_unittests"
+
   tpu_unittests_test_config = test_config.TpuVmTest(
       test_config.Tpu(
           version=tpu_version,
           cores=tpu_cores,
           runtime_version=runtime_version,
           reserved=is_tpu_reserved,
+          network=network,
+          subnetwork=subnetwork,
       ),
-      test_name="bite_unittests",
+      test_name=test_name,
       set_up_cmds=unittest_setupcmds,
       run_model_cmds=unittest_runcmds,
       timeout=datetime.timedelta(minutes=time_out_in_min),
diff --git a/dags/sparsity_diffusion_devx/project_bite_tpu_e2e.py b/dags/sparsity_diffusion_devx/project_bite_tpu_e2e.py
index f0667bc3b..4afd509f3 100644
--- a/dags/sparsity_diffusion_devx/project_bite_tpu_e2e.py
+++ b/dags/sparsity_diffusion_devx/project_bite_tpu_e2e.py
@@ -18,23 +18,57 @@
 from airflow import models
 from dags import composer_env
 from dags.common import test_owner
-from dags.common.vm_resource import TpuVersion, Zone, RuntimeVersion
+from dags.common.vm_resource import TpuVersion, Zone, RuntimeVersion, Project
 from dags.sparsity_diffusion_devx.configs import project_bite_config as config
 
 
 # Run once a day at 6 pm UTC (11 am PST)
-SCHEDULED_TIME = "0 18 * * *" if composer_env.is_prod_env() else None
+SCHEDULED_TIME = '0 18 * * *' if composer_env.is_prod_env() else None
+
+trillium_conf = {
+    'tpu_version': TpuVersion.TRILLIUM,
+    'tpu_cores': 4,
+    'tpu_zone': Zone.EUROPE_WEST4_A.value,
+    'runtime_version': RuntimeVersion.V2_ALPHA_TPUV6.value,
+}
+
+v5p_conf = {
+    'tpu_version': TpuVersion.V5P,
+    'tpu_cores': 8,
+    'tpu_zone': Zone.US_EAST5_A.value,
+    'is_tpu_reserved': True,
+    'runtime_version': RuntimeVersion.V2_ALPHA_TPUV5.value,
+    'project_name': Project.TPU_PROD_ENV_AUTOMATED.value,
+    'network': 'mas-test',
+    'subnetwork': 'mas-test',
+}
+
+v5e_conf = {
+    'tpu_version': TpuVersion.V5E,
+    'tpu_cores': 8,
+    'tpu_zone': Zone.US_EAST1_C.value,
+    'is_tpu_reserved': True,
+    'runtime_version': RuntimeVersion.V2_ALPHA_TPUV5_LITE.value,
+    'project_name': Project.TPU_PROD_ENV_AUTOMATED.value,
+    'network': 'mas-test',
+    'subnetwork': 'mas-test',
+}
+
+common = {
+    'time_out_in_min': 180,
+    'task_owner': test_owner.Andrew_S,
+}
 
 
 with models.DAG(
-    dag_id="project_bite_tpu_e2e",
+    dag_id='project_bite_tpu_e2e',
     schedule=SCHEDULED_TIME,
     tags=[
-        "sparsity_diffusion_devx",
-        "multipod_team",
-        "tpu",
-        "axlearn",
-        "bite",
+        'sparsity_diffusion_devx',
+        'multipod_team',
+        'tpu',
+        'axlearn',
+        'bite',
     ],
     start_date=datetime.datetime(2024, 4, 4),
     catchup=False,
@@ -65,28 +99,37 @@
       task_owner=test_owner.Maggie_Z,
   )
 
-default_unittest_args = {
-    "retries": 0,
-}
-
-with models.DAG(
-    dag_id="project_bite_tpu_unittests",
-    schedule=SCHEDULED_TIME,
-    tags=[
-        "sparsity_diffusion_devx",
-        "tpu",
-        "axlearn",
-        "bite",
-    ],
-    start_date=datetime.datetime(2025, 2, 24),
-    catchup=False,
-    default_args=default_unittest_args,
-) as dag:
-  unittests = config.get_bite_tpu_unittests_config(
-      tpu_version=TpuVersion.TRILLIUM,
-      tpu_cores=4,
-      tpu_zone=Zone.EUROPE_WEST4_A.value,
-      runtime_version=RuntimeVersion.V2_ALPHA_TPUV6.value,
-      time_out_in_min=180,
-      task_owner=test_owner.Andrew_S,
-  )
+  tpu_unittests = [
+      # Trillium (v6e) with JAX 0.5.1
+      config.get_bite_tpu_unittests_config(
+          **trillium_conf,
+          test_suffix='v6e_051',
+          jax_version='0.5.1',
+          **common,
+      ),
+      # Trillium (v6e) with JAX nightly
+      config.get_bite_tpu_unittests_config(
+          **trillium_conf,
+          test_suffix='v6e_nightly',
+          **common,
+      ),
+      # V5P with JAX 0.5.1
+      config.get_bite_tpu_unittests_config(
+          **v5p_conf,
+          test_suffix='v5p_051',
+          jax_version='0.5.1',
+          **common,
+      ),
+      # V5P with JAX nightly
+      config.get_bite_tpu_unittests_config(
+          **v5p_conf,
+          test_suffix='v5p_nightly',
+          **common,
+      ),
+      # V5E with JAX nightly
+      config.get_bite_tpu_unittests_config(
+          **v5e_conf,
+          test_suffix='v5e_nightly',
+          **common,
+      ),
+  ]