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

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

.github/workflows/run_tests_against_package.yml

@@ -71,6 +71,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
@@ -107,6 +108,7 @@ jobs:
           if [ "${{ inputs.device_type }}" != "cuda12" ]; then
             export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
           fi
+          # Use pytest-split to statically split tests across runners, and pytest-xdist to dynamically split across processes within each runner
+          [ "${{ inputs.total_workers }}" -gt 1 ] && .venv/bin/python3 -m pip install --quiet pytest-split pytest-xdist && SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }} -n auto" || SPLIT_ARGS=""
           # 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=""
           .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