Merge pull request #2496 from AI-Hypercomputer:tests_improved

PiperOrigin-RevId: 823644326

Author: Google-ML-Automation

.github/workflows/RunTests.yml

@@ -65,17 +65,22 @@ jobs:
 
   cpu_unit_tests:
     needs: tpu_image
+    strategy:
+      fail-fast: false
+      matrix:
+        worker_group: [1, 2, 3, 4]
     uses: ./.github/workflows/run_tests_internal.yml
     with:
       device_type: cpu
       device_name: X64
-      cloud_runner: linux-x86-n2-16
       image_type: tpu
       pytest_marker: 'cpu_only'
       xla_python_client_mem_fraction: 0.75
       tf_force_gpu_allow_growth: false
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
+      worker_group: ${{ matrix.worker_group }}
+      total_workers: 4
 
   tpu_unit_tests:
     needs: tpu_image

.github/workflows/build_and_test_maxtext.yml

@@ -51,6 +51,7 @@ jobs:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
           image_type: ["py312"]
+          worker_group: [1, 2, 3, 4]
     with:
       device_type: cpu
       device_name: X64
@@ -61,6 +62,8 @@ jobs:
       tf_force_gpu_allow_growth: false
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
+      worker_group: ${{ matrix.worker_group }}
+      total_workers: 4
 
   maxtext_tpu_unit_tests:
     needs: build_and_upload_maxtext_package

.github/workflows/run_tests_against_package.yml

@@ -31,10 +31,6 @@ on:
       pytest_marker:
         required: true
         type: string
-      pytest_addopts:
-        required: false
-        type: string
-        default: ''
       is_scheduled_run:
         required: true
         type: string
@@ -50,12 +46,20 @@ on:
       cloud_runner:
         required: false
         type: string
+      worker_group:
+        required: false
+        type: number
+        default: 1
+      total_workers:
+        required: false
+        type: number
+        default: 1
 
 permissions:
   contents: read
 jobs:
   run:
-    runs-on: ${{ inputs.cloud_runner }}
+    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 }}
       env:
@@ -102,4 +106,5 @@ jobs:
             export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
           fi
           # TODO: Fix the skipped tests and remove the deselect flags
-          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" --durations=0 --deselect "tests/aot_hlo_identical_test.py::AotHloIdenticalTest::test_default_hlo_match" --deselect "tests/tokenizer_test.py::TokenizerTest::test_detokenize"
+          [ "${{ inputs.total_workers }}" -gt 1 ] && .venv/bin/python3 -m pip install --quiet pytest-split && SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}" || SPLIT_ARGS=""
+          .venv/bin/python3 -m pytest -v -m "${FINAL_PYTEST_MARKER}" --durations=0 --deselect "tests/aot_hlo_identical_test.py::AotHloIdenticalTest::test_default_hlo_match" --deselect "tests/tokenizer_test.py::TokenizerTest::test_detokenize" $SPLIT_ARGS

.github/workflows/run_tests_internal.yml

@@ -50,6 +50,14 @@ on:
       cloud_runner:
         required: false
         type: string
+      worker_group:
+        required: false
+        type: number
+        default: 1
+      total_workers:
+        required: false
+        type: number
+        default: 1
 
 jobs:
   run:
@@ -70,5 +78,7 @@ jobs:
           else
             FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and not scheduled_only"
           fi
-          python3 -m pip install -e . --no-dependencies &&
-          LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" --durations=0
+          python3 -m pip install -e . --no-dependencies
+          [ "${{ inputs.total_workers }}" -gt 1 ] && python3 -m pip install --quiet pytest-split && SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}" || SPLIT_ARGS=""
+          export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
+          python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" --durations=0 $SPLIT_ARGS