From ba01c9ec25a94704b390ab9017df1c17c9e32364 Mon Sep 17 00:00:00 2001 From: Surbhi Jain Date: Fri, 21 Aug 2026 17:25:31 +0000 Subject: [PATCH] CI: Streamline release pipeline and Docker image verification --- .github/workflows/ci_pipeline.yml | 3 +- .github/workflows/promote_docker_image.yml | 57 +++++---- .github/workflows/pypi_release.yml | 76 ++++-------- .github/workflows/release_pipeline.yml | 110 ++++++++++++++++-- .github/workflows/run_ci_tests.yml | 19 +-- .../workflows/tpu_docker_images_pipeline.yml | 3 - 6 files changed, 162 insertions(+), 106 deletions(-) diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml index f0ce827ebb..98db52fb6f 100644 --- a/.github/workflows/ci_pipeline.yml +++ b/.github/workflows/ci_pipeline.yml @@ -183,8 +183,7 @@ jobs: needs: [gate_test_run] if: | always() && - needs.gate_test_run.result == 'success' && - github.ref == 'refs/heads/main' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') + needs.gate_test_run.result == 'success' && github.event_name != 'pull_request' uses: ./.github/workflows/run_tests_coordinator.yml strategy: fail-fast: false diff --git a/.github/workflows/promote_docker_image.yml b/.github/workflows/promote_docker_image.yml index 78b037dacf..35f0cf5d7b 100644 --- a/.github/workflows/promote_docker_image.yml +++ b/.github/workflows/promote_docker_image.yml @@ -18,6 +18,7 @@ # 2. If the DAG succeeded, it promotes the Docker images by tagging the artifact with the latest tag. name: Promote MaxText Docker Image +run-name: "Triggered by Airflow ${{ github.event.client_payload.test_type }} Tests DAG" on: repository_dispatch: @@ -39,11 +40,11 @@ jobs: DAG_ID: ${{ github.event.client_payload.dag_id }} DAG_RUN_ID: ${{ github.event.client_payload.dag_run_id }} SHA: ${{ github.event.client_payload.sha }} - GITHUB_RUN_ID: ${{ github.event.client_payload.github_run_id }} + CALLER_RUN_ID: ${{ github.event.client_payload.github_run_id }} TEST_TYPE: ${{ github.event.client_payload.test_type }} run: | echo "================================" - echo "Github Run ID: ${GITHUB_RUN_ID}" + echo "Github Run ID: ${CALLER_RUN_ID}" echo "DAG ID: ${DAG_ID}" echo "DAG Run ID: ${DAG_RUN_ID}" echo "Commit SHA: ${SHA}" @@ -57,16 +58,16 @@ jobs: *) echo "DAG ended with unexpected state: ${STATE}"; exit 1 ;; esac - if [ -n "$GITHUB_RUN_ID" ]; then - echo "Checking nightly build status for run: $GITHUB_RUN_ID" - gh run watch "$GITHUB_RUN_ID" --exit-status --interval 60 --repo "$GITHUB_REPOSITORY" + if [ -n "$CALLER_RUN_ID" ]; then + echo "Checking nightly build status for run: $CALLER_RUN_ID" + gh run watch "$CALLER_RUN_ID" --exit-status --interval 60 --repo "$GITHUB_REPOSITORY" else echo "Error: No Github Run ID provided. Cannot verify nightly build status." exit 1 fi tag_docker_image: - name: Promote ${{ matrix.image_name }} Docker Image + name: Promote ${{ matrix.test_type }} Docker Image needs: handle_result if: needs.handle_result.result == 'success' runs-on: linux-x86-n2-16-buildkit @@ -76,13 +77,11 @@ jobs: matrix: include: - test_type: pre_training - image_name: maxtext_jax_nightly + stable_image_name: maxtext_jax_stable + nightly_image_name: maxtext_jax_nightly - test_type: post_training - image_name: maxtext_post_training_nightly - - test_type: pre_training - image_name: maxtext_jax_stable - - test_type: post_training - image_name: maxtext_post_training_stable + stable_image_name: maxtext_post_training_stable + nightly_image_name: maxtext_post_training_nightly steps: - name: Configure Docker if: ${{ github.event.client_payload.test_type == '' || github.event.client_payload.test_type == matrix.test_type }} @@ -91,15 +90,33 @@ jobs: if: ${{ github.event.client_payload.test_type == '' || github.event.client_payload.test_type == matrix.test_type }} shell: bash env: - GITHUB_RUN_ID: ${{ github.event.client_payload.github_run_id }} + CALLER_RUN_ID: ${{ github.event.client_payload.github_run_id }} PROJECT_NAME: ${{ vars.PROJECT_NAME }} - IMAGE_NAME: ${{ matrix.image_name }} + STABLE_IMAGE_NAME: ${{ matrix.stable_image_name }} + NIGHTLY_IMAGE_NAME: ${{ matrix.nightly_image_name }} run: | - SOURCE_IMAGE="gcr.io/${PROJECT_NAME}/${IMAGE_NAME}" + tag_image() { + local img="$1" + local tag="$2" + echo "Tagging ${img}:${CALLER_RUN_ID} as ${tag}..." + gcloud container images add-tag "${img}:${CALLER_RUN_ID}" "${img}:${tag}" --quiet + } - # Add the traceability tag to confirm it passed validation suite - gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \ - "${SOURCE_IMAGE}:verified-${GITHUB_RUN_ID}" --quiet + # 1. Check if stable image exists and tag it + STABLE_IMAGE="gcr.io/${PROJECT_NAME}/${STABLE_IMAGE_NAME}" + if gcloud container images describe "${STABLE_IMAGE}:${CALLER_RUN_ID}" >/dev/null 2>&1; then + echo "Stable image ${STABLE_IMAGE}:${CALLER_RUN_ID} exists. Tagging stable image..." + tag_image "${STABLE_IMAGE}" "verified-e2e-${CALLER_RUN_ID}" + tag_image "${STABLE_IMAGE}" "latest" + else + echo "Stable image ${STABLE_IMAGE}:${CALLER_RUN_ID} does not exist. Skipping stable image tagging." + fi - # Add "latest" tag - gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" "${SOURCE_IMAGE}:latest" --quiet + # 2. Check if nightly image exists and tag it + NIGHTLY_IMAGE="gcr.io/${PROJECT_NAME}/${NIGHTLY_IMAGE_NAME}" + if gcloud container images describe "${NIGHTLY_IMAGE}:${CALLER_RUN_ID}" >/dev/null 2>&1; then + echo "Nightly image ${NIGHTLY_IMAGE}:${CALLER_RUN_ID} exists. Tagging nightly image..." + tag_image "${NIGHTLY_IMAGE}" "latest" + else + echo "Nightly image ${NIGHTLY_IMAGE}:${CALLER_RUN_ID} does not exist. Skipping nightly image tagging." + fi diff --git a/.github/workflows/pypi_release.yml b/.github/workflows/pypi_release.yml index 4d1e067ad1..62c504e6a3 100644 --- a/.github/workflows/pypi_release.yml +++ b/.github/workflows/pypi_release.yml @@ -16,10 +16,23 @@ name: Publish MaxText to PyPI -# Triggered by Airflow when E2E DAG completes or can be manually triggered via workflow_dispatch. +# Triggered by release_pipeline.yml via workflow_call or manually via workflow_dispatch. on: - repository_dispatch: - types: [airflow-dag-complete] + workflow_call: + 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: true + type: boolean workflow_dispatch: inputs: run_id: @@ -40,77 +53,33 @@ permissions: contents: read 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 - env: - 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 }} - run: | - - 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') && - (github.event_name == 'repository_dispatch' || inputs.publish) + if: inputs.publish == true # "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 publishing MaxText package." + - name: Acknowledge Approval + run: echo "Release approved, proceeding to publishing MaxText package." build_maxtext_package: needs: release_approval - if: | - always() && - (needs.release_approval.result == 'success' || needs.release_approval.result == 'skipped') 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 }} + maxtext_sha: ${{ inputs.maxtext_sha }} publish_maxtext_to_pypi: name: Publish MaxText to PyPI needs: [release_approval, build_maxtext_package] runs-on: ubuntu-latest environment: release - if: needs.build_maxtext_package.result == 'success' && (github.event_name == 'repository_dispatch' || inputs.publish) + if: needs.build_maxtext_package.result == 'success' && inputs.publish permissions: id-token: write # required for PyPI Trusted Publishing (OIDC) contents: read @@ -164,11 +133,10 @@ jobs: env: PROJECT_NAME: ${{ vars.PROJECT_NAME }} IMAGE_NAME: ${{ matrix.image_name }} - RUN_ID_INPUT: ${{ github.event_name == 'workflow_dispatch' && inputs.run_id || github.event.client_payload.github_run_id }} + GITHUB_RUN_ID: ${{ inputs.run_id }} PYPI_VERSION: ${{ needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version }} run: | SOURCE_IMAGE="gcr.io/${PROJECT_NAME}/${IMAGE_NAME}" - GITHUB_RUN_ID="${RUN_ID_INPUT}" gcloud container images add-tag \ "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \ "${SOURCE_IMAGE}:${PYPI_VERSION}" \ diff --git a/.github/workflows/release_pipeline.yml b/.github/workflows/release_pipeline.yml index bb4bd17a5d..877d0a6ae9 100644 --- a/.github/workflows/release_pipeline.yml +++ b/.github/workflows/release_pipeline.yml @@ -64,18 +64,32 @@ jobs: - name: Acknowledge Approval run: echo "Release candidate approved, proceeding to build and test MaxText package." - build_and_test_maxtext_package: - name: Build and Test MaxText Package + build_maxtext_package: + name: Build MaxText Package needs: [release_approval, get_maxtext_sha] if: | always() && needs.get_maxtext_sha.result == 'success' && (needs.release_approval.result == 'success' || needs.release_approval.result == 'skipped') + uses: ./.github/workflows/build_package.yml + with: + device_type: tpu + device_name: v4-8 + cloud_runner: linux-x86-n2-16-buildkit + maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} + + run_ci_tests: + name: Test MaxText Package + needs: [build_maxtext_package, get_maxtext_sha] + if: | + always() && + needs.get_maxtext_sha.result == 'success' && + needs.build_maxtext_package.result == 'success' permissions: issues: write id-token: write pull-requests: write - contents: read + contents: write actions: read uses: ./.github/workflows/ci_pipeline.yml with: @@ -84,14 +98,13 @@ jobs: HF_TOKEN: ${{ secrets.HF_TOKEN }} GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} - build_release_candidate_images: + build_tpu_release_candidate_images: name: Build ${{ matrix.name }} Docker Image - needs: [build_and_test_maxtext_package, get_maxtext_sha] - # Ensure the SHA resolved, and CI pipeline tests succeeded + needs: [build_maxtext_package, get_maxtext_sha] if: | always() && needs.get_maxtext_sha.result == 'success' && - needs.build_and_test_maxtext_package.result == 'success' + needs.build_maxtext_package.result == 'success' strategy: fail-fast: false matrix: @@ -106,6 +119,28 @@ jobs: 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: stable + workflow: ${{ matrix.workflow }} + dockerfile: ${{ matrix.dockerfile }} + maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} + secrets: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + + build_gpu_release_candidate_images: + name: Build ${{ matrix.name }} Docker Image + needs: [build_maxtext_package, get_maxtext_sha] + if: | + always() && + needs.get_maxtext_sha.result == 'success' && + needs.build_maxtext_package.result == 'success' + strategy: + fail-fast: false + matrix: + include: - name: 'GPU Pre-Training Stable' device: gpu image_name: maxtext_gpu_jax_stable @@ -124,12 +159,12 @@ jobs: run_e2e_tests: name: Run E2E tests - needs: [build_release_candidate_images, get_maxtext_sha] + needs: [build_tpu_release_candidate_images, get_maxtext_sha] # Ensure the SHA resolved and all image builds succeeded if: | always() && needs.get_maxtext_sha.result == 'success' && - needs.build_release_candidate_images.result == 'success' + needs.build_tpu_release_candidate_images.result == 'success' uses: ./.github/workflows/run_e2e_tests.yml with: mode: stable @@ -162,3 +197,60 @@ jobs: uses: ./.github/workflows/docs_link_check.yml with: maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} + + wait_for_tests_verification: + name: Wait for Tests Verification + needs: [run_ci_tests, run_e2e_tests] + if: | + always() && + needs.run_ci_tests.result == 'success' && needs.run_e2e_tests.result == 'success' + runs-on: linux-x86-n2-16-buildkit + container: google/cloud-sdk:524.0.0 + steps: + - name: Configure Docker + run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q + - name: Poll for verified Docker images + env: + PROJECT_NAME: ${{ vars.PROJECT_NAME }} + GITHUB_RUN_ID: ${{ github.run_id }} + run: | + check_tag() { + gcloud container images describe "gcr.io/${PROJECT_NAME}/$1" >/dev/null 2>&1 + } + + echo "Waiting for pre-training and post-training tests to complete..." + while true; do + PRETRAIN_READY=false + POSTTRAIN_READY=false + + if check_tag "maxtext_jax_stable:verified-e2e-${GITHUB_RUN_ID}"; then + PRETRAIN_READY=true + fi + if check_tag "maxtext_post_training_stable:verified-e2e-${GITHUB_RUN_ID}"; then + POSTTRAIN_READY=true + fi + + if [ "$PRETRAIN_READY" = "true" ] && [ "$POSTTRAIN_READY" = "true" ]; then + echo "Both pre-training and post-training tests have passed. Proceeding to publish to PyPI." + break + fi + + echo "Either pre-training or post-training tests are still running. Checking again in 5 minutes..." + sleep 300 + done + + publish_to_pypi: + name: Publish MaxText to PyPI + needs: [wait_for_tests_verification, get_maxtext_sha] + if: | + always() && + needs.get_maxtext_sha.result == 'success' && + needs.wait_for_tests_verification.result == 'success' + permissions: + id-token: write + contents: read + uses: ./.github/workflows/pypi_release.yml + with: + maxtext_sha: ${{ needs.get_maxtext_sha.outputs.maxtext_sha }} + run_id: ${{ github.run_id }} + publish: true diff --git a/.github/workflows/run_ci_tests.yml b/.github/workflows/run_ci_tests.yml index e7d98ae956..7b4433fa07 100644 --- a/.github/workflows/run_ci_tests.yml +++ b/.github/workflows/run_ci_tests.yml @@ -31,11 +31,6 @@ on: workflow: required: true type: string - additional_image_to_tag: - description: 'Optional additional image to tag as latest if tests pass.' - required: false - type: string - default: '' permissions: contents: read @@ -97,7 +92,6 @@ jobs: PROJECT_NAME: ${{ vars.PROJECT_NAME }} IMAGE_NAME_EXPR: ${{ inputs.image_name }} IMAGE_TAG_EXPR: ${{ inputs.image_tag }} - ADDITIONAL_IMAGE_TO_TAG: ${{ inputs.additional_image_to_tag }} run: | image_name="$IMAGE_NAME_EXPR" image_tag="$IMAGE_TAG_EXPR" @@ -105,15 +99,4 @@ jobs: # Add the traceability tag to confirm it passed validation suite gcloud container images add-tag "${SOURCE_IMAGE}:${image_tag}" \ - "${SOURCE_IMAGE}:verified-${image_tag}" --quiet - - # TODO: This is a workaround, to be removed when promote_docker_image.yml workflow is stable - # Add "latest" tag - gcloud container images add-tag "${SOURCE_IMAGE}:${image_tag}" "${SOURCE_IMAGE}:latest" --quiet - - # Add "latest" tag to additional image if specified - if [ -n "$ADDITIONAL_IMAGE_TO_TAG" ]; then - ADDITIONAL_IMAGE="gcr.io/${PROJECT_NAME}/${ADDITIONAL_IMAGE_TO_TAG}" - echo "Tagging ${ADDITIONAL_IMAGE}:${image_tag} as ${ADDITIONAL_IMAGE}:latest" - gcloud container images add-tag "${ADDITIONAL_IMAGE}:${image_tag}" "${ADDITIONAL_IMAGE}:latest" --quiet - fi + "${SOURCE_IMAGE}:verified-ci-${image_tag}" --quiet diff --git a/.github/workflows/tpu_docker_images_pipeline.yml b/.github/workflows/tpu_docker_images_pipeline.yml index ab88393988..33c82f8659 100644 --- a/.github/workflows/tpu_docker_images_pipeline.yml +++ b/.github/workflows/tpu_docker_images_pipeline.yml @@ -86,10 +86,8 @@ jobs: - name: "Pre-Training" image_name: maxtext_jax_stable workflow: pre-training - additional_image_to_tag: maxtext_jax_nightly - name: "Post-Training" image_name: maxtext_post_training_stable - additional_image_to_tag: maxtext_post_training_nightly workflow: post-training uses: ./.github/workflows/run_ci_tests.yml with: @@ -97,7 +95,6 @@ jobs: image_tag: ${{ github.run_id }} device: tpu workflow: ${{ matrix.workflow }} - additional_image_to_tag: ${{ matrix.additional_image_to_tag != '' && (inputs.image_suffix != '' && format('{0}_{1}', matrix.additional_image_to_tag, inputs.image_suffix) || matrix.additional_image_to_tag) || '' }} run_e2e_tests: name: Run E2E tests