From fb46bba91685dff0edd82e6fee6df1759ce6a0ab Mon Sep 17 00:00:00 2001
From: Surbhi Jain <sjsurbhi@google.com>
Date: Tue, 11 Aug 2026 22:20:24 +0000
Subject: [PATCH] CI: Run CI tests on stable deps in Docker image pipelines and
 explicitly name test jobs

---
 .../workflows/build_and_push_docker_image.yml | 15 -------
 .github/workflows/build_package.yml           |  1 +
 .github/workflows/check_docs_build.yml        |  1 +
 .github/workflows/ci_pipeline.yml             | 20 ++++++---
 .github/workflows/code_quality.yml            |  2 +-
 .github/workflows/docs_link_check.yml         |  1 +
 .../workflows/gpu_nightly_images_pipeline.yml | 21 ++++++++-
 .github/workflows/require_checklist.yml       |  1 +
 .github/workflows/run_ci_tests.yml            | 34 +++++++++++++-
 .github/workflows/run_jupyter_notebooks.yml   |  1 +
 .github/workflows/run_pathways_tests.yml      |  1 +
 .../workflows/run_tests_against_package.yml   |  1 +
 .github/workflows/run_tests_coordinator.yml   | 18 ++++----
 .../workflows/tpu_nightly_images_pipeline.yml | 44 ++-----------------
 14 files changed, 84 insertions(+), 77 deletions(-)

diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml
index 18bf6081fc..f5dc62d4a5 100644
--- a/.github/workflows/build_and_push_docker_image.yml
+++ b/.github/workflows/build_and_push_docker_image.yml
@@ -42,10 +42,6 @@ on:
         required: false
         type: boolean
         default: false
-      run_tests:
-        required: false
-        type: boolean
-        default: true
     outputs:
       build_result:
         description: 'The result of build_and_push job.'
@@ -195,14 +191,3 @@ jobs:
         env:
           INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
           PROJECT_NAME: ${{ vars.PROJECT_NAME }}
-
-  run_ci_tests:
-    name: Run Unit and Integration Tests
-    needs: [pre_build_check, build_and_push]
-    if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true && inputs.run_tests == true
-    uses: ./.github/workflows/run_ci_tests.yml
-    with:
-      image_name: ${{ inputs.image_name }}
-      image_tag: ${{ github.run_id }}
-      device: ${{ inputs.device }}
-      workflow: ${{ inputs.workflow }}
diff --git a/.github/workflows/build_package.yml b/.github/workflows/build_package.yml
index 0e61db2b6a..e53aa21698 100644
--- a/.github/workflows/build_package.yml
+++ b/.github/workflows/build_package.yml
@@ -41,6 +41,7 @@ permissions:
   contents: read
 jobs:
   build_and_upload:
+    name: Build Wheel
     runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
     container: python:3.12-slim-bookworm
     outputs:
diff --git a/.github/workflows/check_docs_build.yml b/.github/workflows/check_docs_build.yml
index 5979174805..677cd28edf 100644
--- a/.github/workflows/check_docs_build.yml
+++ b/.github/workflows/check_docs_build.yml
@@ -30,6 +30,7 @@ permissions:
 
 jobs:
   build-docs:
+    name: Build Sphinx Docs
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml
index bcdfd51ed8..34d7c0cd25 100644
--- a/.github/workflows/ci_pipeline.yml
+++ b/.github/workflows/ci_pipeline.yml
@@ -77,16 +77,19 @@ jobs:
           bash .github/scripts/analyze_code_changes.sh
 
   code_quality_check:
+    name: Code Quality Check
     uses: ./.github/workflows/code_quality.yml
     with:
       maxtext_sha: ${{ inputs.maxtext_sha || github.sha }}
 
   docs_build_check:
+    name: Documentation Build Check
     uses: ./.github/workflows/check_docs_build.yml
     with:
       maxtext_sha: ${{ inputs.maxtext_sha || github.sha }}
 
   build_and_upload_maxtext_package:
+    name: Build MaxText Package
     needs: [analyze_code_changes, code_quality_check, docs_build_check]
     # Run if either tests or notebooks need to run; on PRs, gate on code quality + docs passing
     if: |
