Merge pull request #3335 from AI-Hypercomputer:docker_image_fix

PiperOrigin-RevId: 879728816

Author: Google-ML-Automation

.github/workflows/UploadDockerImages.yml

@@ -79,31 +79,32 @@ jobs:
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
 
-  tpu-post-training:
-    name: ${{ matrix.image_name }}
-    needs: [setup, tpu-pre-training]
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - device: tpu
-            build_mode: post-training
-            image_name: maxtext_post_training_stable
-            dockerfile: ./dependencies/dockerfiles/maxtext_post_training_dependencies.Dockerfile
-          - device: tpu
-            build_mode: post-training
-            image_name: maxtext_post_training_nightly
-            dockerfile: ./dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile
+  tpu-post-training-stable:
+    name: tpu-post-training-stable
+    needs: setup
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
-      image_name: ${{ matrix.image_name }}
-      device: ${{ matrix.device }}
-      build_mode: ${{ matrix.build_mode }}
-      dockerfile: ${{ matrix.dockerfile }}
+      image_name: maxtext_post_training_stable
+      device: tpu
+      build_mode: stable
+      workflow: post-training
+      dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+      maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
+      image_date: ${{ needs.setup.outputs.image_date }}
+
+  tpu-post-training-nightly:
+    name: tpu-post-training-nightly
+    needs: [setup, tpu-post-training-stable]
+    uses: ./.github/workflows/build_and_push_docker_image.yml
+    with:
+      image_name: maxtext_post_training_nightly
+      device: tpu
+      build_mode: nightly
+      workflow: post-training
+      dockerfile: ./dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
-      base_image: gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:${{ needs.setup.outputs.image_date }}
-      is_post_training: true
+      base_image: gcr.io/tpu-prod-env-multipod/maxtext_post_training_stable:${{ needs.setup.outputs.image_date }}
 
   gpu-pre-training:
     name: ${{ matrix.image_name }}

.github/workflows/build_and_push_docker_image.yml

@@ -41,10 +41,10 @@ on:
         required: false
         type: string
         default: ''
-      is_post_training:
+      workflow:
         required: false
-        type: boolean
-        default: false
+        type: string
+        default: 'pre-training'
 
 permissions:
   contents: read
@@ -121,6 +121,7 @@ jobs:
           build-args: |
             DEVICE=${{ inputs.device }}
             MODE=${{ inputs.build_mode }}
+            WORKFLOW=${{ inputs.workflow }}
             JAX_VERSION=NONE
             LIBTPU_VERSION=NONE
             INCLUDE_TEST_ASSETS=true
@@ -144,7 +145,7 @@ jobs:
 
 
           # Add post-training dependencies tags
-          if [ "${{ inputs.is_post_training }}" == "true" ]; then
+          if [ "${{ inputs.workflow }}" == "post-training" ]; then
             for dir in tunix vllm tpu-inference; do
               if [ -d "./$dir" ]; then
                 dir_hash=$(git -C "$dir" rev-parse --short HEAD)