diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 29ea90f4..c23505c6 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -125,7 +125,7 @@ jobs: # Linux environment with development versions of numpy and scipy - name: pylatest_pip_dev - os : ubuntu-latest + os : ubuntu-24.04 PACKAGER: "pip-dev" PYTHON_VERSION: "*" CC_OUTER_LOOP: "gcc" @@ -134,7 +134,7 @@ jobs: # Linux + Python 3.10 and homogeneous runtime nesting. - name: py310_conda_openblas_clang_clang - os: ubuntu-latest + os: ubuntu-24.04 PACKAGER: "conda" PYTHON_VERSION: "3.10" BLAS: "openblas" @@ -145,7 +145,7 @@ jobs: # Linux environment with MKL and Clang (known to be unsafe for # threadpoolctl) to only test the warning from multiple OpenMP. - name: pylatest_conda_mkl_clang_gcc - os: ubuntu-latest + os: ubuntu-24.04 PYTHON_VERSION: "*" PACKAGER: "conda" BLAS: "mkl" @@ -154,7 +154,7 @@ jobs: TESTS: "libomp_libiomp_warning" # Linux environment with MKL, safe for threadpoolctl. - name: pylatest_conda_mkl_gcc_gcc - os: ubuntu-latest + os: ubuntu-24.04 PYTHON_VERSION: "*" PACKAGER: "conda" BLAS: "mkl" @@ -166,7 +166,7 @@ jobs: # Linux + Python 3.11 with numpy / scipy installed with pip from PyPI # and heterogeneous OpenMP runtimes. - name: py311_pip_openblas_gcc_clang - os: ubuntu-latest + os: ubuntu-24.04 PACKAGER: "pip" PYTHON_VERSION: "3.11" CC_OUTER_LOOP: "gcc" @@ -174,7 +174,7 @@ jobs: # Linux + conda-forge Python RC, extra deps from PyPI. # OpenCV is omitted: no wheels for the current RC ABI yet. - name: pyrc_pip - os: ubuntu-latest + os: ubuntu-24.04 PACKAGER: "pip" # Using the "*" marker does not select the latest RC release # from the conda-forge `python_rc` channel hence the need to pin @@ -185,7 +185,7 @@ jobs: CC_INNER_LOOP: "gcc" # Linux + conda-forge free-threaded Python RC, extra deps from PyPI. - name: pyrc_ft_pip - os: ubuntu-latest + os: ubuntu-24.04 PACKAGER: "pip" # Using the "*" marker does not select the latest RC release # from the conda-forge `python_rc` channel hence the need to pin @@ -197,7 +197,7 @@ jobs: CC_INNER_LOOP: "gcc" # Linux environment with numpy from conda-forge channel and openblas-openmp - name: pylatest_conda_forge - os: ubuntu-latest + os: ubuntu-24.04 PACKAGER: "conda-forge" PYTHON_VERSION: "*" BLAS: "openblas" @@ -207,7 +207,7 @@ jobs: CC_INNER_LOOP: "gcc" # Linux environment with no numpy and heterogeneous OpenMP runtimes. - name: pylatest_conda_nonumpy_gcc_clang - os: ubuntu-latest + os: ubuntu-24.04 PACKAGER: "conda" PYTHON_VERSION: "*" NO_NUMPY: "true" @@ -216,7 +216,7 @@ jobs: # Linux environments with numpy linked to BLIS - name: pylatest_blis_gcc_clang_openmp - os: ubuntu-latest + os: ubuntu-24.04 PYTHON_VERSION: "*" INSTALL_BLAS: "blis" BLIS_NUM_THREAEDS: "4" @@ -225,7 +225,7 @@ jobs: BLIS_CC: "clang-18" BLIS_ENABLE_THREADING: "openmp" - name: pylatest_blis_clang_gcc_pthreads - os: ubuntu-latest + os: ubuntu-24.04 PYTHON_VERSION: "*" INSTALL_BLAS: "blis" BLIS_NUM_THREADS: "4" @@ -234,7 +234,7 @@ jobs: BLIS_CC: "gcc-12" BLIS_ENABLE_THREADING: "pthreads" - name: pylatest_blis_no_threading - os: ubuntu-latest + os: ubuntu-24.04 PYTHON_VERSION: "*" INSTALL_BLAS: "blis" BLIS_NUM_THREADS: "1" @@ -245,7 +245,7 @@ jobs: # Linux env with FlexiBLAS - name: pylatest_flexiblas - os: ubuntu-latest + os: ubuntu-24.04 PYTHON_VERSION: "*" INSTALL_BLAS: "flexiblas" PLATFORM_SPECIFIC_PACKAGES: "mkl" @@ -254,7 +254,7 @@ jobs: # Python 3.14, free-threaded - name: py314_freethreaded_pip_openblas - os: ubuntu-latest + os: ubuntu-24.04 PACKAGER: "conda-forge" BLAS: "openblas" PYTHON_VERSION: "3.14" diff --git a/CHANGES.md b/CHANGES.md index eb539e0d..8f1bbc12 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,6 +1,11 @@ -3.9.0 (under development) +3.8.1 (under development) ========================= +- Fix a number of race conditions in OpenBLAS-pthreads/win32 controlling thread + logic, that could result in unexpected number of threads being set. + https://github.com/joblib/threadpoolctl/pull/300 + + 3.8.0 (2026-11-05) ================== diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 93c32292..372d09f8 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -8,7 +8,7 @@ import subprocess import sys import time -from threading import Lock, Thread, current_thread +from threading import Event, Lock, Thread, current_thread import pytest @@ -1063,7 +1063,7 @@ def test_restore_original_limits_keeps_openblas_controlling_thread(request): libs, lib, orig = _require_resizable_non_openmp_openblas() request.addfinalizer(_restore_openblas_thread_counts(libs)) - limiter = threadpool_limits(limits=1) + limiter = ThreadpoolController().select(internal_api="openblas").limit(limits=1) assert lib.get_num_threads() == 1 assert lib._SHARED_STATE.controlling_thread is current_thread() @@ -1075,6 +1075,144 @@ def test_restore_original_limits_keeps_openblas_controlling_thread(request): assert lib._SHARED_STATE.controlling_thread is None +def _race_condition_skip_problem_platforms(): + flexiblas_controller = ThreadpoolController().select(internal_api="flexiblas") + + if sys.platform == "win32": + pytest.xfail("OpenMP on Windows has its own issues with overlapping setting") + if flexiblas_controller: + pytest.skip( + "FlexiBLAS means there are two controllers for same OpenBLAS " + "(direct + FlexiBLAS) which causes test failures." + ) + + +def test_overlapping_context_manager_restores_correctly(request): + """A thread that failed to set should also not restore.""" + _race_condition_skip_problem_platforms() + + libs, lib, orig_limit = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + openmp = ThreadpoolController().select(user_api="openmp") + original_openmp_info = openmp.info() + + a_in, b_in, a_out = (Event() for _ in range(3)) + + a_limits_correct = [] + b_limits_correct = [] + + def A(): + with threadpool_limits(2): + a_in.set() + b_in.wait() + a_limits_correct.append(openmp.info() == original_openmp_info) + a_out.set() + + def B(): + a_in.wait() + with threadpool_limits(1): # snapshot taken while A's limit is active + b_in.set() + a_out.wait() + b_limits_correct.append(openmp.info() == original_openmp_info) + + ts = [Thread(target=A), Thread(target=B)] + for t in ts: + t.start() + for t in ts: + t.join() + + assert a_limits_correct == [True] + assert b_limits_correct == [True] + assert lib.get_num_threads() == orig_limit + assert openmp.info() == original_openmp_info + + +def test_overlapping_set_and_restore(request): + """A thread that failed to set should also not restore.""" + _race_condition_skip_problem_platforms() + + libs, lib, orig_limit = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + openmp = ThreadpoolController().select(user_api="openmp") + original_openmp_info = openmp.info() + + a_in, b_in, a_out = (Event() for _ in range(3)) + + a_limits_correct = [] + b_limits_correct = [] + + def A(): + limiter = threadpool_limits(2) + a_in.set() + b_in.wait() + limiter.restore_original_limits() + a_limits_correct.append(openmp.info() == original_openmp_info) + a_out.set() + + def B(): + a_in.wait() + limiter = threadpool_limits(1) # snapshot taken while A's limit is active + b_in.set() + a_out.wait() + limiter.restore_original_limits() + b_limits_correct.append(openmp.info() == original_openmp_info) + + ts = [Thread(target=A), Thread(target=B)] + for t in ts: + t.start() + for t in ts: + t.join() + + assert a_limits_correct == [True] + assert b_limits_correct == [True] + assert lib.get_num_threads() == orig_limit + assert openmp.info() == original_openmp_info + + +def test_get_current_info_race_condition(request): + """Test for race condition in multi-threaded OpenBLAS-pthreads usage. + + 1. Thread A sets limit L. + 2. Thread B gets limit L in `_ThreadpoolLimiter.__init__`. + 3. Thread A restores to original limit O and releases control. + 4. Thread B continues `__init__`, setting a limit and taking control. + 5. Thread B restores to limit L. + """ + _race_condition_skip_problem_platforms() + + libs, lib, orig_limit = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + a_in, b_retrieved, a_out = (Event() for _ in range(3)) + + def A(): + with threadpool_limits(2): + a_in.set() + b_retrieved.wait() + a_out.set() + + class racy_threadpool_limits(threadpool_limits): + def _set_threadpool_limits(self): + # Tell A it can finish up: + b_retrieved.set() + # Wait until thread A is finished, so we can take control: + a_out.wait() + super()._set_threadpool_limits() + + def B(): + a_in.wait() + with racy_threadpool_limits(1): + pass + + ts = [Thread(target=A), Thread(target=B)] + for t in ts: + t.start() + for t in ts: + t.join() + + assert lib.get_num_threads() == orig_limit + + # skip test if not run in a azure pipelines job since it relies on a specific flexiblas # installation. @pytest.mark.skipif( diff --git a/threadpoolctl.py b/threadpoolctl.py index 40377c03..d59f795e 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -201,6 +201,21 @@ def get_and_set() -> None: OpenBLASController.clear_controlling_thread(force=True) +# Returned from set_num_threads() to indicate that setting did not happen, +# typically because the current thread is not the controlling thread. +_DID_NOT_SET = object() + + +@dataclass +class _THIS_THREAD_TOOK_CONTROL: + """Indicates the current thread took control when calling ``set_num_threads()``. + + Includes number of threads set previously, right after control was released. + """ + + previous_num_threads: int + + class LibController(ABC): """Abstract base class for the individual library controllers @@ -550,6 +565,9 @@ def clear_controlling_thread(cls, force: bool = False): cls._SHARED_STATE.controlling_thread = None def set_num_threads(self, num_threads): + self._set_num_threads_with_extra_info(num_threads) + + def _set_num_threads_with_extra_info(self, num_threads): if self.threading_layer == "openmp": # Per-thread limits, so no need to be constraining: return self._real_set_num_threads(num_threads) @@ -563,6 +581,7 @@ def set_num_threads(self, num_threads): return self._real_set_num_threads(num_threads) elif controlling_thread is None or not controlling_thread.is_alive(): # No thread controls OpenBLAS, so try to take control: + previous_num_threads = None with self._SHARED_STATE.lock: # Another thread might have changed it, so ensure that's not # the case: @@ -571,12 +590,16 @@ def set_num_threads(self, num_threads): or not self._SHARED_STATE.controlling_thread.is_alive() ): self._SHARED_STATE.controlling_thread = this_thread - # Controlling thread should now be set, to either this thread or - # another, so try again: - return self.set_num_threads(num_threads) + previous_num_threads = self.get_num_threads() + if previous_num_threads is not None: + self._real_set_num_threads(num_threads) + return _THIS_THREAD_TOOK_CONTROL(previous_num_threads) + else: + # Another thread took control: + return _DID_NOT_SET else: # Some other thread controls OpenBLAS: do nothing. - return + return _DID_NOT_SET def _real_set_num_threads(self, num_threads): # The OpenBLAS limiting API is process-wide, and we want current thread @@ -656,6 +679,8 @@ def decrement(self): """Decrement ownership.""" self.count -= 1 if self.count <= 0: + # If we're not currently the controlling thread, this will do + # nothing, which is the desired behavior: OpenBLASController.clear_controlling_thread() @@ -1105,7 +1130,9 @@ def __init__(self, controller, *, limits=None, user_api=None): self._limits, self._user_api = _check_limiter_params( controller, limits, user_api ) - self._original_info = self._controller.info() + self._original_info = dict( + zip(self._controller.lib_controllers, self._controller.info()) + ) self._set_threadpool_limits() def __enter__(self): @@ -1125,9 +1152,7 @@ def wrap(cls, controller, *, limits=None, user_api=None): def restore_original_limits(self): """Set the limits back to their original values""" - for lib_controller, original_info in zip( - self._controller.lib_controllers, self._original_info - ): + for lib_controller, original_info in self._original_info.items(): lib_controller.set_num_threads(original_info["num_threads"]) # Alias of `restore_original_limits` for backward compatibility @@ -1144,7 +1169,7 @@ def get_original_num_threads(self): for user_api in self._user_api: limits = [ lib_info["num_threads"] - for lib_info in self._original_info + for lib_info in self._original_info.values() if lib_info["user_api"] == user_api ] limits = set(limits) @@ -1188,7 +1213,22 @@ def _set_threadpool_limits(self): continue if num_threads is not None: - lib_controller.set_num_threads(num_threads) + if type(lib_controller) is OpenBLASController: + # Special-case API that returns additional information: + result = lib_controller._set_num_threads_with_extra_info( + num_threads + ) + else: + result = lib_controller.set_num_threads(num_threads) + if result is _DID_NOT_SET: + # Didn't set, so shouldn't restore either; if this thread + # can become the controlling thread during restoration, it + # might restore the wrong value. + self._original_info.pop(lib_controller, None) + elif isinstance(result, _THIS_THREAD_TOOK_CONTROL): + self._original_info[lib_controller]["num_threads"] = ( + result.previous_num_threads + ) class _ThreadpoolLimiterDecorator(ContextDecorator):