@@ -101,6 +104,7 @@ jobs:
       maxtext_sha: ${{ inputs.maxtext_sha || github.sha }}
 
   maxtext_jupyter_notebooks:
+    name: Jupyter Notebook Tests
     needs: [analyze_code_changes, build_and_upload_maxtext_package]
     if: |
       always() &&
@@ -140,7 +144,7 @@ jobs:
           echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
 
   tpu-pretrain-tests:
-    name: ${{ matrix.flavor || 'TPU Pretrain' }} tests
+    name: TPU Pretrain Tests
     needs: [gate_test_run, analyze_code_changes]
     if: |
       always() &&
@@ -158,7 +162,7 @@ jobs:
       maxtext_sha: ${{ needs.gate_test_run.outputs.maxtext_sha }}
 
   tpu-posttrain-tests:
-    name: ${{ matrix.flavor || 'TPU Posttrain' }} tests
+    name: TPU Posttrain Tests
     needs: [gate_test_run, analyze_code_changes]
     if: |
       always() &&
@@ -177,7 +181,7 @@ jobs:
 
 
   tpu7x-tests:
-    name: TPU7X tests
+    name: TPU7X Tests
     needs: [gate_test_run]
     if: |
       always() &&
@@ -195,7 +199,7 @@ jobs:
       maxtext_sha: ${{ needs.gate_test_run.outputs.maxtext_sha }}
 
   gpu-tests:
-    name: ${{ matrix.flavor || 'GPU' }} tests
+    name: GPU Tests
     needs: [gate_test_run, analyze_code_changes]
     if: |
       always() &&
@@ -213,7 +217,7 @@ jobs:
       maxtext_sha: ${{ needs.gate_test_run.outputs.maxtext_sha }}
 
   cpu-pretrain-tests:
-    name: ${{ matrix.flavor || 'CPU Pretrain' }} tests
+    name: CPU Pretrain Tests
     needs: [gate_test_run, analyze_code_changes]
     if: |
       always() &&
@@ -223,7 +227,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        flavor: [cpu-unit]
+        flavor: [cpu-unit, cpu-integration]
     with:
       flavor: ${{ matrix.flavor }}
       base_image: maxtext-unit-test-tpu:py312
@@ -231,7 +235,7 @@ jobs:
       maxtext_sha: ${{ needs.gate_test_run.outputs.maxtext_sha }}
 
   cpu-posttrain-tests:
-    name: ${{ matrix.flavor || 'CPU Posttrain' }} tests
+    name: CPU Posttrain Tests
     needs: [gate_test_run, analyze_code_changes]
     if: |
       always() &&
@@ -249,6 +253,7 @@ jobs:
       maxtext_sha: ${{ needs.gate_test_run.outputs.maxtext_sha }}
 
   maxtext_tpu_pathways_unit_tests:
+    name: TPU Pathways Unit Tests
     needs: [gate_test_run, analyze_code_changes]
     if: |
       always() &&
@@ -275,6 +280,7 @@ jobs:
       worker_group: ${{ matrix.group }}
 
   maxtext_tpu_pathways_integration_tests:
+    name: TPU Pathways Integration Tests
     needs: [gate_test_run, analyze_code_changes]
     if: |
       always() &&
diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml
index 0b5db7ea44..85ab121430 100644
--- a/.github/workflows/code_quality.yml
+++ b/.github/workflows/code_quality.yml
@@ -26,7 +26,7 @@ on:
 
 jobs:
   qa:
-    name: "Static code-quality checkers"
+    name: "Pre-commit Linters"
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0
diff --git a/.github/workflows/docs_link_check.yml b/.github/workflows/docs_link_check.yml
index da303603e1..0c033b163f 100644
--- a/.github/workflows/docs_link_check.yml
+++ b/.github/workflows/docs_link_check.yml
@@ -33,6 +33,7 @@ permissions:
 
 jobs:
   linkcheck:
+    name: Check Documentation Links
     runs-on: ubuntu-latest
 
     steps:
