[DO NOT REVIEW] Publish TPU pre-training and post-training stable docker images to cloud-tpu-images

.github/workflows/UploadDockerImages.yml

@@ -84,6 +84,7 @@ jobs:
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
       image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
+      docker_image_prefix: gcr.io/${{ vars.PROJECT_NAME }}
       device: ${{ matrix.device }}
       build_mode: ${{ matrix.build_mode }}
       workflow: ${{ matrix.workflow }}

.github/workflows/build_and_push_docker_image.yml

@@ -38,6 +38,9 @@ on:
         required: false
         type: string
         default: 'pre-training'
+      docker_image_prefix:
+        required: true
+        type: string
       version_name:
         required: false
         type: string
@@ -148,7 +151,7 @@ jobs:
           push: true
           context: .
           file: .venv/lib/python3.12/site-packages/dependencies/dockerfiles/${{ inputs.dockerfile }}
-          tags: gcr.io/${{ vars.PROJECT_NAME }}/${{ inputs.image_name }}:${{ github.run_id }}
+          tags: ${{ inputs.docker_image_prefix }}/${{ inputs.image_name }}:${{ github.run_id }}
           cache-from: type=gha
           outputs: type=image,compression=zstd,force-compression=true
           build-args: |
@@ -161,7 +164,7 @@ jobs:
       - name: Add tags to Docker image
         shell: bash
         run: |
-          SOURCE_IMAGE="gcr.io/${{ vars.PROJECT_NAME }}/${INPUTS_IMAGE_NAME}"
+          SOURCE_IMAGE="${{ inputs.docker_image_prefix }}/${{ inputs.image_name }}"
           TEMP_IMG="${SOURCE_IMAGE}:${{ github.run_id }}"
 
           if [[ $INPUTS_VERSION_NAME ]]; then

.github/workflows/publish_public_docker_image.yml

