Switch from pytest-split to pytest-xdist for parallel test execution

Previously, CPU tests are distributed across multiple workers using pytest-split,
which assigns the same number of tests to each worker. However, since the runtime
of tests is different, some workers end up finishing fast and stand idle while
others take a long time, so we're not utilizing the workers fully. This change
replaces the use of pytest-split with pytest-xdist which dynamically assigns work to workers.

Author: shralex

.github/workflows/build_and_test_maxtext.yml

@@ -47,23 +47,16 @@ jobs:
   maxtext_cpu_unit_tests:
     needs: build_and_upload_maxtext_package
     uses: ./.github/workflows/run_tests_against_package.yml
-    strategy:
-        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
       cloud_runner: linux-x86-n2-16
-      image_type: ${{ matrix.image_type }}
+      image_type: "py312"
       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
 
   maxtext_tpu_unit_tests:
     needs: build_and_upload_maxtext_package

.github/workflows/run_tests_against_package.yml

@@ -50,14 +50,6 @@ 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
@@ -71,6 +63,7 @@ jobs:
         TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
         TPU_SKIP_MDS_QUERY: ${{ inputs.device_type == 'cpu' && '1' || '' }}
         MAXTEXT_PACKAGE_EXTRA: ${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}
+        ALLOW_MULTIPLE_LIBTPU_LOAD: ${{ inputs.device_type == 'cpu' && 'true' || '' }} # bypass /tmp/libtpu_lockfile check for cpu tests, which don't actually use accelerators (to allow concurrency)
       options: ${{ inputs.container_resource_option }}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -108,5 +101,5 @@ jobs:
             export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
           fi
           # TODO: Fix the skipped tests and remove the deselect flags
-          [ "${{ 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=""
+          [ "${{ inputs.device_type }}" == "cpu" ] && .venv/bin/python3 -m pip install --quiet pytest-xdist && SPLIT_ARGS="-n auto" || SPLIT_ARGS=""
           .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" --durations=0 --deselect "tests/tokenizer_test.py::TokenizerTest::test_detokenize" $SPLIT_ARGS