From 34db38e770a5542bf549c2c0f8373a11b17bdc52 Mon Sep 17 00:00:00 2001 From: Surbhi Jain Date: Mon, 11 May 2026 23:54:38 +0000 Subject: [PATCH] Publish MaxText stable docker images to cloud-tpu-images --- .github/workflows/UploadDockerImages.yml | 1 + .../workflows/build_and_push_docker_image.yml | 7 +- .../workflows/publish_public_docker_image.yml | 121 ++++++++++++++++++ .github/workflows/pypi_release.yml | 65 ---------- 4 files changed, 127 insertions(+), 67 deletions(-) create mode 100644 .github/workflows/publish_public_docker_image.yml diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml index 93f0b081bc..ed4f6f57c5 100644 --- a/.github/workflows/UploadDockerImages.yml +++ b/.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 }} diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml index 5cdc2aaf41..9059bb136f 100644 --- a/.github/workflows/build_and_push_docker_image.yml +++ b/.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 diff --git a/.github/workflows/publish_public_docker_image.yml b/.github/workflows/publish_public_docker_image.yml new file mode 100644 index 0000000000..1f8f02b891 --- /dev/null +++ b/.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 }} diff --git a/.github/workflows/pypi_release.yml b/.github/workflows/pypi_release.yml index 0eba10c310..3461d357a3 100644 --- a/.github/workflows/pypi_release.yml +++ b/.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 }}