Merge pull request #14441 from rgommers/testhook-threadpoolctl

ENH: TST: add a threadpoolctl hook to limit OpenBLAS parallelism

.github/workflows/linux.yml

@@ -37,7 +37,7 @@ jobs:
           python3.8-dbg -c 'import sys; print("Python debug build:", hasattr(sys, "gettotalrefcount"))'
           python3.8-dbg -m pip install --upgrade pip 'setuptools<58.5' wheel
           python3.8-dbg -m pip install --upgrade numpy cython pytest pytest-xdist pybind11
-          python3.8-dbg -m pip install --upgrade mpmath gmpy2==2.1.0rc1 pythran
+          python3.8-dbg -m pip install --upgrade mpmath gmpy2==2.1.0rc1 pythran threadpoolctl
           python3.8-dbg -m pip uninstall -y nose
           cd ..
       - name: Building SciPy

azure-pipelines.yml

@@ -291,6 +291,7 @@ stages:
         pytest-env
         pytest-timeout
         pytest-xdist==1.34.0
+        threadpoolctl
       displayName: 'Install dependencies'
     # DLL resolution mechanics were changed in
     # Python 3.8: https://bugs.python.org/issue36085

ci/azure-travis-template.yaml

@@ -68,6 +68,7 @@ steps:
     pip install ${{parameters.other_spec}}
     cython
     gmpy2==2.1.0rc1
+    threadpoolctl
     mpmath
     pythran
     pybind11

doc/source/dev/core-dev/distributing.rst.inc

@@ -28,7 +28,8 @@ are:
 - Pillow_ (for image loading/saving)
 - scikits.umfpack_ (optionally used in ``sparse.linalg``)
 - mpmath_ (for more extended tests in ``special``)
-- pydata/sparse  (compatibility support in ``scipy.sparse``)
+- pydata/sparse (compatibility support in ``scipy.sparse``)
+- threadpoolctl_ (to control BLAS/LAPACK threading in test suite)
 
 *Unconditional build-time dependencies:*
 
@@ -221,3 +222,4 @@ Wheelhouse_, see at the wheel_ and Wheelhouse_ docs.
 .. _six: https://pypi.org/project/six
 .. _decorator: https://github.com/micheles/decorator
 .. _manylinux: https://github.com/pypa/manylinux/
+.. _threadpoolctl: https://github.com/joblib/threadpoolctl

environment.yml

@@ -36,3 +36,4 @@ dependencies:
   # Some optional test dependencies
   - mpmath
   - gmpy2=2.1.0rc1
+  - threadpoolctl

mypy.ini

@@ -70,6 +70,9 @@ ignore_missing_imports = True
 [mypy-mpmath]
 ignore_missing_imports = True
 
+[mypy-threadpoolctl]
+ignore_missing_imports = True
+
 #
 # Extension modules without stubs.
 #

pyproject.toml

@@ -86,6 +86,7 @@ test = [
     "asv",
     "mpmath",
     "gmpy2==2.1.0rc1",
+    "threadpoolctl",
     "scikit-umfpack",
 ]
 doc = [

scipy/conftest.py

@@ -5,6 +5,7 @@
 
 from distutils.version import LooseVersion
 import numpy as np
+import numpy.testing as npt
 from scipy._lib._fpumode import get_fpu_mode
 from scipy._lib._testutils import FPUModeChangeWarning
 
@@ -39,6 +40,40 @@ def pytest_runtest_setup(item):
     if mark is not None and np.intp(0).itemsize < 8:
         pytest.xfail('Fails on our 32-bit test platform(s): %s' % (mark.args[0],))
 
+    # Older versions of threadpoolctl have an issue that may lead to this
+    # warning being emitted, see gh-14441
+    with npt.suppress_warnings() as sup:
+        sup.filter(pytest.PytestUnraisableExceptionWarning)
+
+        try:
+            from threadpoolctl import threadpool_limits
+
+            HAS_THREADPOOLCTL = True
+        except Exception:  # observed in gh-14441: (ImportError, AttributeError)
+            # Optional dependency only. All exceptions are caught, for robustness
+            HAS_THREADPOOLCTL = False
+
+        if HAS_THREADPOOLCTL:
+            # Set the number of openmp threads based on the number of workers
+            # xdist is using to prevent oversubscription. Simplified version of what
+            # sklearn does (it can rely on threadpoolctl and its builtin OpenMP helper
+            # functions)
+            try:
+                xdist_worker_count = int(os.environ['PYTEST_XDIST_WORKER_COUNT'])
+            except KeyError:
+                # raises when pytest-xdist is not installed
+                return
+
+            if not os.getenv('OMP_NUM_THREADS'):
+                max_openmp_threads = os.cpu_count() // 2  # use nr of physical cores
+                threads_per_worker = max(max_openmp_threads // xdist_worker_count, 1)
+                try:
+                    threadpool_limits(threads_per_worker, user_api='blas')
+                except Exception:
+                    # May raise AttributeError for older versions of OpenBLAS.
+                    # Catch any error for robustness.
+                    return
+
 
 @pytest.fixture(scope="function", autouse=True)
 def check_fpu_mode(request):