[CI] Split tpu-unit and pathways unit tests tests into 2 groups

TPU tests need to avoid spawning parallel xdist workers on the same TPU device

Author: xibinliu

.github/workflows/ci_pipeline.yml

@@ -199,12 +199,31 @@ jobs:
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
+  setup-pathways-parameters:
+    name: Setup Pathways Parameters
+    runs-on: ubuntu-latest
+    outputs:
+      total_workers: ${{ steps.set-params.outputs.total_workers }}
+      worker_groups: ${{ steps.set-params.outputs.worker_groups }}
+    steps:
+      - id: set-params
+        name: Set Pathways Worker Parameters
+        run: |
+          # TPU worker constants
+          TPU_UNIT_TOTAL_WORKERS=2
+          TPU_UNIT_WORKER_GROUPS='[1, 2]'
+
+          echo "total_workers=${TPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+
   maxtext_tpu_pathways_unit_tests:
-    needs: build_and_upload_maxtext_package
+    needs: [build_and_upload_maxtext_package, setup-pathways-parameters]
     if: needs.analyze_code_changes.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
+        matrix:
+          group: ${{ fromJSON(needs.setup-pathways-parameters.outputs.worker_groups) }}
     with:
       device_type: tpu
       device_name: v6e-4
@@ -217,6 +236,8 @@ jobs:
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+      total_workers: ${{ needs.setup-pathways-parameters.outputs.total_workers }}
+      worker_group: ${{ matrix.group }}
 
   maxtext_tpu_pathways_integration_tests:
     needs: build_and_upload_maxtext_package

.github/workflows/run_pathways_tests.yml

@@ -53,6 +53,15 @@ on:
       maxtext_sha:
         required: true
         type: string
+      total_workers:
+        required: false
+        type: string
+        default: '1'
+      worker_group:
+        required: false
+        type: string
+        default: '1'
+
 
 permissions:
   contents: read
@@ -95,6 +104,12 @@ jobs:
         run : gcloud storage cp gs://maxtext-test-assets/* tests/assets
       - name: Run Tests
         run: |
+          if [ "${{ inputs.total_workers }}" -gt 1 ]; then
+            .venv/bin/python3 -m pip install --quiet pytest-split
+            SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}"
+          else
+            SPLIT_ARGS=""
+          fi
           if [ "${{ inputs.is_scheduled_run }}" = "true" ]; then
             FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}"
           else
@@ -105,7 +120,7 @@ jobs:
           export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
           export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
           # TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported.
-          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0
+          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0 ${SPLIT_ARGS}
         env:
           PYTHONPATH: "${{ github.workspace }}/src"
     services:

.github/workflows/run_tests_against_package.yml

@@ -176,8 +176,13 @@ jobs:
             fi
           fi
           if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then
-            $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist
-            SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP} -n auto"
+            $PYTHON_EXE -m pip install --quiet pytest-split
+            if [ "${INPUTS_DEVICE_TYPE}" = "tpu" ]; then
+              SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP}"
+            else
+              $PYTHON_EXE -m pip install --quiet pytest-xdist
+              SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP} -n auto"
+            fi
           else
             SPLIT_ARGS=""
           fi

.github/workflows/run_tests_coordinator.yml

@@ -61,12 +61,48 @@ permissions:
   contents: read
 
 jobs:
+  setup-parameters:
+    name: Setup Parameters
+    runs-on: ubuntu-latest
+    outputs:
+      worker_groups: ${{ steps.set-params.outputs.worker_groups }}
+      total_workers: ${{ steps.set-params.outputs.total_workers }}
+    steps:
+      - id: set-params
+        name: Set Worker Parameters
+        run: |
+          # CPU worker constants
+          CPU_UNIT_TOTAL_WORKERS=4
+          CPU_UNIT_WORKER_GROUPS='[1, 2, 3, 4]'
+
+          # TPU worker constants
+          TPU_UNIT_TOTAL_WORKERS=2
+          TPU_UNIT_WORKER_GROUPS='[1, 2]'
+
+          # Default fallback constants
+          DEFAULT_TOTAL_WORKERS=1
+          DEFAULT_WORKER_GROUPS='[1]'
+
+          FLAVOR="${{ inputs.flavor }}"
+
+          if [[ "$FLAVOR" == "cpu-unit" || "$FLAVOR" == "cpu-post-training-unit" ]]; then
+            echo "worker_groups=${CPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+            echo "total_workers=${CPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          elif [[ "$FLAVOR" == "tpu-unit" ]]; then
+            echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+            echo "total_workers=${TPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          else
+            echo "worker_groups=${DEFAULT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+            echo "total_workers=${DEFAULT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          fi
+
   execute-test-package:
+    needs: setup-parameters
     name: ${{ inputs.flavor }}
     strategy:
       fail-fast: false
       matrix:
-        worker_group: ${{ fromJSON(contains(inputs.flavor, 'cpu-unit') && '[1, 2, 3, 4]' || '[1]') }}
+        worker_group: ${{ fromJSON(needs.setup-parameters.outputs.worker_groups) }}
 
     uses: ./.github/workflows/run_tests_against_package.yml
     with:
@@ -158,6 +194,6 @@ jobs:
       is_scheduled_run: ${{ inputs.is_scheduled_run }}
       maxtext_installed: ${{ inputs.maxtext_installed }}
       worker_group: ${{ matrix.worker_group }}
-      total_workers: ${{ contains(inputs.flavor, 'cpu-unit') && 4 || 1 }}
+      total_workers: ${{ fromJSON(needs.setup-parameters.outputs.total_workers) }}
       maxtext_sha: ${{ inputs.maxtext_sha }}
       is_update_hlo: ${{ inputs.is_update_hlo }}