Add run_tests_coordinator

This file is called by build_and_test_maxtext.yml to run
tpu/cpu/gpu unit and integration tests.
It simplifies the caller's input.

Author: xibinliu

.github/workflows/build_and_push_docker_image.yml

@@ -41,10 +41,6 @@ on:
         required: false
         type: string
         default: ''
-      is_post_training:
-        required: false
-        type: boolean
-        default: false
 
 permissions:
   contents: read
@@ -87,7 +83,9 @@ jobs:
           ref: ${{ inputs.maxtext_sha }}
 
       - name: Checkout post-training dependencies
-        if: steps.check.outputs.should_run == 'true' && inputs.image_name == 'maxtext_post_training_nightly'
+        if: |
+          steps.check.outputs.should_run == 'true' &&
+          contains(inputs.image_name, 'post_training_nightly')
         run: |
           git clone https://github.com/google/tunix.git ./tunix
           git clone https://github.com/vllm-project/vllm.git ./vllm
@@ -115,8 +113,7 @@ jobs:
           push: true
           context: .
           file: ${{ inputs.dockerfile }}
-          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:latest
-          cache-from: type=gha
+          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ inputs.image_date }}-build-${{ github.run_id }}
           outputs: type=image,compression=zstd,force-compression=true
           build-args: |
             DEVICE=${{ inputs.device }}
@@ -131,27 +128,22 @@ jobs:
         shell: bash
         run: |
           SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${INPUTS_IMAGE_NAME}"
-
-          # Add date tag
-          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
+          TEMP_IMG="${SOURCE_IMAGE}:${INPUTS_IMAGE_DATE}-build-${{ github.run_id }}"
 
           # Convert date to YYYYMMDD format
           clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
 
           # Add MaxText tag
           maxtext_hash=$(git rev-parse --short HEAD)
-          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
-
+          gcloud container images add-tag "${TEMP_IMG}" "${SOURCE_IMAGE}:maxtext_${maxtext_hash}_${clean_date}" --quiet
 
           # Add post-training dependencies tags
-          if [ "${{ inputs.is_post_training }}" == "true" ]; then
-            for dir in tunix vllm tpu-inference; do
-              if [ -d "./$dir" ]; then
-                dir_hash=$(git -C "$dir" rev-parse --short HEAD)
-                gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
-              fi
-            done
-          fi
+          for dir in tunix vllm tpu-inference; do
+            if [ -d "./$dir" ]; then
+              dir_hash=$(git -C "$dir" rev-parse --short HEAD)
+              gcloud container images add-tag "${TEMP_IMG}" "${SOURCE_IMAGE}:${dir}_${dir_hash}_${clean_date}" --quiet
+            fi
+          done
         env:
           INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
-          INPUTS_IMAGE_DATE: ${{ inputs.image_date }}
+          INPUTS_IMAGE_DATE: ${{ inputs.image_date }}

.github/workflows/build_and_test_maxtext.yml

@@ -113,72 +113,47 @@ jobs:
     with:
       device_type: tpu
       device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
+      base_image: maxtext-unit-test-tpu:${{ matrix.image_type }}
       cloud_runner: linux-x86-ct6e-180-4tpu
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
-  maxtext_cpu_unit_tests:
-    needs: build_and_upload_maxtext_package
+  tpu-tests:
+    needs: [build_and_upload_maxtext_package]
     if: needs.doc_only_check.outputs.run_tests == 'true'
-    uses: ./.github/workflows/run_tests_against_package.yml
     strategy:
-        fail-fast: false # don't cancel all jobs on failure
-        matrix:
-          image_type: ["py312"]
-          worker_group: [1, 2]
+      fail-fast: false
+      matrix:
+        flavor: [tpu-unit, tpu-integration]
+    uses: ./.github/workflows/run_tests_coordinator.yml
     with:
-      device_type: cpu
-      device_name: X64
-      cloud_runner: linux-x86-n2-16
-      image_type: ${{ matrix.image_type }}
-      pytest_marker: 'cpu_only'
-      xla_python_client_mem_fraction: 0.75
-      tf_force_gpu_allow_growth: false
-      container_resource_option: "--privileged"
+      flavor: ${{ matrix.flavor }}
+      base_image: maxtext-unit-test-tpu:py312
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
-      worker_group: ${{ matrix.worker_group }}
-      total_workers: 2
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
-  maxtext_tpu_unit_tests:
-    needs: build_and_upload_maxtext_package
+  gpu-tests:
+    needs: [build_and_upload_maxtext_package]
     if: needs.doc_only_check.outputs.run_tests == 'true'
-    uses: ./.github/workflows/run_tests_against_package.yml
     strategy:
-        fail-fast: false
-        matrix:
-          image_type: ["py312"]
+      fail-fast: false
+      matrix:
+        flavor: [gpu-unit, gpu-integration]
+    uses: ./.github/workflows/run_tests_coordinator.yml
     with:
