Only run gke mxla

GoogleCloudPlatform · Jan 10, 2025 · 155c05a · 155c05a
1 parent d628936
commit 155c05a
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 149 deletions.
diff --git a/dags/multipod/maxtext_trillium_configs_perf.py b/dags/multipod/maxtext_trillium_configs_perf.py
@@ -27,10 +27,14 @@
 
 # Run once a day at 4 am UTC (8 pm PST / 9 pm PDT)
 SCHEDULED_TIME = "0 4 * * *" if composer_env.is_prod_env() else None
-MODEL_CONFIGS = ["gpt3_175b", "llama2_7b_4096", "mixtral_8x7b"]
+#MODEL_CONFIGS = ["gpt3_175b", "llama2_7b_4096", "mixtral_8x7b"]
+MODEL_CONFIGS = ["llama2_7b_4096"]
+# DOCKER_IMAGES = [
+#     (SetupMode.STABLE, DockerImage.MAXTEXT_TPU_JAX_STABLE_STACK),
+#     (SetupMode.NIGHTLY, DockerImage.MAXTEXT_TPU_JAX_NIGHTLY),
+# ]
 DOCKER_IMAGES = [
     (SetupMode.STABLE, DockerImage.MAXTEXT_TPU_JAX_STABLE_STACK),
-    (SetupMode.NIGHTLY, DockerImage.MAXTEXT_TPU_JAX_NIGHTLY),
 ]
 QUANTIZATION_SWEEP = {"M_QUANTIZATION": ["", "int8"]}
 BASE_OUTPUT_DIRECTORY = "gs://runner-maxtext-logs"

diff --git a/dags/multipod/mxla_maxtext_nightly.py b/dags/multipod/mxla_maxtext_nightly.py
diff --git a/dags/multipod/mxla_maxtext_nightly_gke.py b/dags/multipod/mxla_maxtext_nightly_gke.py
@@ -24,8 +24,7 @@
 from dags.multipod.configs import gke_config
 
 # Run once a day at 9 am UTC (1 am PST)
-# Pause test on GKE
-SCHEDULED_TIME = None
+SCHEDULED_TIME = "0 9 * * *" if composer_env.is_prod_env() else None
 
 with models.DAG(
     dag_id="mxla_maxtext_nightly_gke",