diff --git a/.github/workflows/gpu_nightly_images_pipeline.yml b/.github/workflows/gpu_nightly_images_pipeline.yml
index 41d8a9f2aa..509c95e5c0 100644
--- a/.github/workflows/gpu_nightly_images_pipeline.yml
+++ b/.github/workflows/gpu_nightly_images_pipeline.yml
@@ -68,9 +68,26 @@ jobs:
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
+  run_ci_tests:
+    name: Run ${{ matrix.name }} CI Tests
+    needs: build_and_push_docker_images
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: "Pre-Training"
+            image_name: maxtext_gpu_jax_stable
+            workflow: pre-training
+    uses: ./.github/workflows/run_ci_tests.yml
+    with:
+      image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
+      image_tag: ${{ github.run_id }}
+      device: gpu
+      workflow: ${{ matrix.workflow }}
+
   notify_failure:
     name: Notify failed build
-    needs: [build_and_push_docker_images]
+    needs: [build_and_push_docker_images, run_ci_tests]
     if: ${{ failure() && inputs.image_suffix == '' }}
     runs-on: ubuntu-latest
     permissions:
@@ -80,5 +97,5 @@ jobs:
         uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          title-template: "MaxText Docker Image Build Failure"
+          title-template: "MaxText GPU Docker Image Build Failure"
           label-name: "docker-image-build-failure"
diff --git a/.github/workflows/require_checklist.yml b/.github/workflows/require_checklist.yml
index aa41fbd61b..3091b01f9e 100644
--- a/.github/workflows/require_checklist.yml
+++ b/.github/workflows/require_checklist.yml
@@ -20,6 +20,7 @@ on:
     types: [opened, edited, synchronize]
 jobs:
   check_pr_body:
+    name: Require PR Checklist
     runs-on: ubuntu-latest
     permissions:
       pull-requests: read
diff --git a/.github/workflows/run_ci_tests.yml b/.github/workflows/run_ci_tests.yml
index bdada36650..787992246d 100644
--- a/.github/workflows/run_ci_tests.yml
+++ b/.github/workflows/run_ci_tests.yml
@@ -36,6 +36,7 @@ permissions:
 
 jobs:
   check_image:
+    name: Check Docker Image Existence
     runs-on: linux-x86-n2-16-buildkit
     container: google/cloud-sdk:524.0.0
     outputs:
@@ -56,6 +57,7 @@ jobs:
           fi
 
   test:
+    name: CI Tests
     needs: check_image
     if: needs.check_image.outputs.image_exists == 'true'
     strategy:
@@ -65,7 +67,7 @@ jobs:
             ${{ fromJSON('{
               "gpu-pre-training": ["gpu-unit", "gpu-integration"],
               "tpu-post-training": ["tpu-post-training-unit", "tpu-post-training-integration", "cpu-post-training-unit"],
-              "tpu-pre-training": ["tpu-unit", "tpu-integration", "cpu-unit"]
+              "tpu-pre-training": ["tpu-unit", "tpu-integration", "cpu-unit", "cpu-integration"]
             }')[format('{0}-{1}', inputs.device, inputs.workflow)] }}
     uses: ./.github/workflows/run_tests_coordinator.yml
     with:
@@ -73,4 +75,32 @@ jobs:
       base_image: ${{ inputs.image_name }}:${{ inputs.image_tag }}
       is_scheduled_run: true
       maxtext_installed: true
-      additional_pytest_args: ${{ contains(inputs.image_name, 'maxtext_jax_nightly') && '--ignore=tests/integration/hlo_diff_test.py' || '' }}
+
+  tag_docker_image:
+    name: Tag Docker Image
+    needs: test
+    if: needs.test.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: Add tags
+        shell: bash
+        env:
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          PROJECT_NAME: ${{ vars.PROJECT_NAME }}
+          IMAGE_NAME_EXPR: ${{ inputs.image_name }}
+          IMAGE_TAG_EXPR: ${{ inputs.image_tag }}
+        run: |
+          image_name="$IMAGE_NAME_EXPR"
+          image_tag="$IMAGE_TAG_EXPR"
+          SOURCE_IMAGE="gcr.io/${PROJECT_NAME}/${image_name}"
+
+          # 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
diff --git a/.github/workflows/run_jupyter_notebooks.yml b/.github/workflows/run_jupyter_notebooks.yml
index d57ae7a71c..4837d72d2d 100644
--- a/.github/workflows/run_jupyter_notebooks.yml
+++ b/.github/workflows/run_jupyter_notebooks.yml
@@ -47,6 +47,7 @@ permissions:
   contents: read
 jobs:
   run:
+    name: Execute Notebooks
     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 }} # zizmor: ignore[unpinned-images]