-      device_type: tpu
-      device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
-      cloud_runner: linux-x86-ct6e-180-4tpu
-      pytest_marker: 'not cpu_only and not gpu_only and not integration_test'
-      xla_python_client_mem_fraction: 0.75
-      tf_force_gpu_allow_growth: false
-      container_resource_option: "--privileged"
+      flavor: ${{ matrix.flavor }}
+      base_image: maxtext-unit-test-cuda12:py312
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
-  maxtext_tpu_integration_tests:
-    needs: build_and_upload_maxtext_package
+  cpu-tests:
+    needs: [build_and_upload_maxtext_package]
     if: needs.doc_only_check.outputs.run_tests == 'true'
-    uses: ./.github/workflows/run_tests_against_package.yml
-    strategy:
-        fail-fast: false
-        matrix:
-          image_type: ["py312"]
+    uses: ./.github/workflows/run_tests_coordinator.yml
     with:
-      device_type: tpu
-      device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
-      cloud_runner: linux-x86-ct6e-180-4tpu
-      pytest_marker: 'not cpu_only and not gpu_only and integration_test'
-      xla_python_client_mem_fraction: 0.75
-      tf_force_gpu_allow_growth: false
-      container_resource_option: "--privileged"
+      flavor: cpu-unit
+      base_image: maxtext-unit-test-tpu:py312
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
@@ -188,12 +163,10 @@ jobs:
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
-        matrix:
-          image_type: ["py312"]
     with:
       device_type: tpu
       device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
+      base_image: maxtext-unit-test-tpu:py312
       cloud_runner: linux-x86-ct6e-180-4tpu
       pytest_marker: 'not cpu_only and not gpu_only and not integration_test'
       xla_python_client_mem_fraction: 0.75
@@ -208,12 +181,10 @@ jobs:
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
-        matrix:
-          image_type: ["py312"]
     with:
       device_type: tpu
       device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
+      base_image: maxtext-unit-test-tpu:py312
       cloud_runner: linux-x86-ct6e-180-4tpu
       pytest_marker: 'not cpu_only and not gpu_only and integration_test'
       xla_python_client_mem_fraction: 0.75
@@ -222,51 +193,9 @@ jobs:
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
-  maxtext_gpu_unit_tests:
-    needs: build_and_upload_maxtext_package
-    if: needs.doc_only_check.outputs.run_tests == 'true'
-    uses: ./.github/workflows/run_tests_against_package.yml
-    strategy:
-        fail-fast: false
-        matrix:
-          image_type: ["py312"]
-          cuda: ["cuda12"]
-    with:
-      device_type: ${{ matrix.cuda }}
-      device_name: a100-40gb-4
-      image_type: ${{ matrix.image_type }}
-      cloud_runner: linux-x86-a2-48-a100-4gpu
-      pytest_marker: 'not cpu_only and not tpu_only and not integration_test'
-      xla_python_client_mem_fraction: 0.65
-      tf_force_gpu_allow_growth: true
-      container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
-      is_scheduled_run: ${{ github.event_name == 'schedule' }}
-      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
-
-  maxtext_gpu_integration_tests:
-    needs: build_and_upload_maxtext_package
-    if: needs.doc_only_check.outputs.run_tests == 'true'
-    uses: ./.github/workflows/run_tests_against_package.yml
-    strategy:
-        fail-fast: false
-        matrix:
-          image_type: ["py312"]
-          cuda: ["cuda12"]
-    with:
-      device_type: ${{ matrix.cuda }}
-      device_name: a100-40gb-4
-      image_type: ${{ matrix.image_type }}
-      cloud_runner: linux-x86-a2-48-a100-4gpu
-      pytest_marker: 'not cpu_only and not tpu_only and integration_test'
-      xla_python_client_mem_fraction: 0.65
-      tf_force_gpu_allow_growth: true
-      container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
-      is_scheduled_run: ${{ github.event_name == 'schedule' }}
-      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
-
   all_tests_passed:
     name: All Required Tests Passed
-    needs: [doc_only_check, build_and_upload_maxtext_package, maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests]
+    needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests]
     if: always()
     runs-on: ubuntu-latest
     steps:
@@ -280,13 +209,11 @@ jobs:
 
           # Otherwise, check that build and all tests passed or were skipped
           echo "Build result: ${NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT}"
-          echo "CPU tests: ${NEEDS_MAXTEXT_CPU_UNIT_TESTS_RESULT}"
-          echo "TPU tests: ${NEEDS_MAXTEXT_TPU_UNIT_TESTS_RESULT}"
-          echo "TPU integration: ${NEEDS_MAXTEXT_TPU_INTEGRATION_TESTS_RESULT}"
-          echo "TPU pathways: ${NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT}"
-          echo "TPU pathways integration: ${NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT}"
-          echo "GPU tests: ${NEEDS_MAXTEXT_GPU_UNIT_TESTS_RESULT}"
-          echo "GPU integration: ${NEEDS_MAXTEXT_GPU_INTEGRATION_TESTS_RESULT}"
+          echo "TPU Tests (Matrix) result: ${NEEDS_TPU_TESTS_RESULT}"
+          echo "GPU Tests (Matrix) result: ${NEEDS_GPU_TESTS_RESULT}"
+          echo "CPU Tests (Matrix) result: ${NEEDS_CPU_TESTS_RESULT}"
+          echo "Pathways Unit result: ${NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT}"
+          echo "Pathways Integration result: ${NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT}"
 
           # Fail only if any job failed or was cancelled (skipped is OK)
           if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
