TST Speed-up test suite when using pytest-xdist (#25918)

.github/workflows/wheels.yml

@@ -116,9 +116,7 @@ jobs:
         env:
           CONFTEST_PATH: ${{ github.workspace }}/conftest.py
           CONFTEST_NAME: conftest.py
-          CIBW_ENVIRONMENT: OMP_NUM_THREADS=2
-                            OPENBLAS_NUM_THREADS=2
-                            SKLEARN_SKIP_NETWORK_TESTS=1
+          CIBW_ENVIRONMENT: SKLEARN_SKIP_NETWORK_TESTS=1
                             SKLEARN_BUILD_PARALLEL=3
           CIBW_BUILD: cp${{ matrix.python }}-${{ matrix.platform_id }}
           CIBW_ARCHS: all
@@ -173,8 +171,6 @@ jobs:
       - name: Test source distribution
         run: bash build_tools/github/test_source.sh
         env:
-          OMP_NUM_THREADS: 2
-          OPENBLAS_NUM_THREADS: 2
           SKLEARN_SKIP_NETWORK_TESTS: 1
 
       - name: Store artifacts

build_tools/azure/posix-docker.yml

@@ -16,9 +16,6 @@ jobs:
     VIRTUALENV: 'testvenv'
     TEST_DIR: '$(Agent.WorkFolder)/tmp_folder'
     JUNITXML: 'test-data.xml'
-    OMP_NUM_THREADS: '2'
-    OPENBLAS_NUM_THREADS: '2'
-    CPU_COUNT: '2'
     SKLEARN_SKIP_NETWORK_TESTS: '1'
     PYTEST_XDIST_VERSION: 'latest'
     COVERAGE: 'false'
@@ -71,11 +68,8 @@ jobs:
         -e JUNITXML
         -e VIRTUALENV
         -e PYTEST_XDIST_VERSION
-        -e OMP_NUM_THREADS
-        -e OPENBLAS_NUM_THREADS
         -e SKLEARN_SKIP_NETWORK_TESTS
         -e SELECTED_TESTS
-        -e CPU_COUNT
         -e CCACHE_COMPRESS
         -e BUILD_SOURCEVERSIONMESSAGE
         -e BUILD_REASON

build_tools/azure/posix.yml

@@ -16,9 +16,6 @@ jobs:
     TEST_DIR: '$(Agent.WorkFolder)/tmp_folder'
     VIRTUALENV: 'testvenv'
     JUNITXML: 'test-data.xml'
-    OMP_NUM_THREADS: '2'
-    OPENBLAS_NUM_THREADS: '2'
-    CPU_COUNT: '2'
     SKLEARN_SKIP_NETWORK_TESTS: '1'
     CCACHE_DIR: $(Pipeline.Workspace)/ccache
     CCACHE_COMPRESS: '1'

build_tools/azure/test_script.sh

@@ -34,7 +34,8 @@ mkdir -p $TEST_DIR
 cp setup.cfg $TEST_DIR
 cd $TEST_DIR
 
-python -c "import joblib; print(f'Number of cores: {joblib.cpu_count()}')"
+python -c "import joblib; print(f'Number of cores (physical): \
+{joblib.cpu_count()} ({joblib.cpu_count(only_physical_cores=True)})')"
 python -c "import sklearn; sklearn.show_versions()"
 
 show_installed_libraries
@@ -70,7 +71,8 @@ if [[ -n "$CHECK_WARNINGS" ]]; then
 fi
 
 if [[ "$PYTEST_XDIST_VERSION" != "none" ]]; then
-    TEST_CMD="$TEST_CMD -n$CPU_COUNT"
+    XDIST_WORKERS=$(python -c "import joblib; print(joblib.cpu_count(only_physical_cores=True))")
+    TEST_CMD="$TEST_CMD -n$XDIST_WORKERS"
 fi
 
 if [[ "$SHOW_SHORT_SUMMARY" == "true" ]]; then

build_tools/azure/windows.yml

@@ -19,7 +19,6 @@ jobs:
     PYTEST_XDIST_VERSION: 'latest'
     TEST_DIR: '$(Agent.WorkFolder)/tmp_folder'
     SHOW_SHORT_SUMMARY: 'false'
-    CPU_COUNT: '2'
   strategy:
     matrix:
       ${{ insert }}: ${{ parameters.matrix }}

build_tools/cirrus/arm_tests.yml

@@ -8,8 +8,6 @@ linux_aarch64_test_task:
     memory: 6G
   env:
     CONDA_ENV_NAME: testenv
-    OMP_NUM_THREADS: 2
-    OPENBLAS_NUM_THREADS: 2
     LOCK_FILE: build_tools/cirrus/py39_conda_forge_linux-aarch64_conda.lock
     CONDA_PKGS_DIRS: /root/.conda/pkgs
     HOME: /  # $HOME is not defined in image and is required to install mambaforge

build_tools/cirrus/arm_wheel.yml

@@ -4,11 +4,8 @@ macos_arm64_wheel_task:
   env:
     CONFTEST_PATH: ${CIRRUS_WORKING_DIR}/conftest.py
     CONFTEST_NAME: conftest.py
-    CIBW_ENVIRONMENT: OMP_NUM_THREADS=2
-                      OPENBLAS_NUM_THREADS=2
-                      SKLEARN_SKIP_NETWORK_TESTS=1
+    CIBW_ENVIRONMENT: SKLEARN_SKIP_NETWORK_TESTS=1
                       SKLEARN_BUILD_PARALLEL=5
-                      CPU_COUNT=2
     CIBW_TEST_COMMAND: bash {project}/build_tools/wheels/test_wheels.sh
     CIBW_TEST_REQUIRES: pytest pandas threadpoolctl pytest-xdist
     CIBW_BUILD_VERBOSITY: 1
@@ -54,11 +51,8 @@ linux_arm64_wheel_task:
   env:
     CONFTEST_PATH: ${CIRRUS_WORKING_DIR}/conftest.py
     CONFTEST_NAME: conftest.py
-    CIBW_ENVIRONMENT: OMP_NUM_THREADS=2
-                      OPENBLAS_NUM_THREADS=2
-                      SKLEARN_SKIP_NETWORK_TESTS=1
+    CIBW_ENVIRONMENT: SKLEARN_SKIP_NETWORK_TESTS=1
                       SKLEARN_BUILD_PARALLEL=5
-                      CPU_COUNT=4
     CIBW_TEST_COMMAND: bash {project}/build_tools/wheels/test_wheels.sh
     CIBW_TEST_REQUIRES: pytest pandas threadpoolctl pytest-xdist
     CIBW_BUILD_VERBOSITY: 1

build_tools/github/test_windows_wheels.sh

@@ -11,7 +11,5 @@ docker container run \
 
 docker container run \
     -e SKLEARN_SKIP_NETWORK_TESTS=1 \
-    -e OMP_NUM_THREADS=2 \
-    -e OPENBLAS_NUM_THREADS=2 \
     --rm scikit-learn/minimal-windows \
     powershell -Command "pytest --pyargs sklearn"

build_tools/wheels/test_wheels.sh

@@ -11,12 +11,16 @@ if [[ "$UNAME" != "Linux" ]]; then
     cp $CONFTEST_PATH $CONFTEST_NAME
 fi
 
+python -c "import joblib; print(f'Number of cores (physical): \
+{joblib.cpu_count()} ({joblib.cpu_count(only_physical_cores=True)})')"
+
 # Test that there are no links to system libraries in the
 # threadpoolctl output section of the show_versions output:
 python -c "import sklearn; sklearn.show_versions()"
 
-if [ ! -z "$CPU_COUNT" ]; then
-    pytest --pyargs sklearn -n $CPU_COUNT
+if pip show -qq pytest-xdist; then
+    XDIST_WORKERS=$(python -c "import joblib; print(joblib.cpu_count(only_physical_cores=True))")
+    pytest --pyargs sklearn -n $XDIST_WORKERS
 else
     pytest --pyargs sklearn
 fi

sklearn/conftest.py

@@ -5,13 +5,13 @@
 from contextlib import suppress
 from unittest import SkipTest
 
+import joblib
 import pytest
 import numpy as np
 from threadpoolctl import threadpool_limits
 from _pytest.doctest import DoctestItem
 
 from sklearn.utils import _IS_32BIT
-from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
 from sklearn._min_dependencies import PYTEST_MIN_VERSION
 from sklearn.utils.fixes import sp_version
 from sklearn.utils.fixes import parse_version
@@ -233,27 +233,6 @@ def pyplot():
     pyplot.close("all")
 
 
-def pytest_runtest_setup(item):
-    """Set the number of openmp threads based on the number of workers
-    xdist is using to prevent oversubscription.
-
-    Parameters
-    ----------
-    item : pytest item
-        item to be processed
-    """
-    xdist_worker_count = environ.get("PYTEST_XDIST_WORKER_COUNT")
-    if xdist_worker_count is None:
-        # returns if pytest-xdist is not installed
-        return
-    else:
-        xdist_worker_count = int(xdist_worker_count)
-
-    openmp_threads = _openmp_effective_n_threads()
-    threads_per_worker = max(openmp_threads // xdist_worker_count, 1)
-    threadpool_limits(threads_per_worker, user_api="openmp")
-
-
 def pytest_configure(config):
     # Use matplotlib agg backend during the tests including doctests
     try:
@@ -263,6 +242,14 @@ def pytest_configure(config):
     except ImportError:
         pass
 
+    allowed_parallelism = joblib.cpu_count(only_physical_cores=True)
+    xdist_worker_count = environ.get("PYTEST_XDIST_WORKER_COUNT")
+    if xdist_worker_count is not None:
+        # Set the number of OpenMP and BLAS threads based on the number of workers
+        # xdist is using to prevent oversubscription.
+        allowed_parallelism = max(allowed_parallelism // int(xdist_worker_count), 1)
+    threadpool_limits(allowed_parallelism)
+
     # Register global_random_seed plugin if it is not already registered
     if not config.pluginmanager.hasplugin("sklearn.tests.random_seed"):
         config.pluginmanager.register(random_seed)

sklearn/utils/_openmp_helpers.pyx

@@ -12,7 +12,7 @@ def _openmp_parallelism_enabled():
     return SKLEARN_OPENMP_PARALLELISM_ENABLED
 
 
-cpdef _openmp_effective_n_threads(n_threads=None):
+cpdef _openmp_effective_n_threads(n_threads=None, only_physical_cores=False):
     """Determine the effective number of threads to be used for OpenMP calls
 
     - For ``n_threads = None``,
@@ -47,7 +47,10 @@ cpdef _openmp_effective_n_threads(n_threads=None):
         # to exceed the number of cpus.
         max_n_threads = omp_get_max_threads()
     else:
-        max_n_threads = min(omp_get_max_threads(), cpu_count())
+        max_n_threads = min(
+            omp_get_max_threads(),
+            cpu_count(only_physical_cores=only_physical_cores)
+        )
 
     if n_threads is None:
         return max_n_threads