Trigger unit tests for docker images upload

Author: xibinliu

.github/workflows/UploadDockerImages.yml

@@ -56,7 +56,7 @@ jobs:
       device_type: cpu
       device_name: X64
       cloud_runner: linux-x86-n2-16
-      image_type: ${{ matrix.image_type }}
+      image_name_and_tag: maxtext-unit-test-tpu:${{ matrix.image_type }}
       pytest_marker: 'cpu_only'
       xla_python_client_mem_fraction: 0.75
       tf_force_gpu_allow_growth: false
@@ -75,7 +75,7 @@ jobs:
     with:
       device_type: tpu
       device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
+      image_name_and_tag: maxtext-unit-test-tpu:${{ 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
@@ -93,7 +93,7 @@ jobs:
     with:
       device_type: tpu
       device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
+      image_name_and_tag: maxtext-unit-test-tpu:${{ 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
@@ -111,7 +111,7 @@ jobs:
     with:
       device_type: tpu
       device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
+      image_name_and_tag: maxtext-unit-test-tpu:${{ 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
@@ -129,7 +129,7 @@ jobs:
     with:
       device_type: tpu
       device_name: v6e-4
-      image_type: ${{ matrix.image_type }}
+      image_name_and_tag: maxtext-unit-test-tpu:${{ 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
@@ -148,7 +148,7 @@ jobs:
     with:
       device_type: ${{ matrix.cuda }}
       device_name: a100-40gb-4
-      image_type: ${{ matrix.image_type }}
+      image_name_and_tag: maxtext-unit-test-${{ metrix.cuda }}:${{ matrix.image_type }}
       cloud_runner: linux-x86-a2-48-a100-4gpu
       pytest_marker: 'not cpu_only and not tpu_only and not integration_test'
       pytest_addopts: '--ignore=tests/sft_hooks_test.py'
@@ -168,7 +168,7 @@ jobs:
     with:
       device_type: ${{ matrix.cuda }}
       device_name: a100-40gb-4
-      image_type: ${{ matrix.image_type }}
+      image_name_and_tag: maxtext-unit-test-${{ metrix.cuda }}:${{ matrix.image_type }}
       cloud_runner: linux-x86-a2-48-a100-4gpu
       pytest_marker: 'not cpu_only and not tpu_only and integration_test'
       pytest_addopts: '--ignore=tests/sft_hooks_test.py'

.github/workflows/build_and_test_maxtext.yml

@@ -25,8 +25,8 @@ on:
       device_name:
         required: true
         type: string
-      image_type:
-        required: false
+      image_name_and_tag:
+        required: true
         type: string
       pytest_marker:
         required: true
@@ -58,7 +58,7 @@ 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-tpu:${{ inputs.image_type != '' && inputs.image_type }}
+      image: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name_and_tag }}
       env:
         XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
         TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}

.github/workflows/run_pathways_tests.yml

@@ -25,8 +25,8 @@ on:
       device_name:
         required: true
         type: string
-      image_type:
-        required: false
+      image_name_and_tag:
+        required: true
         type: string
       pytest_marker:
         required: true
@@ -58,14 +58,19 @@ on:
         required: false
         type: number
         default: 1
+      # Flag to skip source checkout and wheel installation
+      maxtext_installed:
+        required: false
+        type: boolean
+        default: false
 
 permissions:
   contents: read
 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.image_name_and_tag }}
       env:
         XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
         TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
@@ -75,11 +80,14 @@ jobs:
       options: ${{ inputs.container_resource_option }}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        if: ${{ !inputs.maxtext_installed }}
       - name: Download the maxtext wheel
+        if: ${{ !inputs.maxtext_installed }}
         uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
         with:
           name: maxtext-wheel
       - name: Install the maxtext wheel
+        if: ${{ !inputs.maxtext_installed }}
         shell: bash
         run: |
           python3 -m uv venv --seed
@@ -91,10 +99,17 @@ jobs:
           python3 -m pip freeze
           uv pip install pytest-cov
       - name: Copy test assets files
+        if: ${{ !inputs.maxtext_installed }}
         run : gcloud storage cp gs://maxtext-test-assets/* src/MaxText/test_assets
       - name: Run Tests
         shell: bash
         run: |
+          if [ "${{ inputs.maxtext_installed }}" == "true" ]; then
+            PYTHON_EXE="python3"
+          else
+            PYTHON_EXE=".venv/bin/python3"
+          fi
+
           if [ "${{ inputs.is_scheduled_run }}" == "true" ]; then
             FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}"
           else
@@ -110,13 +125,13 @@ jobs:
             export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
           fi
           if [ "${{ inputs.total_workers }}" -gt 1 ]; then
-            .venv/bin/python3 -m pip install --quiet pytest-split pytest-xdist
+            $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist
             SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }} -n auto"
           else
             SPLIT_ARGS=""
           fi
           # TODO: Fix the skipped tests and remove the deselect flags
-          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} \
+          $PYTHON_EXE -m pytest ${{ inputs.pytest_addopts }} \
             -v \
             -m "${FINAL_PYTEST_MARKER}" \
             --durations=0 \