From 0adbeb64070108b27fca9808966a98c9eb98ffb2 Mon Sep 17 00:00:00 2001 From: Surbhi Jain Date: Wed, 27 May 2026 21:33:17 +0000 Subject: [PATCH] Publish MaxText to PyPI when Airflow E2E tests succeed --- .../workflows/build_and_push_docker_image.yml | 36 ++-- .github/workflows/build_package.yml | 8 +- .github/workflows/check_docs_build.yml | 12 ++ .github/workflows/ci_pipeline.yml | 11 ++ .github/workflows/docs_link_check.yml | 5 + .github/workflows/pypi_release.yml | 185 ++++++++++-------- .github/workflows/release_pipeline.yml | 98 ++++++++-- .github/workflows/run_e2e_tests.yml | 3 +- 8 files changed, 242 insertions(+), 116 deletions(-) diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml index a9e8d6547d..b64629d214 100644 --- a/.github/workflows/build_and_push_docker_image.yml +++ b/.github/workflows/build_and_push_docker_image.yml @@ -38,14 +38,14 @@ on: required: false type: string default: 'pre-training' - version_name: - required: false - type: string - default: '' include_test_assets: required: false type: boolean default: false + outputs: + build_result: + description: 'The result of build_and_push job.' + value: ${{ jobs.build_and_push.result }} secrets: HF_TOKEN: required: true @@ -65,7 +65,8 @@ jobs: shell: bash run: | EVENT_NAME="${{ github.event_name }}" - TARGET_DEVICE="${{ github.event.inputs.target_device }}" + # Default to 'all' if the input is null/empty + TARGET_DEVICE="${{ github.event.inputs.target_device || 'all' }}" INPUT_DEVICE="${{ inputs.device }}" SHOULD_RUN="false" @@ -121,6 +122,7 @@ jobs: uses: actions/download-artifact@v4 with: name: maxtext-wheel + path: . - name: Install uv and set Python version uses: astral-sh/setup-uv@v7 @@ -166,26 +168,18 @@ jobs: SOURCE_IMAGE="gcr.io/${{ vars.PROJECT_NAME }}/${INPUTS_IMAGE_NAME}" TEMP_IMG="${SOURCE_IMAGE}:${{ github.run_id }}" - if [[ $INPUTS_VERSION_NAME ]]; then - echo "Tagging docker images corresponding to PyPI release..." - gcloud container images add-tag "${TEMP_IMG}" "${SOURCE_IMAGE}:${INPUTS_VERSION_NAME}" --quiet - else - echo "Tagging docker images corresponding to nightly release..." - - # Add date tag - IMAGE_DATE="$(date +%Y-%m-%d)" - gcloud container images add-tag "${TEMP_IMG}" "$SOURCE_IMAGE:$IMAGE_DATE" --quiet + # Add date tag + IMAGE_DATE="$(date +%Y-%m-%d)" + gcloud container images add-tag "${TEMP_IMG}" "$SOURCE_IMAGE:$IMAGE_DATE" --quiet - # Convert date to YYYYMMDD format - clean_date=$(echo "$IMAGE_DATE" | sed 's/[-:]//g' | cut -c1-8) + # Convert date to YYYYMMDD format + clean_date=$(echo "$IMAGE_DATE" | sed 's/[-:]//g' | cut -c1-8) - # Add MaxText tag - MAXTEXT_SHA=$(git rev-parse --short HEAD) - gcloud container images add-tag "${TEMP_IMG}" "${SOURCE_IMAGE}:maxtext_${MAXTEXT_SHA}_${clean_date}" --quiet - fi + # Add MaxText tag + MAXTEXT_SHA=$(git rev-parse --short HEAD) + gcloud container images add-tag "${TEMP_IMG}" "${SOURCE_IMAGE}:maxtext_${MAXTEXT_SHA}_${clean_date}" --quiet env: INPUTS_IMAGE_NAME: ${{ inputs.image_name }} - INPUTS_VERSION_NAME: ${{ inputs.version_name }} run_ci_tests: name: Run Unit and Integration Tests diff --git a/.github/workflows/build_package.yml b/.github/workflows/build_package.yml index 25f2fd2229..b6126e2a67 100644 --- a/.github/workflows/build_package.yml +++ b/.github/workflows/build_package.yml @@ -29,6 +29,9 @@ on: cloud_runner: required: false type: string + maxtext_sha: + required: false + type: string outputs: maxtext_sha: description: "MaxText short SHA used for the build" @@ -45,12 +48,15 @@ jobs: steps: - name: Checkout MaxText uses: actions/checkout@v5 + with: + ref: ${{ inputs.maxtext_sha || github.sha }} - name: Get metadata id: vars shell: bash run: | # MaxText SHA used to build the package - echo "maxtext_sha=${GITHUB_SHA}" >> $GITHUB_OUTPUT + MAXTEXT_SHA="${{ inputs.maxtext_sha || github.sha }}" + echo "maxtext_sha=${MAXTEXT_SHA}" >> $GITHUB_OUTPUT - name: Install build tools run: | python -m pip install --upgrade pip build uv diff --git a/.github/workflows/check_docs_build.yml b/.github/workflows/check_docs_build.yml index 6136256530..a3030ea8ad 100644 --- a/.github/workflows/check_docs_build.yml +++ b/.github/workflows/check_docs_build.yml @@ -20,6 +20,10 @@ on: pull_request: types: [opened, synchronize] workflow_call: + inputs: + maxtext_sha: + required: true + type: string workflow_dispatch: permissions: @@ -34,10 +38,18 @@ jobs: with: persist-credentials: false fetch-depth: 0 + ref: ${{ inputs.maxtext_sha }} - name: Check if only documentation changed id: check run: | + # Force build_docs=true for non-pull-request events + if [ "${{ github.event_name }}" != "pull_request" ]; then + echo "Not a pull request, enabling docs build." + echo "build_docs=true" >> $GITHUB_OUTPUT + exit 0 + fi + git fetch origin ${GITHUB_BASE_REF} CHANGED_FILES=$(git diff --name-only origin/${GITHUB_BASE_REF}...HEAD) diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml index 095f885c28..c0908c4cdd 100644 --- a/.github/workflows/ci_pipeline.yml +++ b/.github/workflows/ci_pipeline.yml @@ -20,6 +20,15 @@ name: MaxText Package Tests on: pull_request: workflow_call: + inputs: + maxtext_sha: + description: 'The specific MaxText commit SHA.' + required: true + type: string + outputs: + tests_result: + description: 'The result of all_tests_passed and all_notebooks_passed jobs.' + value: ${{ (jobs.all_tests_passed.result == 'success' && jobs.all_notebooks_passed.result == 'success') && 'success' || 'failure' }} workflow_dispatch: schedule: # Run the job every 4 hours @@ -48,6 +57,7 @@ jobs: - uses: actions/checkout@v4 with: fetch-depth: 0 + ref: ${{ inputs.maxtext_sha || github.sha }} - name: Check for Code Changes id: check run: | @@ -127,6 +137,7 @@ jobs: device_type: tpu device_name: v4-8 cloud_runner: linux-x86-n2-16-buildkit + maxtext_sha: ${{ inputs.maxtext_sha || github.sha }} maxtext_jupyter_notebooks: needs: build_and_upload_maxtext_package diff --git a/.github/workflows/docs_link_check.yml b/.github/workflows/docs_link_check.yml index c865dc5b43..97ce3b8903 100644 --- a/.github/workflows/docs_link_check.yml +++ b/.github/workflows/docs_link_check.yml @@ -21,6 +21,10 @@ on: # Run every Monday at 9:00 AM UTC - cron: '0 9 * * 1' workflow_call: + inputs: + maxtext_sha: + required: true + type: string workflow_dispatch: # Allow manual triggering permissions: @@ -37,6 +41,7 @@ jobs: with: persist-credentials: false fetch-depth: 0 + ref: ${{ inputs.maxtext_sha }} - name: Install uv and set the Python version uses: astral-sh/setup-uv@eb1897b8dc4b5d5bfe39a428a8f2304605e0983c # v7.0.0 diff --git a/.github/workflows/pypi_release.yml b/.github/workflows/pypi_release.yml index 6bd5126dfa..0436df1811 100644 --- a/.github/workflows/pypi_release.yml +++ b/.github/workflows/pypi_release.yml @@ -12,127 +12,152 @@ # See the License for the specific language governing permissions and # limitations under the License. -# This workflow will build, test and automatically release MaxText package to PyPI using Trusted Publishing (OIDC). +# This workflow publishes MaxText to PyPI once the Airflow E2E DAG reports success. name: Publish MaxText to PyPI -# Triggers when a new "release" is published in the GitHub UI +# Triggered by Airflow when E2E DAG completes or can be manually triggered via workflow_dispatch. on: - release: - types: [published] + repository_dispatch: + types: [airflow-dag-complete] + workflow_dispatch: + inputs: + run_id: + description: 'GitHub workflow run_id passed to the Airflow DAG run.' + required: true + type: string + maxtext_sha: + description: 'MaxText SHA to release to PyPI.' + required: true + type: string + publish: + description: 'Set to true to publish to PyPI.' + required: false + default: 'false' + type: boolean permissions: contents: read - issues: write - id-token: write + id-token: write # required for PyPI Trusted Publishing (OIDC) jobs: + handle_result: + name: Handle Airflow DAG Result + runs-on: ubuntu-latest + # Only run this job if the event that triggered the workflow is 'repository_dispatch' + if: github.event_name == 'repository_dispatch' + steps: + - name: Report DAG result + run: | + STATE="${{ github.event.client_payload.state }}" + DAG_ID="${{ github.event.client_payload.dag_id }}" + DAG_RUN_ID="${{ github.event.client_payload.dag_run_id }}" + SHA="${{ github.event.client_payload.sha }}" + + echo "================================" + echo "DAG ID: ${DAG_ID}" + echo "DAG Run ID: ${DAG_RUN_ID}" + echo "Commit SHA: ${SHA}" + echo "State: ${STATE}" + echo "================================" + + # Evaluate state + case "${STATE}" in + success) + echo "DAG '${DAG_ID}' completed successfully." + exit 0 + ;; + failed|upstream_failed) + echo "DAG '${DAG_ID}' failed with state: ${STATE}." + exit 1 + ;; + *) + echo "DAG '${DAG_ID}' ended with unexpected state: ${STATE}." + exit 1 + ;; + esac + release_approval: name: Approve Release runs-on: ubuntu-latest + needs: handle_result + if: | + always() && + (needs.handle_result.result == 'success' || needs.handle_result.result == 'skipped') # "release" environment is configured in MaxText repository settings. # This environment requires manual approval before proceeding to the next job. environment: release steps: - name: Acknowledge Approval - run: echo "Release approved, proceeding to build, test and publish MaxText package." - - build_and_test_maxtext_package: - name: Build and Test MaxText Package - needs: [release_approval] - uses: ./.github/workflows/ci_pipeline.yml - secrets: inherit - - build_docs: - name: Build Documentation - needs: [build_and_test_maxtext_package] - uses: ./.github/workflows/check_docs_build.yml - secrets: inherit - - check_links: - name: Check Links in Documentation - uses: ./.github/workflows/docs_link_check.yml - secrets: inherit + run: echo "Release approved, proceeding to publishing MaxText package." - publish_maxtext_package_to_pypi: - name: Publish MaxText package to PyPI - needs: [build_and_test_maxtext_package] + build_maxtext_package: + needs: release_approval + if: needs.release_approval.result == 'success' + uses: ./.github/workflows/build_package.yml + with: + device_type: tpu + device_name: v4-8 + cloud_runner: linux-x86-n2-16-buildkit + # Build package from a particular commit to ensure we are publishing the same code that was tested + maxtext_sha: ${{ github.event.client_payload.sha || inputs.maxtext_sha }} + + publish_maxtext_to_pypi: + name: Publish MaxText to PyPI + needs: release_approval runs-on: ubuntu-latest environment: release + if: github.event_name == 'repository_dispatch' || github.event.inputs.publish == 'true' steps: - - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 - - name: Download MaxText wheel - uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 - with: - name: maxtext-wheel - path: dist/ - - name: Publish MaxText wheel to PyPI - # Official action for PyPI Trusted Publishing - uses: pypa/gh-action-pypi-publish@release/v1 - with: - packages-dir: dist/ + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 + with: + name: maxtext-wheel + - name: Publish MaxText wheel to PyPI + 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] + needs: publish_maxtext_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 + latest_version=$(echo "$pypi_json" | python3 -c "import sys,json; print(json.load(sys.stdin)['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 "Latest MaxText version on PyPI: $latest_version" 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 }} + promote_docker_images: + name: Promote ${{ matrix.name }} Docker Image needs: [get_latest_maxtext_pypi_version] + runs-on: linux-x86-n2-16-buildkit + container: google/cloud-sdk:524.0.0 strategy: fail-fast: false matrix: include: - - device: tpu - build_mode: stable + - name: 'TPU Pre-Training 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 + - name: 'TPU Post-Training 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 }} + steps: + - name: Configure Docker + run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q + - name: Add tags to Docker image + shell: bash + run: | + SOURCE_IMAGE="gcr.io/${{ vars.PROJECT_NAME }}/${{ matrix.image_name }}" + GITHUB_RUN_ID="${{ github.event.client_payload.github_run_id }}" + gcloud container images add-tag \ + "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \ + "${SOURCE_IMAGE}:${{ needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version }}" \ + --quiet diff --git a/.github/workflows/release_pipeline.yml b/.github/workflows/release_pipeline.yml index c5da61988f..3cabfd118a 100644 --- a/.github/workflows/release_pipeline.yml +++ b/.github/workflows/release_pipeline.yml @@ -17,18 +17,44 @@ name: MaxText Release Pipeline -# Triggers when a new "release" is published in the GitHub UI +# Triggers when a new "release" is published in the GitHub UI or when manually triggered via workflow_dispatch. on: release: types: [published] workflow_dispatch: + inputs: + release_tag: + description: 'MaxText Release tag to build and test (e.g., maxtext-v0.2.3).' + required: true + type: string permissions: contents: read issues: write id-token: write + actions: read + pull-requests: write jobs: + get_maxtext_sha: + name: Resolve MaxText SHA + runs-on: ubuntu-latest + outputs: + maxtext_sha: ${{ steps.vars.outputs.maxtext_sha }} + steps: + - name: Checkout MaxText + uses: actions/checkout@v5 + with: + ref: ${{ inputs.release_tag || github.event.release.tag_name }} + fetch-depth: 1 + - name: Resolve SHA + id: vars + shell: bash + run: | + MAXTEXT_SHA=$(git rev-parse HEAD) + echo "Using MaxText SHA: ${MAXTEXT_SHA}" + echo "maxtext_sha=${MAXTEXT_SHA}" >> $GITHUB_OUTPUT + release_approval: name: Approve Release Candidate if: github.event_name != 'workflow_dispatch' @@ -42,41 +68,89 @@ jobs: build_and_test_maxtext_package: name: Build and Test MaxText Package - needs: [release_approval] + needs: [release_approval, get_maxtext_sha] if: | always() && (needs.release_approval.result == 'success' || needs.release_approval.result == 'skipped') uses: ./.github/workflows/ci_pipeline.yml + with: + maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} secrets: inherit build_release_candidate_images: - name: ${{ matrix.image_name }} - needs: [build_and_test_maxtext_package] + name: Build ${{ matrix.name }} Docker Image + needs: [build_and_test_maxtext_package, get_maxtext_sha] + # Ensure the SHA resolved, and CI pipeline tests are not failed or cancelled if: | always() && - (needs.build_and_test_maxtext_package.result == 'success' || needs.build_and_test_maxtext_package.result == 'skipped') + needs.get_maxtext_sha.result == 'success' && + needs.build_and_test_maxtext_package.outputs.tests_result == 'success' strategy: fail-fast: false matrix: include: - - device: tpu - build_mode: stable + - name: 'TPU Pre-Training Stable' + device: tpu image_name: maxtext_jax_stable workflow: pre-training dockerfile: maxtext_tpu_dependencies.Dockerfile - - device: tpu - build_mode: stable + - name: 'TPU Post-Training Stable' + device: tpu image_name: maxtext_post_training_stable workflow: post-training dockerfile: maxtext_tpu_dependencies.Dockerfile + - name: 'GPU Pre-Training Stable' + device: gpu + image_name: maxtext_gpu_jax_stable + workflow: pre-training + dockerfile: maxtext_gpu_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 }} + build_mode: stable workflow: ${{ matrix.workflow }} dockerfile: ${{ matrix.dockerfile }} - maxtext_sha: ${{ github.sha }} - include_test_assets: true + maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} secrets: HF_TOKEN: ${{ secrets.HF_TOKEN }} + + run_e2e_tests: + name: Run E2E tests + needs: [build_release_candidate_images, get_maxtext_sha] + # Ensure the SHA resolved, the build job itself wasn't skipped or cancelled, and all builds succeeded + if: | + always() && + needs.get_maxtext_sha.result == 'success' && + needs.build_release_candidate_images.result != 'skipped' && + needs.build_release_candidate_images.result != 'cancelled' && + !contains(needs.build_release_candidate_images.outputs.*.build_result, 'failure') && + !contains(needs.build_release_candidate_images.outputs.*.build_result, 'cancelled') + uses: ./.github/workflows/run_e2e_tests.yml + with: + mode: stable + run_id: ${{ github.run_id }} + maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} + secrets: inherit + + build_docs: + name: Build Documentation + needs: [release_approval, get_maxtext_sha] + if: | + always() && + (needs.release_approval.result == 'success' || needs.release_approval.result == 'skipped') + uses: ./.github/workflows/check_docs_build.yml + with: + maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} + secrets: inherit + + check_links: + name: Check Links in Documentation + needs: [release_approval, get_maxtext_sha] + if: | + always() && + (needs.release_approval.result == 'success' || needs.release_approval.result == 'skipped') + uses: ./.github/workflows/docs_link_check.yml + with: + maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} + secrets: inherit diff --git a/.github/workflows/run_e2e_tests.yml b/.github/workflows/run_e2e_tests.yml index 5e76bfb774..94c6d8d373 100644 --- a/.github/workflows/run_e2e_tests.yml +++ b/.github/workflows/run_e2e_tests.yml @@ -21,9 +21,8 @@ on: inputs: mode: description: 'Build mode: stable or nightly' - required: false + required: true type: string - default: 'stable' run_id: description: 'GitHub workflow run_id to be passed to the Airflow DAG run.' required: true