Adding support for building Maxtext images with nightly version of JS…

…S for GPUs
AI-Hypercomputer · Nov 12, 2024 · a70fba1 · a70fba1
1 parent 2780510
commit a70fba1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
@@ -64,3 +64,6 @@ jobs:
     - name: build jax stable stack image
       run : |
         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxtext_gpu_jax_stable_stack_0.4.35 MODE=stable_stack DEVICE=gpu PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxtext_jax_stable_stack_0.4.35 BASEIMAGE=gcr.io/tpu-prod-env-multipod/jax-stable-stack/gpu:jax0.4.35-rev1 MAXTEXT_REQUIREMENTS_FILE=requirements_with_jax_stable_stack.txt
+    - name: build image with nightly jax stable stack
+    run: |
+      bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxtext_gpu_jax_stable_stack_nightly MODE=stable_stack DEVICE=gpu PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxtext_gpu_jax_stable_stack_nightly BASEIMAGE=us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack-gpu/nightly:latest MAXTEXT_REQUIREMENTS_FILE=requirements_with_jax_stable_stack.txt