Publish MaxText to PyPI when Airflow E2E tests succeed

.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

.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

.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)
 

.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

.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