init v6e mxla nightly (GoogleCloudPlatform#769)

Author: notabee

dags/common/vm_resource.py

@@ -263,6 +263,13 @@ class XpkClusters:
       project=Project.TPU_PROD_ENV_MULTIPOD.value,
       zone=Zone.EUROPE_WEST4_B.value,
   )
+  TPU_V6E_8_CLUSTER = XpkClusterConfig(
+      name="bodaborg-v6e-8-yucmhab-c",
+      device_version=TpuVersion.TRILLIUM,
+      core_count=8,
+      project=Project.TPU_PROD_ENV_ONE_VM.value,
+      zone=Zone.US_EAST5_B.value,
+  )
   TPU_V6E_256_CLUSTER = XpkClusterConfig(
       name="bodaborg-v6e-256",
       device_version=TpuVersion.TRILLIUM,

dags/multipod/mxla_maxtext_nightly_gke.py

@@ -46,6 +46,7 @@
       group_id="Quarantine", dag=dag, prefix_group_id=False
   )
 
+  # v5p tests
   maxtext_nightly_1slice_v5p_8 = gke_config.get_gke_maxtext_nightly_config(
       cluster=XpkClusters.TPU_V5P_8_CLUSTER,
       time_out_in_min=60,
@@ -81,9 +82,54 @@
       test_owner=test_owner.RAYMOND_Z,
   ).run_with_quarantine(quarantine_task_group)
 
+  # v6e tests
+  maxtext_nightly_1slice_v6e_8 = gke_config.get_gke_maxtext_nightly_config(
+      cluster=XpkClusters.TPU_V6E_8_CLUSTER,
+      time_out_in_min=60,
+      test_name=default_test_name,
+      docker_image=jax_nightly_image.value,
+      test_owner=test_owner.RISHABH_B,
+  ).run_with_quarantine(quarantine_task_group)
+
+  maxtext_nightly_2slice_v6e_8 = gke_config.get_gke_maxtext_nightly_config(
+      num_slices=2,
+      cluster=XpkClusters.TPU_V6E_8_CLUSTER,
+      time_out_in_min=60,
+      test_name=default_test_name,
+      docker_image=jax_nightly_image.value,
+      test_owner=test_owner.RISHABH_B,
+  ).run_with_quarantine(quarantine_task_group)
+
+  maxtext_nightly_4slice_v6e_8 = gke_config.get_gke_maxtext_nightly_config(
+      num_slices=4,
+      cluster=XpkClusters.TPU_V6E_8_CLUSTER,
+      time_out_in_min=60,
+      test_name=default_test_name,
+      docker_image=jax_nightly_image.value,
+      test_owner=test_owner.RISHABH_B,
+  ).run_with_quarantine(quarantine_task_group)
+
+  maxtext_nightly_8slice_v6e_8 = gke_config.get_gke_maxtext_nightly_config(
+      num_slices=8,
+      cluster=XpkClusters.TPU_V6E_8_CLUSTER,
+      time_out_in_min=60,
+      test_name=default_test_name,
+      docker_image=jax_nightly_image.value,
+      test_owner=test_owner.RISHABH_B,
+  ).run_with_quarantine(quarantine_task_group)
+
+  # Define dependencies for v5p tests to run sequentially
   (
       maxtext_nightly_1slice_v5p_8
       >> maxtext_nightly_2slice_v5p_8
       >> maxtext_nightly_4slice_v5p_8
       >> maxtext_nightly_8slice_v5p_8
   )
+
+  # Define dependencies for v6e tests to run sequentially
+  (
+      maxtext_nightly_1slice_v6e_8
+      >> maxtext_nightly_2slice_v6e_8
+      >> maxtext_nightly_4slice_v6e_8
+      >> maxtext_nightly_8slice_v6e_8
+  )