CI: Run CI tests on stable deps in Docker image pipelines and explicitly name test jobs

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

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

.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

.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,15 +227,15 @@ 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
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       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() &&

.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

.github/workflows/docs_link_check.yml

@@ -33,6 +33,7 @@ permissions:
 
 jobs:
   linkcheck:
+    name: Check Documentation Links
     runs-on: ubuntu-latest
 
     steps:

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

.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

.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,12 +67,40 @@ 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:
       flavor: ${{ matrix.flavor }}
       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

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

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

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

.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' }}"