Change Gemma test path (#238)

GoogleCloudPlatform · Apr 3, 2024 · 283d0a5 · 283d0a5
1 parent 280081c
commit 283d0a5
Showing 1 changed file with 48 additions and 49 deletions.
diff --git a/dags/multipod/maxtext_end_to_end.py b/dags/multipod/maxtext_end_to_end.py
@@ -35,54 +35,53 @@
 ) as dag:
   test_name_prefix = "maxtext"
   test_models = {
-      "llama2": ["test_llama2_7b"],
-      "mistral": ["test_mistral"],
-      "gemma": ["test_gemma"],
-      "gpt3": ["test_gpt3"],
+      "llama2-7b": "test_llama2_7b",
+      "mistral": "test_mistral",
+      "gemma-2b": "gemma/2b/test_gemma",
+      "gpt3": "test_gpt3",
   }
 
-  for model in test_models.keys():
-    for test_script in test_models[model]:
-      stable_tpu = gke_config.get_gke_config(
-          tpu_version=TpuVersion.V4,
-          tpu_cores=8,
-          tpu_zone=Zone.US_CENTRAL2_B.value,
-          time_out_in_min=60,
-          test_name=f"{test_name_prefix}-stable-{test_script}",
-          run_model_cmds=(f"bash end_to_end/{test_script}.sh",),
-          docker_image=DockerImage.MAXTEXT_TPU_JAX_STABLE.value,
-          test_owner=test_owner.JON_B,
-      ).run()
-      nightly_tpu = gke_config.get_gke_config(
-          tpu_version=TpuVersion.V4,
-          tpu_cores=8,
-          tpu_zone=Zone.US_CENTRAL2_B.value,
-          time_out_in_min=60,
-          test_name=f"{test_name_prefix}-nightly-{test_script}",
-          run_model_cmds=(f"bash end_to_end/{test_script}.sh",),
-          docker_image=DockerImage.MAXTEXT_TPU_JAX_NIGHTLY.value,
-          test_owner=test_owner.JON_B,
-      ).run()
-      stable_gpu = gke_config.get_maxtext_end_to_end_gpu_gke_test_config(
-          accelerator_type=GpuVersion.XPK_H100,
-          gpu_zone=Zone.US_CENTRAL1_C.value,
-          time_out_in_min=300,
-          test_name=f"{test_name_prefix}-stable-{test_script}",
-          test_script=test_script,
-          num_slices=2,
-          cluster_name=ClusterName.A3_CLUSTER.value,
-          docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE.value,
-          test_owner=test_owner.NINA_C,
-      ).run()
-      nightly_gpu = gke_config.get_maxtext_end_to_end_gpu_gke_test_config(
-          accelerator_type=GpuVersion.XPK_H100,
-          gpu_zone=Zone.US_CENTRAL1_C.value,
-          time_out_in_min=300,
-          test_name=f"{test_name_prefix}-nightly-{test_script}",
-          test_script=test_script,
-          num_slices=2,
-          cluster_name=ClusterName.A3_CLUSTER.value,
-          docker_image=DockerImage.MAXTEXT_GPU_JAX_NIGHTLY.value,
-          test_owner=test_owner.NINA_C,
-      ).run()
-      stable_tpu >> nightly_tpu >> stable_gpu >> nightly_gpu
+  for model, test_script in test_models.items():
+    stable_tpu = gke_config.get_gke_config(
+        tpu_version=TpuVersion.V4,
+        tpu_cores=8,
+        tpu_zone=Zone.US_CENTRAL2_B.value,
+        time_out_in_min=60,
+        test_name=f"{test_name_prefix}-stable-{model}",
+        run_model_cmds=(f"bash end_to_end/{test_script}.sh",),
+        docker_image=DockerImage.MAXTEXT_TPU_JAX_STABLE.value,
+        test_owner=test_owner.JON_B,
+    ).run()
+    nightly_tpu = gke_config.get_gke_config(
+        tpu_version=TpuVersion.V4,
+        tpu_cores=8,
+        tpu_zone=Zone.US_CENTRAL2_B.value,
+        time_out_in_min=60,
+        test_name=f"{test_name_prefix}-nightly-{model}",
+        run_model_cmds=(f"bash end_to_end/{test_script}.sh",),
+        docker_image=DockerImage.MAXTEXT_TPU_JAX_NIGHTLY.value,
+        test_owner=test_owner.JON_B,
+    ).run()
+    stable_gpu = gke_config.get_maxtext_end_to_end_gpu_gke_test_config(
+        accelerator_type=GpuVersion.XPK_H100,
+        gpu_zone=Zone.US_CENTRAL1_C.value,
+        time_out_in_min=300,
+        test_name=f"{test_name_prefix}-stable-{model}",
+        test_script=test_script,
+        num_slices=2,
+        cluster_name=ClusterName.A3_CLUSTER.value,
+        docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE.value,
+        test_owner=test_owner.NINA_C,
+    ).run()
+    nightly_gpu = gke_config.get_maxtext_end_to_end_gpu_gke_test_config(
+        accelerator_type=GpuVersion.XPK_H100,
+        gpu_zone=Zone.US_CENTRAL1_C.value,
+        time_out_in_min=300,
+        test_name=f"{test_name_prefix}-nightly-{model}",
+        test_script=test_script,
+        num_slices=2,
+        cluster_name=ClusterName.A3_CLUSTER.value,
+        docker_image=DockerImage.MAXTEXT_GPU_JAX_NIGHTLY.value,
+        test_owner=test_owner.NINA_C,
+    ).run()
+    stable_tpu >> nightly_tpu >> stable_gpu >> nightly_gpu