feat(ci): allow re-running tests without recreating docker images

.github/workflows/build_and_push_docker_image.yml

@@ -46,6 +46,10 @@
         required: false
         type: boolean
         default: true
+      fast_rebuild:
+        required: false
+        type: boolean
+        default: false
     outputs:
       build_result:
         description: 'The result of build_and_push job.'
@@ -70,8 +74,8 @@
         run: |
           EVENT_NAME="${{ github.event_name }}"
           # Default to 'all' if the input is null/empty
           TARGET_DEVICE="${{ github.event.inputs.target_device || 'all' }}"
           INPUT_DEVICE="${{ inputs.device }}"
 
           SHOULD_RUN="false"
           if [[ "$EVENT_NAME" == "release" || "$EVENT_NAME" == "schedule" || "$EVENT_NAME" == "pull_request" ]]; then
@@ -84,10 +88,10 @@
 
           if [[ "$SHOULD_RUN" == "true" ]]; then
             echo "should_run=true" >> $GITHUB_OUTPUT
             echo "Building ${{ inputs.image_name }} for device: ${{ inputs.device }} in ${{ inputs.build_mode }} mode."
           else
             echo "should_run=false" >> $GITHUB_OUTPUT
             echo "Skipping ${{ inputs.image_name }} build for device: ${{ inputs.device }} in ${{ inputs.build_mode }} mode."
           fi
 
   build_and_push:
@@ -99,8 +103,8 @@
     steps:
       - name: Matrix Debugger
         run: |
           echo "device: ${{ inputs.device }}"
           echo "workflow: ${{ inputs.workflow }}"
           echo "build_mode: ${{ inputs.build_mode }}"
           echo "image_name: ${{ inputs.image_name }}"
           echo "dockerfile: ${{ inputs.dockerfile }}"
@@ -158,6 +162,7 @@
           file: .venv/lib/python3.12/site-packages/dependencies/dockerfiles/${{ inputs.dockerfile }}
           tags: gcr.io/${{ vars.PROJECT_NAME }}/${{ inputs.image_name }}:${{ github.run_id }}
           cache-from: type=gha
+          cache-to: type=gha,mode=max
           outputs: type=image,compression=zstd,force-compression=true
           build-args: |
             DEVICE=${{ inputs.device }}

.github/workflows/ci_pipeline.yml

@@ -30,6 +30,16 @@
         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:
+    inputs:
+      mode:
+        description: 'Test execution mode: build-all (default), test-only (skip package build and run tests), or fast-rebuild'
+        required: false
+        default: 'build-all'
+        type: choice
+        options:
+          - build-all
+          - test-only
+          - fast-rebuild
   schedule:
     # Run the job every 4 hours
     - cron:  '0 */4 * * *'
@@ -142,6 +152,7 @@
     # Run if either tests or notebooks need to run; on PRs, gate on code quality + docs passing
     if: |
       always() &&
+      inputs.mode != 'test-only' &&
       (needs.analyze_code_changes.outputs.run_tests == 'true' || needs.analyze_code_changes.outputs.run_notebooks == 'true') &&
       (github.event_name != 'pull_request' || (needs.code_quality_check.result == 'success' && needs.docs_build_check.result == 'success'))
     uses: ./.github/workflows/build_package.yml
@@ -156,7 +167,7 @@
     if: |
       always() &&
       needs.analyze_code_changes.outputs.run_notebooks == 'true' &&
-      needs.build_and_upload_maxtext_package.result == 'success'
+      (inputs.mode == 'test-only' || needs.build_and_upload_maxtext_package.result == 'success')
     uses: ./.github/workflows/run_jupyter_notebooks.yml
     strategy:
         fail-fast: false
@@ -165,7 +176,7 @@
       device_name: v6e-8
       base_image: maxtext-unit-test-tpu:py312
       cloud_runner: linux-x86-ct6e-180-8tpu
-      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha || inputs.maxtext_sha || github.sha }}
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
@@ -175,12 +186,12 @@
     if: |
       always() &&
       needs.analyze_code_changes.outputs.run_tests == 'true' &&
-      needs.build_and_upload_maxtext_package.result == 'success'
+      (inputs.mode == 'test-only' || needs.build_and_upload_maxtext_package.result == 'success')
     runs-on: ubuntu-latest
     outputs:
       total_workers: ${{ steps.set-params.outputs.total_workers }}
       worker_groups: ${{ steps.set-params.outputs.worker_groups }}
-      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha || inputs.maxtext_sha || github.sha }}
     steps:
       - id: set-params
         name: Formalize Test Suite Parameters
@@ -406,7 +417,7 @@
     name: Investigate failed build # investigates failure of scheduled run and comments on tracking issue
     needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check, notify_failure]
     if: ${{ always() && contains(needs.*.result, 'failure') && github.event_name == 'schedule' }}
     uses: ./.github/workflows/gemini_investigate.yml
     permissions:
       contents: 'read'
       id-token: 'write'
