move to correct nightly image (GoogleCloudPlatform#784)

Author: khatwanimohit

dags/multipod/maxtext_checkpointing.py

@@ -46,7 +46,7 @@
   current_datetime = current_time.strftime("%Y-%m-%d-%H-%M-%S")
   docker_images = [
       (SetupMode.STABLE, DockerImage.MAXTEXT_TPU_JAX_STABLE_STACK_CANDIDATE),
-      (SetupMode.NIGHTLY, DockerImage.MAXTEXT_TPU_JAX_NIGHTLY),
+      (SetupMode.NIGHTLY, DockerImage.MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX),
   ]
   test_configs = {
       # accelerator: list of slices to test

dags/multipod/maxtext_end_to_end.py

@@ -74,7 +74,7 @@
             f"export HF_TOKEN={HF_TOKEN}",
             f"bash end_to_end/{test_script}.sh",
         ),
-        docker_image=DockerImage.MAXTEXT_TPU_JAX_NIGHTLY.value,
+        docker_image=DockerImage.MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX.value,
         cluster=XpkClusters.TPU_V5P_8_CLUSTER,
         test_owner=test_owner.MOHIT_K,
     ).run_with_quarantine(quarantine_task_group)
@@ -147,7 +147,7 @@ def convert_checkpoint_and_run_training(
 
   docker_image = {
       "stable": DockerImage.MAXTEXT_TPU_JAX_STABLE_STACK_CANDIDATE.value,
-      "nightly": DockerImage.MAXTEXT_TPU_JAX_NIGHTLY.value,
+      "nightly": DockerImage.MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX.value,
   }
   tests = []
   for model, test_scripts_details in multicluster_test_models.items():

dags/multipod/maxtext_sft_trainer.py

@@ -37,7 +37,7 @@
   base_output_directory = f'{gcs_bucket.BASE_OUTPUT_DIR}/maxtext_sft_trainer'
   docker_images = [
       (SetupMode.STABLE, DockerImage.MAXTEXT_TPU_JAX_STABLE_STACK_CANDIDATE),
-      (SetupMode.NIGHTLY, DockerImage.MAXTEXT_TPU_JAX_NIGHTLY),
+      (SetupMode.NIGHTLY, DockerImage.MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX),
   ]
 
   for mode, image in docker_images: