Build post-training docker images in github workflow dailiy

Author: SurbhiJainUSC

.github/workflows/UploadDockerImages.yml

@@ -49,19 +49,36 @@ jobs:
             build_mode: stable
             image_name: maxtext_jax_stable
             dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
+            run_secondary: false
           - device: tpu
             build_mode: nightly
             image_name: maxtext_jax_nightly
             dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
+            run_secondary: false
+          # Post-Training Image Builds
+          - device: tpu
+            build_mode: stable
+            image_name: maxtext_post_training
+            dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
+            secondary_dockerfile: ./dependencies/dockerfiles/maxtext_post_training_dependencies.Dockerfile
+            run_secondary: true
+          - device: tpu
+            build_mode: stable
+            image_name: maxtext_post_training_local
+            dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
+            secondary_dockerfile: ./dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile
+            run_secondary: true
           # GPU Image Builds
           - device: gpu
             build_mode: stable
             image_name: maxtext_gpu_jax_stable
             dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+            run_secondary: false
           - device: gpu
             build_mode: nightly
             image_name: maxtext_gpu_jax_nightly
             dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+            run_secondary: false
 
     if: >
       github.event_name == 'schedule' ||
@@ -86,13 +103,20 @@ jobs:
             echo "Building for device: ${{ matrix.device }} in ${{ matrix.build_mode }} mode."
           fi
 
-      - name: Checkout git repository
+      - name: Checkout MaxText
         uses: actions/checkout@v5
         if: steps.check.outputs.should_run == 'true'
 
-      - name: Mark git repository as safe
+      - name: Checkout post-training dependencies
+        if: steps.check.outputs.should_run == 'true' && matrix.image_name == "maxtext_post_training_local"
+        run: |
+          git clone https://github.com/google/tunix.git ./tunix
+          git clone https://github.com/vllm-project/vllm.git ./vllm
+          git clone https://github.com/google/tpu-inference.git ./tpu-inference
+
+      - name: Mark git repositories as safe
         if: steps.check.outputs.should_run == 'true'
-        run: git config --global --add safe.directory ${GITHUB_WORKSPACE}
+        run: git config --global --add safe.directory '*'
 
       - name: Configure Docker
         if: steps.check.outputs.should_run == 'true'
@@ -110,21 +134,31 @@ jobs:
         id: vars
         if: steps.check.outputs.should_run == 'true'
         run: |
-          echo "commit_hash=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
           echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT
 
-      # Docker BuildX command config
-      - name: Build and Push Docker Image
+          # MaxText commit hash
+          echo "maxtext_hash=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
+
+          # Post-training dependencies commit hashes
+          for dir in tunix vllm tpu-inference; do
+            if [ -d "./$dir" ]; then
+              cd $dir
+              echo "${dir}_hash=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
+              cd ..
+            else
+              echo "${dir}_hash=" >> $GITHUB_OUTPUT
+            fi
+          done
+
+      # Docker BuildX command config for pre-training images
+      - name: Build and Push Primary Docker Image
         uses: docker/build-push-action@v6
         if: steps.check.outputs.should_run == 'true'
         with:
-          push: true
+          push: {{ !matrix.run_secondary }}
           context: .
           file: ${{ matrix.dockerfile }}
-          tags: |
-            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:maxtext_${{ steps.vars.outputs.commit_hash }}
-            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.vars.outputs.image_date }}
-            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest
+          tags: gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest
           cache-from: type=gha
           cache-to: type=gha,mode=max
           provenance: false
@@ -133,3 +167,40 @@ jobs:
             MODE=${{ matrix.build_mode }}
             JAX_VERSION=NONE
             LIBTPU_GCS_PATH=NONE
+
+      # Docker BuildX command config for post-training images
+      - name: Build and Push Secondary Docker Image
+        uses: docker/build-push-action@v6
+        if: steps.check.outputs.should_run == 'true' && matrix.run_secondary
+        with:
+          push: true
+          context: .
+          file: ${{ matrix.secondary_dockerfile }}
+          tags: gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+          provenance: false
+          build-args: |
+            MODE=post-training
+            BASEIMAGE=maxtext_base_image
+
+      - name: Add tags to Docker Image
+        if: steps.check.outputs.should_run == 'true'
+        run: |
+          SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}"
+
+          # Add date tag
+          gcloud container images add-tag "$SOURCE_IMAGE:latest" \
+            "gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.vars.outputs.image_date }}" --quiet
+
+          # Add MaxText tag
+          gcloud container images add-tag "$SOURCE_IMAGE:latest" \
+            "gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:maxtext_${{ steps.vars.outputs.maxtext_hash }}" --quiet
+
+          # Add post-training dependencies tags
+          for dep in tunix vllm tpu_inference; do
+            HASH=$(echo "${{ steps.vars.outputs }}" | jq -r ".$dep" || echo "${{ steps.vars.outputs[format('{0}_hash', dep)] }}")
+            if [ -n "$HASH" ]; then
+              gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${dep//_/-}_$HASH" --quiet
+            fi
+          done