@@ -421,7 +432,7 @@
     name: Track Test Performance
     needs: [tpu-tests, gpu-tests, cpu-tests]
     if: ${{ always() && (needs.cpu-tests.result == 'success' || needs.gpu-tests.result == 'success' || needs.tpu-tests.result == 'success') }}
     uses: ./.github/workflows/track_performance.yml
     permissions:
       contents: write
       id-token: write

.github/workflows/gpu_nightly_images_pipeline.yml

@@ -28,6 +28,15 @@ on:
         required: false
         type: string
         default: ""
+      mode:
+        description: 'Docker test execution mode: build-all (rebuild image and test), test-only (run tests against existing image without rebuild), or fast-rebuild (rebuild code only)'
+        required: false
+        default: 'build-all'
+        type: choice
+        options:
+          - build-all
+          - test-only
+          - fast-rebuild
 
 permissions:
   contents: read
@@ -44,6 +53,7 @@ jobs:
   build_and_push_docker_images:
     name: Build ${{ matrix.name }} Docker Image
     needs: build_maxtext_package
+    if: inputs.mode != 'test-only'
     strategy:
       fail-fast: false
       matrix:
@@ -65,6 +75,7 @@ jobs:
       dockerfile: maxtext_gpu_dependencies.Dockerfile
       maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
       include_test_assets: true
+      fast_rebuild: ${{ inputs.mode == 'fast-rebuild' }}
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 

.github/workflows/tpu_nightly_images_pipeline.yml

@@ -28,6 +28,15 @@
         required: false
         type: string
         default: ""
+      mode:
+        description: 'Docker test execution mode: build-all (rebuild image and test), test-only (run tests against existing image without rebuild), or fast-rebuild (rebuild code only)'
+        required: false
+        default: 'build-all'
+        type: choice
+        options:
+          - build-all
+          - test-only
+          - fast-rebuild
 
 permissions:
   contents: read
@@ -44,6 +53,7 @@
   build_and_push_docker_images:
     name: Build ${{ matrix.name }} Docker Image
     needs: build_maxtext_package
+    if: inputs.mode != 'test-only'
     strategy:
       fail-fast: false
       matrix:
@@ -70,12 +80,14 @@
       maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
       include_test_assets: true
       run_tests: false
+      fast_rebuild: ${{ inputs.mode == 'fast-rebuild' }}
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
   run_ci_tests:
     name: Run ${{ matrix.name }} CI Tests
     needs: build_and_push_docker_images
+    if: always() && (inputs.mode == 'test-only' || needs.build_and_push_docker_images.result == 'success')
     strategy:
       fail-fast: false
       matrix:
@@ -92,7 +104,7 @@
     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 }}
+      image_tag: ${{ inputs.mode == 'test-only' && 'latest' || github.run_id }}
       device: tpu
       workflow: ${{ matrix.workflow }}
 
@@ -100,7 +112,7 @@
     name: Run E2E tests
     needs: build_and_push_docker_images
     if: github.event_name == 'schedule'
     uses: ./.github/workflows/run_e2e_tests.yml
     with:
       mode: ${{ inputs.image_suffix != '' && format('{0}_{1}', 'nightly', inputs.image_suffix) || 'nightly' }}
       run_id: ${{ github.run_id }}

docs/tutorials/build_maxtext.md

