fix: Update the image for maxtext_gpu_end_to_end (GoogleCloudPlatform#1117)

This change updates the Docker image used by the `maxtext_gpu_end_to_end` DAG to align with changes in the MaxText repository.

Author: ooops678

dags/common/vm_resource.py

@@ -368,20 +368,12 @@ class DockerImage(enum.Enum):
       "gcr.io/tpu-prod-env-multipod/maxdiffusion_stable_stack_candidate:"
       f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
   )
-  MAXTEXT_GPU_JAX_STABLE_STACK = (
-      "gcr.io/tpu-prod-env-multipod/maxtext_gpu_jax_stable_stack:"
-      f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
-  )
   MAXTEXT_GPU_JAX_STABLE = (
       "gcr.io/tpu-prod-env-multipod/maxtext_gpu_jax_stable:"
       f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
   )
-  MAXTEXT_GPU_STABLE_STACK_NIGHTLY_JAX = (
-      "gcr.io/tpu-prod-env-multipod/maxtext_gpu_stable_stack_nightly_jax:"
-      f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
-  )
-  MAXTEXT_GPU_JAX_AI_CANDIDATE_IMAGE = (
-      "gcr.io/tpu-prod-env-multipod/maxtext_stable_stack_candidate_gpu:"
+  MAXTEXT_GPU_JAX_NIGHTLY = (
+      "gcr.io/tpu-prod-env-multipod/maxtext_gpu_jax_nightly:"
       f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
   )
   CLOUD_HYBRIDSIM_NIGHTLY = (

dags/multipod/maxtext_configs_aot_gpu.py

@@ -20,9 +20,8 @@
 from airflow.utils.task_group import TaskGroup
 from dags import composer_env
 from dags.common import test_owner
-from dags.common.vm_resource import GpuVersion, TpuVersion, Zone, DockerImage, XpkClusters
+from dags.common.vm_resource import DockerImage, XpkClusters
 from dags.multipod.configs import gke_config
-from dags.multipod.configs.common import SetupMode
 
 # Run once a day at 5 am UTC (9 pm PST / 10 pm PDT)
 SCHEDULED_TIME = "45 4 * * *" if composer_env.is_prod_env() else None
@@ -56,6 +55,6 @@
       run_model_cmds=(cmd,),
       num_slices=1,
       cluster=XpkClusters.GPU_A3PLUS_CLUSTER,
-      docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE_STACK.value,
+      docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE.value,
       test_owner=test_owner.NUOJIN_C,
   ).run_with_quarantine(quarantine_task_group)

dags/multipod/maxtext_gpu_end_to_end.py

@@ -24,7 +24,7 @@
 from airflow.utils.task_group import TaskGroup
 from dags import composer_env
 from dags.common import test_owner
-from dags.common.vm_resource import XpkClusters, CpuVersion, DockerImage, GpuVersion, Project, TpuVersion, Zone
+from dags.common.vm_resource import XpkClusters, DockerImage, Project
 from dags.multipod.configs import gke_config
 from xlml.utils import gke
 
@@ -206,7 +206,7 @@ def run_maxtext_tests(dag: models.DAG):
         run_model_cmds=(test_script,),
         num_slices=nnodes,
         cluster=XpkClusters.GPU_A3_CLUSTER,
-        docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE_STACK.value,
+        docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE.value,
         test_owner=test_owner.ANISHA_M,
     ).run_with_quarantine(quarantine_task_group)
     stable_a3plus_gpu = gke_config.get_maxtext_end_to_end_gpu_gke_test_config(
@@ -215,7 +215,7 @@ def run_maxtext_tests(dag: models.DAG):
         run_model_cmds=(test_script,),
         num_slices=nnodes,
         cluster=XpkClusters.GPU_A3PLUS_CLUSTER,
-        docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE_STACK.value,
+        docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE.value,
         test_owner=test_owner.ANISHA_M,
     ).run_with_quarantine(quarantine_task_group)
     stable_a3_gpu >> stable_a3plus_gpu

dags/sparsity_diffusion_devx/jax_ai_image_gpu_e2e.py

@@ -17,12 +17,11 @@
 
 import datetime
 from airflow import models
-from dags import composer_env, gcs_bucket
+from dags import composer_env
 from dags.common import test_owner
-from dags.common.vm_resource import Project, TpuVersion, CpuVersion, Zone, DockerImage, GpuVersion, XpkClusters
+from dags.common.vm_resource import DockerImage, XpkClusters
 from airflow.utils.task_group import TaskGroup
 from dags.sparsity_diffusion_devx.configs import gke_config as config
-from xlml.utils import name_format
 from dags.multipod.configs.common import SetupMode
 
 # Run once a day at 7 am UTC (11 pm PST)
@@ -70,7 +69,7 @@
   )
 
   docker_images = [
-      (SetupMode.NIGHTLY, DockerImage.MAXTEXT_GPU_STABLE_STACK_NIGHTLY_JAX),
+      (SetupMode.NIGHTLY, DockerImage.MAXTEXT_GPU_JAX_NIGHTLY),
   ]
 
   for model, (test_script, nnodes) in test_models_gpu.items():

dags/sparsity_diffusion_devx/maxtext_moe_gpu_e2e.py

@@ -63,7 +63,7 @@ def run_maxtext_tests():
             run_model_cmds=(test_script,),
             num_slices=nnodes,
             cluster=XpkClusters.GPU_A3PLUS_CLUSTER,
-            docker_image=DockerImage.MAXTEXT_GPU_JAX_AI_CANDIDATE_IMAGE.value,
+            docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE.value,
             test_owner=test_owner.MICHELLE_Y,
         ).run_with_quarantine(quarantine_task_group)
     )
@@ -73,7 +73,7 @@ def run_maxtext_tests():
         run_model_cmds=(test_script,),
         num_slices=nnodes,
         cluster=XpkClusters.GPU_A3PLUS_CLUSTER,
-        docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE_STACK.value,
+        docker_image=DockerImage.MAXTEXT_GPU_JAX_STABLE.value,
         test_owner=test_owner.MICHELLE_Y,
     ).run_with_quarantine(quarantine_task_group)
     candidate_a3plus_gpu >> stable_a3plus_gpu