Add job to build and upload docker image to pypi_release workflow

Author: SurbhiJainUSC

.github/workflows/build_and_push_docker_image.yml

@@ -43,8 +43,17 @@ on:
         default: ''
       workflow:
         required: false
+<<<<<<< HEAD
         type: string
         default: 'pre-training'
+=======
+        type: boolean
+        default: false
+      version_name:
+        required: false
+        type: string
+        default: ''
+>>>>>>> 46b04fed2 (Add job to build and upload docker image to pypi_release workflow)
 
 permissions:
   contents: read
@@ -115,7 +124,7 @@ jobs:
           push: true
           context: .
           file: ${{ inputs.dockerfile }}
-          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:latest
+          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ github.run_id }}
           cache-from: type=gha
           outputs: type=image,compression=zstd,force-compression=true
           build-args: |
@@ -133,26 +142,32 @@ jobs:
         run: |
           SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${INPUTS_IMAGE_NAME}"
 
-          # Add date tag
-          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
+          if [[ $INPUTS_VERSION_NAME ]]; then
+            echo "Tagging docker images corresponding to PyPI release..."
+            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_VERSION_NAME}" --quiet
+          else
+            echo "Tagging docker images corresponding to nightly release..."
 
-          # Convert date to YYYYMMDD format
-          clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
+            # Add date tag
+            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
 
-          # Add MaxText tag
-          maxtext_hash=$(git rev-parse --short HEAD)
-          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
+            # Convert date to YYYYMMDD format
+            clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
 
+            # Add MaxText tag
+            maxtext_hash=$(git rev-parse --short HEAD)
+            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
 
           # Add post-training dependencies tags
           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)
-                gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
-              fi
-            done
+                gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
+                fi
+              done
+            fi
           fi
         env:
           INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
           INPUTS_IMAGE_DATE: ${{ inputs.image_date }}
+          INPUTS_VERSION_NAME: ${{ inputs.version_name }}

.github/workflows/pypi_release.yml

@@ -44,9 +44,7 @@ jobs:
 
   publish_maxtext_package_to_pypi:
     name: Publish MaxText package to PyPI
-    # Temporarily only require release_approval for a one-time upload.
-    # Immediately revert this to `needs: [build_and_test_maxtext_package]`.
-    needs: [release_approval]
+    needs: [build_and_test_maxtext_package]
     runs-on: ubuntu-latest
     environment: release
     steps:
@@ -61,3 +59,38 @@ jobs:
       uses: pypa/gh-action-pypi-publish@release/v1
       with:
         packages-dir: dist/
+
+  # This job builds and pushes MaxText stable Docker images for both TPU and GPU devices.
+  # It runs only after the release is approved, ensuring that the Docker images are built and pushed only for approved releases.
+  # Creates docker image for MaxText commit corresponding to the release.
+  upload_maxtext_docker_images:
+    name: Build and upload MaxText stable Docker images
+    needs: [publish_maxtext_package_to_pypi]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - device: tpu
+            build_mode: stable
+            image_name: maxtext_jax_stable
+            workflow: pre-training
+            dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+          - device: gpu
+            build_mode: stable
+            image_name: maxtext_gpu_jax_stable
+            workflow: pre-training
+            dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+          - device: tpu
+            build_mode: stable
+            image_name: maxtext_post_training_stable
+            workflow: post-training
+            dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+    uses: ./.github/workflows/build_and_push_docker_image.yml
+    with:
+      image_name: ${{ matrix.image_name }}
+      device: ${{ matrix.device }}
+      build_mode: ${{ matrix.build_mode }}
+      workflow: ${{ matrix.workflow }}
+      dockerfile: ${{ matrix.dockerfile }}
+      maxtext_sha: ${{ github.sha }}
+      version_name: ${{ github.event.release.tag_name }}