Skip to content

CI: Streamline release pipeline and Docker image verification #4965

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SurbhiJainUSC wants to merge 1 commit into
base: main
Choose a base branch
from
Open

CI: Streamline release pipeline and Docker image verification #4965

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions .github/workflows/ci_pipeline.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
57 changes: 37 additions & 20 deletions .github/workflows/promote_docker_image.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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}"
Expand All @@ -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
Expand All @@ -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 }}
Expand All @@ -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
76 changes: 22 additions & 54 deletions .github/workflows/pypi_release.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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}" \
Expand Down
110 changes: 101 additions & 9 deletions .github/workflows/release_pipeline.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Loading
Loading