diff --git a/.github/workflows/run_pathways_tests.yml b/.github/workflows/run_pathways_tests.yml
index c35d7ee1bd..97de8772eb 100644
--- a/.github/workflows/run_pathways_tests.yml
+++ b/.github/workflows/run_pathways_tests.yml
@@ -68,6 +68,7 @@ permissions:
 
 jobs:
   run:
+    name: ${{ format('tpu-pathways-{0}', !contains(inputs.pytest_marker, 'not integration_test') && 'integration' || 'unit') }}
     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 }} # zizmor: ignore[unpinned-images]
diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml
index 9fb9bdb356..8990535ea3 100644
--- a/.github/workflows/run_tests_against_package.yml
+++ b/.github/workflows/run_tests_against_package.yml
@@ -86,6 +86,7 @@ permissions:
   contents: read
 jobs:
   run:
+    name: ${{ inputs.flavor }}
     runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
     container:
       image: gcr.io/${{ vars.PROJECT_NAME || 'tpu-prod-env-multipod' }}/${{ inputs.base_image }} # zizmor: ignore[unpinned-images]
diff --git a/.github/workflows/run_tests_coordinator.yml b/.github/workflows/run_tests_coordinator.yml
index 17b0749a28..96be70f92f 100644
--- a/.github/workflows/run_tests_coordinator.yml
+++ b/.github/workflows/run_tests_coordinator.yml
@@ -28,7 +28,7 @@ on:
             tpu7x-unit, tpu7x-integration,
             tpu7x-post-training-unit, tpu7x-post-training-integration,
             gpu-unit, gpu-integration,
-            cpu-unit,
+            cpu-unit, cpu-integration,
             cpu-post-training-unit
           )
         required: true
@@ -53,11 +53,6 @@ on:
         required: false
         type: boolean
         default: false
-      additional_pytest_args:
-        description: 'Additional pytest arguments to pass during test execution'
-        required: false
-        type: string
-        default: ''
 
 permissions:
   contents: read
@@ -99,8 +94,8 @@ jobs:
           fi
 
   execute-test-package:
+    name: Execute Tests
     needs: setup-parameters
-    name: ${{ inputs.flavor }}
     strategy:
       fail-fast: false
       matrix:
@@ -123,6 +118,7 @@ jobs:
           "gpu-unit": "cuda12",
           "gpu-integration": "cuda12",
           "cpu-unit": "cpu",
+          "cpu-integration": "cpu",
           "cpu-post-training-unit": "cpu"
         }')[inputs.flavor] }}
 
@@ -139,6 +135,7 @@ jobs:
           "gpu-unit": "a100-40gb-4",
           "gpu-integration": "a100-40gb-4",
           "cpu-unit": "X64",
+          "cpu-integration": "X64",
           "cpu-post-training-unit": "X64"
         }')[inputs.flavor] }}
 
@@ -155,6 +152,7 @@ jobs:
           "gpu-unit": "linux-x86-a2-48-a100-4gpu",
           "gpu-integration": "linux-x86-a2-48-a100-4gpu",
           "cpu-unit": "linux-x86-n2-32",
+          "cpu-integration": "linux-x86-n2-32",
           "cpu-post-training-unit": "linux-x86-n2-32"
         }')[inputs.flavor] }}
       # Pytest Marker Mapping