@@ -0,0 +1,121 @@
+# Copyright 2023-2026 Google LLC
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This workflow will build and push MaxText stable Docker images for both TPU and GPU devices.
+# It runs only after a new release is published to PyPI.
+# Creates docker image for MaxText commit corresponding to the release.
+
+name: Publish MaxText Stable Docker Images
+
+on:
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  get_latest_maxtext_pypi_version:
+    name: Get latest MaxText PyPI version
+    runs-on: ubuntu-latest
+    outputs:
+      latest_pypi_version: ${{ steps.get_version.outputs.version }}
+    steps:
+      - name: Install jq
+        run: sudo apt-get update && sudo apt-get install -y jq
+      - name: Fetch latest version of maxtext from PyPI
+        id: get_version
+        run: |
+          # Fetch JSON from PyPI for 'maxtext'
+          echo "Fetching latest version from https://pypi.org/pypi/maxtext/json"
+          pypi_json=$(curl -s https://pypi.org/pypi/maxtext/json)
+
+          # Extract the version from the "info" section using jq
+          latest_version=$(echo "$pypi_json" | jq -r ".info.version")
+
+          if [ -z "$latest_version" ] || [ "$latest_version" == "null" ]; then
+            echo "Error: Could not parse latest version from PyPI JSON."
+            exit 1
+          fi
+
+          echo "Successfully fetched latest MaxText version on PyPI: $latest_version"
+          # Set the output variable for other jobs to consume
+          echo "version=$latest_version" >> "$GITHUB_OUTPUT"
+
+  build_maxtext_package:
+    name: Build MaxText Package
+    needs: get_latest_maxtext_pypi_version
+    runs-on: ubuntu-latest
+    outputs:
+      commit_sha: ${{ steps.get_sha.outputs.commit_sha }}
+    steps:
+      - name: Checkout MaxText
+        uses: actions/checkout@v5
+        with:
+          fetch-depth: 0 # Fetch all history for all tags
+      - name: Get commit SHA for tag
+        id: get_sha
+        run: |
+          PYPI_VERSION="${{ needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version }}"
+          TAG_NAME="maxtext-v${PYPI_VERSION}"
+          echo "Looking for tag: ${TAG_NAME}"
+
+          # Get the commit SHA for the tag
+          COMMIT_SHA=$(git rev-parse --verify "${TAG_NAME}^{commit}" 2>/dev/null)
+
+          echo "Found commit SHA: ${COMMIT_SHA}"
+          echo "commit_sha=${COMMIT_SHA}" >> "$GITHUB_OUTPUT"
+      - name: Install build tools
+        run: |
+          python -m pip install --upgrade pip build uv
+      - name: Build maxtext wheel
+        run: |
+          uv build --wheel
+      - name: Upload the built maxtext wheel
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02
+        with:
+          name: maxtext-wheel
+          path: dist/*
+
+  upload_maxtext_docker_images:
+    name: ${{ matrix.image_name }}
+    needs: [get_latest_maxtext_pypi_version, build_maxtext_package]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - device: tpu
+            build_mode: stable
+            image_name: tpu_pre_training
+            workflow: pre-training
+            dockerfile: maxtext_tpu_dependencies.Dockerfile
+          - device: gpu
+            build_mode: stable
+            image_name: gpu_pre_training
+            workflow: pre-training
+            dockerfile: maxtext_gpu_dependencies.Dockerfile
+          - device: tpu
+            build_mode: stable
+            image_name: tpu_post_training
+            workflow: post-training
+            dockerfile: maxtext_tpu_dependencies.Dockerfile
+    uses: ./.github/workflows/build_and_push_docker_image.yml
+    with:
+      image_name: ${{ matrix.image_name }}
+      docker_image_prefix: us-docker.pkg.dev/${{ vars.PUBLIC_IMAGE_PROJECT_NAME }}/maxtext-images
+      device: ${{ matrix.device }}
+      build_mode: ${{ matrix.build_mode }}
+      workflow: ${{ matrix.workflow }}
+      dockerfile: ${{ matrix.dockerfile }}
+      maxtext_sha: ${{ needs.build_maxtext_package.outputs.commit_sha }}
+      version_name: needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}

.github/workflows/pypi_release.yml

@@ -60,68 +60,3 @@ jobs:
       uses: pypa/gh-action-pypi-publish@release/v1
       with:
         packages-dir: dist/
-
-  get_latest_maxtext_pypi_version:
-    name: Get latest MaxText PyPI version
-    needs: [publish_maxtext_package_to_pypi]
-    runs-on: ubuntu-latest
-    outputs:
-      latest_pypi_version: ${{ steps.get_version.outputs.version }}
-    steps:
-      - name: Install jq
-        run: sudo apt-get update && sudo apt-get install -y jq
-      - name: Fetch latest version of maxtext from PyPI
-        id: get_version
-        run: |
-          # Fetch JSON from PyPI for 'maxtext'
-          echo "Fetching latest version from https://pypi.org/pypi/maxtext/json"
-          pypi_json=$(curl -s https://pypi.org/pypi/maxtext/json)
-
-          # Extract the version from the "info" section using jq
-          latest_version=$(echo "$pypi_json" | jq -r ".info.version")
-
-          if [ -z "$latest_version" ] || [ "$latest_version" == "null" ]; then
-            echo "Error: Could not parse latest version from PyPI JSON."
-            exit 1
-          fi
-
-          echo "Successfully fetched latest MaxText version on PyPI: $latest_version"
-          # Set the output variable for other jobs to consume
-          echo "version=$latest_version" >> "$GITHUB_OUTPUT"
-
-  # This job builds and pushes MaxText stable Docker images for both TPU and GPU devices.
-  # It runs only after a new release is published to PyPI.
-  # Creates docker image for MaxText commit corresponding to the release.
-  upload_maxtext_docker_images:
-    name: ${{ matrix.image_name }}
-    needs: [get_latest_maxtext_pypi_version]
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          - device: tpu
-            build_mode: stable
-            image_name: maxtext_jax_stable
-            workflow: pre-training
-            dockerfile: maxtext_tpu_dependencies.Dockerfile
-          - device: gpu
-            build_mode: stable
-            image_name: maxtext_gpu_jax_stable
-            workflow: pre-training
-            dockerfile: maxtext_gpu_dependencies.Dockerfile
-          - device: tpu
-            build_mode: stable
-            image_name: maxtext_post_training_stable
-            workflow: post-training
-            dockerfile: 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: ${{ needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version }}
-    secrets:
-      HF_TOKEN: ${{ secrets.HF_TOKEN }}