@@ -146,6 +146,31 @@ If you get a `command not found: upload_maxtext_docker_image` error, it means yo
 **Note on Permissions:** You will need the [**Artifact Registry Writer**](https://docs.cloud.google.com/artifact-registry/docs/access-control#permissions) role to push Docker images to Artifact Registry. If you don't have this permission, contact your project administrator to grant you this role through "Google Cloud Console -> IAM -> Grant access".
 ````
 
+## Re-Running Tests Without Rebuilding Images
+
+When MaxText tests fail or flake, users often want to re-run their tests with the same dependencies (or even the same code) without waiting to rebuild the Docker image from scratch.
+
+In an ideal setup, users can choose between three execution modes using `run_maxtext_docker_test` (or `bash src/dependencies/scripts/run_docker_test.sh`):
+
+1. **Re-run all steps including building image**:
+   ```bash
+   run_maxtext_docker_test --mode=build-all --command="python3 -m pytest -vv tests/unit/"
+   ```
+2. **Re-run just the tests, using the same image (including same code)**:
+   ```bash
+   run_maxtext_docker_test --mode=test-only --command="python3 -m pytest -vv tests/unit/"
+   ```
+3. **Rebuild the image using the same dependencies, only new code (fast rebuild)**:
+   ```bash
+   run_maxtext_docker_test --mode=fast-rebuild --command="python3 -m pytest -vv tests/unit/"
+   ```
+
+### XPK Workload Option
+
+When running benchmarks or tests on XPK using `benchmarks/maxtext_xpk_runner.py`:
+- Set `reuse_image = True` in `WorkloadConfig` to re-run tests using an existing `--docker-image` without rebuilding or deleting any intermediate images.
+- Set `fast_rebuild = True` in `WorkloadConfig` to fast-rebuild only the runner image over cached dependencies (`docker_upload_runner.sh FAST_REBUILD=true`).
+
 ## Troubleshooting
 
 1. If you see the following error while building or uploading your Docker image, try adding the listed file path to `.dockerignore`. Do not include the `./` prefix in the `.dockerignore` file:

pyproject.toml

@@ -48,4 +48,5 @@ install_cuda12_pre_train_extra_deps = "dependencies.scripts.install_pre_train_ex
 install_tpu_post_train_extra_deps = "dependencies.scripts.install_post_train_extra_deps:main"
 build_maxtext_docker_image = "dependencies.scripts.build_maxtext_docker_image:main"
 upload_maxtext_docker_image = "dependencies.scripts.upload_maxtext_docker_image:main"
+run_maxtext_docker_test = "dependencies.scripts.run_docker_test:main"
 run_tutorial = "dependencies.scripts.run_tutorial:main"

src/dependencies/scripts/build_upload_colocated_python_image.sh

@@ -125,8 +125,13 @@ FULL_BASE_IMAGE="${BASE_IMAGE_REPO}:${BASE_IMAGE_TAG}"
 echo "  Found latest compatible base image: ${FULL_BASE_IMAGE}"
 
 # --- Step 3: Build the Docker image using --build-arg ---
+CACHE_FLAG="--no-cache"
+if [[ "${FAST_REBUILD:-false}" == "true" || "${NO_CACHE:-true}" == "false" ]]; then
+  CACHE_FLAG=""
+  echo "Fast rebuild enabled: reusing Docker layer cache."
+fi
 echo "$(date): Running docker build with local tag '${LOCAL_IMAGE_NAME}' using ${TMP_DOCKERFILE}..."
-docker build --no-cache \
+docker build ${CACHE_FLAG} \
   -f ${DOCKERFILE_TEMPLATE} \
   --build-arg BASE_IMAGE="${FULL_BASE_IMAGE}" \
   -t "${LOCAL_IMAGE_NAME}" \
@@ -140,7 +145,11 @@ echo "$(date): Pushing '${IMAGE_LOCATION}'..."
 docker push "${IMAGE_LOCATION}"
 
 # --- Step 5: Cleanup ---
-echo "$(date): Cleaning up local tag '${LOCAL_IMAGE_NAME}'..."
-docker image rm "${LOCAL_IMAGE_NAME}"
+if [[ "${CLEANUP_IMAGE:-true}" == "true" ]]; then
+  echo "$(date): Cleaning up local tag '${LOCAL_IMAGE_NAME}'..."
+  docker image rm "${LOCAL_IMAGE_NAME}"
+else
+  echo "$(date): Preserving local tag '${LOCAL_IMAGE_NAME}' (CLEANUP_IMAGE=false)..."
+fi
 
 echo "$(date): Build and push complete for ${IMAGE_LOCATION}"

src/dependencies/scripts/docker_upload_runner.sh

@@ -116,7 +116,13 @@ if ! docker image inspect "${LOCAL_IMAGE_NAME}" &> /dev/null; then
   exit 1
 fi
 
-docker build --no-cache --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} \
+CACHE_FLAG="--no-cache"
+if [[ "${FAST_REBUILD:-false}" == "true" || "${NO_CACHE:-true}" == "false" ]]; then
+  CACHE_FLAG=""
+  echo "Fast rebuild enabled: reusing Docker layer cache for runner image build."
+fi
+
+docker build ${CACHE_FLAG} --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} \
              --build-arg PACKAGE_DIR=${PACKAGE_DIR} \
              -f "$PACKAGE_DIR"'/dependencies/dockerfiles/maxtext_runner.Dockerfile' \
              -t ${LOCAL_IMAGE_NAME_RUNNER} .

src/dependencies/scripts/run_docker_test.py

@@ -0,0 +1,36 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Wrapper to run run_docker_test.sh from pip install or module execution."""
+
+import os
+import sys
+
+
+def main():
+  current_dir = os.path.dirname(os.path.abspath(__file__))
+  repo_root = os.path.abspath(os.path.join(current_dir, "..", ".."))
+  # Use relative path for Docker
+  os.environ["PACKAGE_DIR"] = os.path.relpath(repo_root, os.getcwd())
+
+  script_path = os.path.join(current_dir, "run_docker_test.sh")
+  if not os.path.exists(script_path):
+    raise FileNotFoundError(f"Script not found at {script_path}")
+
+  cmd = ["bash", script_path] + sys.argv[1:]
+  os.execvp("bash", cmd)
+
+
+if __name__ == "__main__":
+  main()