Revert "Add run_tests_coordinator"

This reverts commit 8f501bd.

Author: Charles Li

.github/workflows/build_and_push_docker_image.yml

@@ -41,6 +41,10 @@ on:
         required: false
         type: string
         default: ''
+      is_post_training:
+        required: false
+        type: boolean
+        default: false
 
 permissions:
   contents: read
@@ -78,9 +82,7 @@ jobs:
           ref: ${{ inputs.maxtext_sha }}
 
       - name: Checkout post-training dependencies
-        if: |
-          steps.check.outputs.should_run == 'true' &&
-          contains(inputs.image_name, 'post_training_nightly')
+        if: steps.check.outputs.should_run == 'true' && inputs.image_name == 'maxtext_post_training_nightly'
         run: |
           git clone https://github.com/google/tunix.git ./tunix
           git clone https://github.com/vllm-project/vllm.git ./vllm
@@ -108,7 +110,8 @@ jobs:
           push: true
           context: .
           file: ${{ inputs.dockerfile }}
-          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ inputs.image_date }}-build-${{ github.run_id }}
+          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:latest
+          cache-from: type=gha
           outputs: type=image,compression=zstd,force-compression=true
           build-args: |
             DEVICE=${{ inputs.device }}
@@ -123,19 +126,23 @@ jobs:
         shell: bash
         run: |
           SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}"
-          TEMP_IMG="$SOURCE_IMAGE:${{ inputs.image_date }}-build-${{ github.run_id }}"
+
+          # Add date tag
+          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${{ inputs.image_date }}" --quiet
 
           # 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 "$TEMP_IMG" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
+          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
 
           # Add post-training dependencies tags
-          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
+          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

.github/workflows/build_and_test_maxtext.yml

@@ -113,47 +113,72 @@ jobs:
     with:
       device_type: tpu
       device_name: v6e-4
-      base_image: maxtext-unit-test-tpu:${{ matrix.image_type }}
+      image_type: ${{ 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 }}
 
-  tpu-tests:
-    needs: [build_and_upload_maxtext_package]
+  maxtext_cpu_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:
-        flavor: [tpu-unit, tpu-integration]
-    uses: ./.github/workflows/run_tests_coordinator.yml
+        fail-fast: false # don't cancel all jobs on failure
+        matrix:
+          image_type: ["py312"]
+          worker_group: [1, 2]
     with:
-      flavor: ${{ matrix.flavor }}
-      base_image: maxtext-unit-test-tpu:py312
+      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"
       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 }}
 
-  gpu-tests:
-    needs: [build_and_upload_maxtext_package]
+  maxtext_tpu_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:
-        flavor: [gpu-unit, gpu-integration]
-    uses: ./.github/workflows/run_tests_coordinator.yml
+        fail-fast: false
+        matrix:
+          image_type: ["py312"]
     with:
-      flavor: ${{ matrix.flavor }}
-      base_image: maxtext-unit-test-cuda12:py312
+      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"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
-  cpu-tests:
-    needs: [build_and_upload_maxtext_package]
+  maxtext_tpu_integration_tests:
+    needs: build_and_upload_maxtext_package
     if: needs.doc_only_check.outputs.run_tests == 'true'
-    uses: ./.github/workflows/run_tests_coordinator.yml
+    uses: ./.github/workflows/run_tests_against_package.yml
+    strategy:
+        fail-fast: false
+        matrix:
+          image_type: ["py312"]
     with:
-      flavor: cpu-unit
-      base_image: maxtext-unit-test-tpu:py312
+      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"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
@@ -163,12 +188,14 @@ jobs:
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
+        matrix:
+          image_type: ["py312"]
     with:
       device_type: tpu
       device_name: v6e-4
-      base_image: maxtext-unit-test-tpu:py312
+      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 and not post_training'
+      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"
@@ -181,38 +208,85 @@ jobs:
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
+        matrix:
+          image_type: ["py312"]
     with:
       device_type: tpu
       device_name: v6e-4
-      base_image: maxtext-unit-test-tpu:py312
+      image_type: ${{ matrix.image_type }}
       cloud_runner: linux-x86-ct6e-180-4tpu
-      pytest_marker: 'not cpu_only and not gpu_only and integration_test and not post_training'
+      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"
       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: [tpu-tests, gpu-tests, cpu-tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests]
+    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]
     if: always()
     runs-on: ubuntu-latest
     steps:
       - name: Check test results
         run: |
+          # If doc-only, all tests should be skipped
           if [ "${{ needs.doc_only_check.outputs.run_tests }}" == "false" ]; then
             echo "Documentation-only changes detected, tests were skipped"
             exit 0
           fi
 
           # Otherwise, check that build and all tests passed or were skipped
           echo "Build result: ${{ needs.build_and_upload_maxtext_package.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 }}"
+          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 }}"
 
           # Fail only if any job failed or was cancelled (skipped is OK)
           if [ "${{ contains(needs.*.result, 'failure') }}" == "true" ] || [ "${{ contains(needs.*.result, 'cancelled') }}" == "true" ]; then
@@ -249,14 +323,14 @@ jobs:
 
   notify_failure:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build
-    needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests]
+    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]
     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_name == 'schedule' }}
-        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.pull_request == null && github.event_name != 'workflow_dispatch' }}
+      uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
+      with:
+        github-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/run_jupyter_notebooks.yml

@@ -25,20 +25,15 @@ on:
       device_name:
         required: true
         type: string
-      base_image:
-        required: true
+      image_type:
+        required: false
         type: string
       cloud_runner:
         required: false
         type: string
       maxtext_sha:
-        required: false
+        required: true
         type: string
-      # Flag to skip source checkout and wheel installation
-      maxtext_installed:
-        required: false
-        type: boolean
-        default: false
     secrets:
       HF_TOKEN:
         required: true
@@ -49,20 +44,17 @@ 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/${{ inputs.base_image }}
+      image: gcr.io/tpu-prod-env-multipod/maxtext-unit-test-${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}:${{ inputs.image_type != '' && inputs.image_type }}
     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: |
           python3 -m uv venv --seed
@@ -73,6 +65,10 @@ jobs:
           uv pip install ${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest
           uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
 
+          # Install dependencies for running notebooks
+          uv pip install papermill ipykernel ipywidgets
+          .venv/bin/python3 -m ipykernel install --user --name maxtext_venv
+
           # Install Tunix for post-training notebooks
           git clone https://github.com/google/tunix
           uv pip install ./tunix
@@ -94,25 +90,12 @@ jobs:
           PYTHONPATH: "${{ github.workspace }}/src"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          if [ "${{ inputs.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"
-          fi
+          source .venv/bin/activate
 
           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")
             output_name="${filename%.ipynb}_output.ipynb"
@@ -121,7 +104,7 @@ jobs:
             echo "Running $filename ..."
             echo "------------------------------------------------------"
 
-            $PAPERMILL_EXE "$notebook" "$output_name" -k maxtext_venv
+            papermill "$notebook" "$output_name" -k maxtext_venv
           done
       - name: Record Commit IDs
         shell: bash