@@ -170,7 +168,8 @@ jobs:
             "tpu7x-post-training-integration": "not cpu_only and not gpu_only and integration_test and post_training",
             "gpu-unit": "not cpu_only and not tpu_only and not integration_test and not post_training",
             "gpu-integration": "not cpu_only and not tpu_only and integration_test and not post_training",
-            "cpu-unit": "cpu_only and not post_training",
+            "cpu-unit": "cpu_only and not post_training and not integration_test",
+            "cpu-integration": "cpu_only and not post_training and integration_test",
             "cpu-post-training-unit": "cpu_only and post_training"
           }')[inputs.flavor] }}
 
@@ -187,6 +186,7 @@ jobs:
             "gpu-unit": "",
             "gpu-integration": "",
             "cpu-unit": "",
+            "cpu-integration": "",
             "cpu-post-training-unit": "tests/post_training/unit tests/unit"
           }')[inputs.flavor] }}
 
@@ -203,9 +203,9 @@ jobs:
             "gpu-unit": "--ignore=tests/post_training",
             "gpu-integration": "--ignore=tests/post_training",
             "cpu-unit": "--ignore=tests/post_training",
+            "cpu-integration": "--ignore=tests/post_training",
             "cpu-post-training-unit": ""
           }')[inputs.flavor] }}
-        ${{ inputs.additional_pytest_args }}
 
       # Resource Scaling
       xla_python_client_mem_fraction: "${{ contains(inputs.flavor, 'gpu') && '0.65' || '0.75' }}"
diff --git a/.github/workflows/tpu_nightly_images_pipeline.yml b/.github/workflows/tpu_nightly_images_pipeline.yml
index 21735d25da..5a82aaa927 100644
--- a/.github/workflows/tpu_nightly_images_pipeline.yml
+++ b/.github/workflows/tpu_nightly_images_pipeline.yml
@@ -69,7 +69,6 @@ jobs:
       dockerfile: maxtext_tpu_dependencies.Dockerfile
       maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
       include_test_assets: true
-      run_tests: false
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
@@ -80,13 +79,10 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - name: "Pre-Training Stable"
+          - name: "Pre-Training"
             image_name: maxtext_jax_stable
             workflow: pre-training
-          - name: "Pre-Training Nightly"
-            image_name: maxtext_jax_nightly
-            workflow: pre-training
-          - name: "Post-Training Nightly"
+          - name: "Post-Training"
             image_name: maxtext_post_training_nightly
             workflow: post-training
     uses: ./.github/workflows/run_ci_tests.yml
@@ -108,40 +104,6 @@ jobs:
     secrets:
       AIRFLOW_CALLBACK_TOKEN: ${{ secrets.AIRFLOW_CALLBACK_TOKEN }}
 
-  # TODO: This is a workaround, to be removed when promote_docker_image.yml workflow is stable
-  tag_docker_image:
-    name: Promote ${{ matrix.image_name }} Docker Image
-    needs: run_ci_tests
-    if: needs.run_ci_tests.result == 'success'
-    runs-on: linux-x86-n2-16-buildkit
-    container: google/cloud-sdk:524.0.0
-    strategy:
-      fail-fast: false
-      matrix:
-        image_name:
-          - maxtext_jax_stable
-          - maxtext_jax_nightly
-          - maxtext_post_training_nightly
-    steps:
-      - name: Configure Docker
-        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
-      - name: Add tags to Docker image
-        shell: bash
-        env:
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          PROJECT_NAME: ${{ vars.PROJECT_NAME }}
-          IMAGE_NAME_EXPR: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
-        run: |
-          image_name="$IMAGE_NAME_EXPR"
-          SOURCE_IMAGE="gcr.io/${PROJECT_NAME}/${image_name}"
-
-          # 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
-
-          # Add "latest" tag
-          gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" "${SOURCE_IMAGE}:latest" --quiet
-
   notify_failure:
     name: Notify failed build
     needs: [build_and_push_docker_images, run_ci_tests, run_e2e_tests]
@@ -154,5 +116,5 @@ jobs:
         uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          title-template: "MaxText Docker Image Build Failure"
+          title-template: "MaxText TPU Docker Image Build Failure"
           label-name: "docker-image-build-failure"