@@ -298,13 +225,11 @@ jobs:
         env:
           NEEDS_DOC_ONLY_CHECK_OUTPUTS_RUN_TESTS: ${{ needs.doc_only_check.outputs.run_tests }}
           NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }}
-          NEEDS_MAXTEXT_CPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_cpu_unit_tests.result }}
-          NEEDS_MAXTEXT_TPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_unit_tests.result }}
-          NEEDS_MAXTEXT_TPU_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_integration_tests.result }}
+          NEEDS_CPU_TESTS_RESULT: ${{ needs.cpu-tests.result }}
+          NEEDS_TPU_TESTS_RESULT: ${{ needs.tpu-tests.result }}
+          NEEDS_GPU_TESTS_RESULT: ${{ needs.gpu-tests.result }}
           NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_unit_tests.result }}
           NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_integration_tests.result }}
-          NEEDS_MAXTEXT_GPU_UNIT_TESTS_RESULT: ${{ needs.maxtext_gpu_unit_tests.result }}
-          NEEDS_MAXTEXT_GPU_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_gpu_integration_tests.result }}
 
   all_notebooks_passed:
     name: All Notebooks Passed
@@ -337,14 +262,14 @@ jobs:
 
   notify_failure:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build
-    needs: [maxtext_jupyter_notebooks, maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests]
+    needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:
       issues: write
     steps:
-    - name: Check whether one of the jobs failed
-      if: ${{ contains(needs.*.result, 'failure') && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }}
-      uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
-      with:
-        github-token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Check whether one of the jobs failed
+        if: ${{ contains(needs.*.result, 'failure') && github.event_name == 'schedule' }}
+        uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/run_jupyter_notebooks.yml

@@ -25,15 +25,20 @@ on:
       device_name:
         required: true
         type: string
-      image_type:
-        required: false
+      base_image:
+        required: true
         type: string
       cloud_runner:
         required: false
         type: string
       maxtext_sha:
-        required: true
+        required: false
         type: string
+      # Flag to skip source checkout and wheel installation
+      maxtext_installed:
+        required: false
+        type: boolean
+        default: false
     secrets:
       HF_TOKEN:
         required: true
@@ -44,17 +49,20 @@ jobs:
   run:
     runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
     container:
-      image: gcr.io/tpu-prod-env-multipod/maxtext-unit-test-${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}:${{ inputs.image_type != '' && inputs.image_type }}
+      image: gcr.io/tpu-prod-env-multipod/${{ inputs.base_image }}
     steps:
       - name: Checkout MaxText
+        if: ${{ !inputs.maxtext_installed }}
         uses: actions/checkout@v5
         with:
           ref: ${{ inputs.maxtext_sha }}
       - name: Download the MaxText wheel
+        if: ${{ !inputs.maxtext_installed }}
         uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0
         with:
           name: maxtext-wheel
       - name: Install MaxText and Dependencies
+        if: ${{ !inputs.maxtext_installed }}
         shell: bash
         run: |
           # 1. Create virtual environment
@@ -75,13 +83,28 @@ jobs:
         env:
           PYTHONPATH: "${{ github.workspace }}/src"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          MAXTEXT_INSTALLED: ${{ inputs.maxtext_installed }}
         run: |
-          source .venv/bin/activate
+          if [ "${MAXTEXT_INSTALLED}" == "true" ]; then
+            # Move to the directory where code is baked into the image. See the Dockerfile.
+            # This is necessary because GHA sets an empty workspace by default.
+            cd /deps
+            PYTHON_EXE="python3"
+            PAPERMILL_EXE="papermill"
+          else
+            PYTHON_EXE=".venv/bin/python3"
+            PAPERMILL_EXE=".venv/bin/papermill"
+            source .venv/bin/activate
+          fi
 
           export MAXTEXT_REPO_ROOT=$(pwd)
           export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
           export MAXTEXT_NOTEBOOKS_ROOT="$MAXTEXT_REPO_ROOT/src/maxtext/examples"
 
+          # Install dependencies for running notebooks
+          $PYTHON_EXE -m pip install papermill ipykernel ipywidgets
+          $PYTHON_EXE -m ipykernel install --user --name maxtext_venv
+
           for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/{sft,rl}*.ipynb; do
             filename=$(basename "$notebook")
             if [[ "$filename" == "sft_qwen3_demo.ipynb" ]]; then
@@ -94,7 +117,7 @@ jobs:
             echo "Running $filename ..."
             echo "------------------------------------------------------"
 
-            papermill "$notebook" "$output_name" -k maxtext_venv
+            $PAPERMILL_EXE "$notebook" "$output_name" -k maxtext_venv
           done
       - name: Record Commit IDs
         shell: bash