From e2caf0c76e97c9b97333b1f2716c1d94983cafd5 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Thu, 29 Oct 2026 18:39:57 +0100 Subject: [PATCH 01/21] FIX Do not drop an OpenBLAS claim a limiter scope did not take. A later context manager cleared control taken by a plain threadpool_limits() call or a direct set_num_threads(), and every OpenBLAS copy shared one controlling thread. Co-authored-by: Cursor --- CHANGES.md | 12 +- README.md | 12 +- .../check_no_test_skipped.py | 2 + tests/test_threadpoolctl.py | 234 +++++++++++- threadpoolctl.py | 334 ++++++++++++++---- 5 files changed, 494 insertions(+), 100 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 7317c063..72fd9bd2 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -8,10 +8,14 @@ - When OpenBLAS uses a non-OpenMP threading layer, you can only set the size of its thread pool from a single controlling thread; since it has a shared process-wide thread pool, having multiple Python threads fight over control - wouldn't work well. A new API - ``OpenBLASController.clear_controlling_thread()`` allows resetting which - thread controlls OpenBLAS in these cases, and decorator and context manager - limit APIs will automatically call that API when they are done. + wouldn't work well. Each loaded OpenBLAS library has its own pool and its own + controlling thread. A plain ``threadpool_limits(...)`` call or a direct + ``set_num_threads()`` keeps that role until + ``OpenBLASController.clear_controlling_thread()``. Context managers and + decorators release it when the outermost scope that took the claim exits. A + later scope does not drop a claim it did not take, for example an OpenMP-only + limit or ``limits=None``. A thread that has finished no longer holds the + claim; the next ``set_num_threads()`` can take it. https://github.com/joblib/threadpoolctl/pull/275 https://github.com/joblib/threadpoolctl/pull/289 diff --git a/README.md b/README.md index fcc26221..989abbb2 100644 --- a/README.md +++ b/README.md @@ -322,7 +322,7 @@ with ( You can also operate without a context manager: ```python -from threadpoolctl import ThreadpoolController +from threadpoolctl import OpenBLASController, ThreadpoolController controller = ThreadpoolController() try: @@ -332,6 +332,10 @@ try: pool.map(somefunc, someargs) finally: limiter.restore_original_limits() + # restore puts the thread counts back. It leaves this thread in control + # of any non-OpenMP OpenBLAS it claimed. Drop that when other threads + # should be allowed to change those limits again. + OpenBLASController.clear_controlling_thread() ``` ### Switching Back And Forth Between Main Thread and Python Threads @@ -381,8 +385,10 @@ with CONTROLLER.limit(limits=2): ### OpenBLAS with pthreads/Windows As mentioned above, OpenBLAS with pthreads or Windows threading layer uses a -process-wide shared thread pool. As a result, there are only two limits that -make sense in most circumstances: +process-wide shared thread pool. Each loaded copy has its own pool and its own +controlling thread, so a limit on one copy does not decide who may limit +another. As a result, there are only two limits that make sense in most +circumstances: * A limit of `1`, which effectively means no OpenBLAS thread pool: all operations run inline in the calling thread. This is a good setting when you diff --git a/continuous_integration/check_no_test_skipped.py b/continuous_integration/check_no_test_skipped.py index 02343d19..1e1f0d9e 100644 --- a/continuous_integration/check_no_test_skipped.py +++ b/continuous_integration/check_no_test_skipped.py @@ -36,6 +36,8 @@ # numpy or scipy and we don't want to pin old versions of these libraries. SAFE_SKIPPED_TESTS = [ "test_multiple_shipped_openblas", + # Runs only when two distinct non-OpenMP OpenBLAS copies are loaded. + "test_openblas_controlling_thread_is_per_library", "test_threadpool_limits_by_prefix[1-libblas]", "test_threadpool_limits_by_prefix[3-libblas]", ] diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 82c918f9..8c3c08f6 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -828,18 +828,44 @@ def set_and_exit_after_lock_released(): t.join() -def _openblas_controlling_thread(): - openblas_controller = ThreadpoolController().select(internal_api="openblas") - - if not (openblas_controller): +def _non_openmp_openblas(): + openblas = ThreadpoolController().select(internal_api="openblas") + if not openblas: pytest.skip("requires OpenBLAS.") - if {lc.threading_layer for lc in openblas_controller.lib_controllers} == {"openmp"}: + libs = [] + seen = set() + for lib in openblas.lib_controllers: + if lib.threading_layer == "openmp" or lib.filepath in seen: + continue + seen.add(lib.filepath) + libs.append(lib) + if not libs: pytest.skip("require non-OpenMP OpenBLAS.") - return OpenBLASController._SHARED_STATE.controlling_thread + return libs -def test_openblas_controlling_thread_cleared_on_limiter_context_manager_exit(): +def _openblas_controlling_thread(): + libs = _non_openmp_openblas() + holders = {lib._runtime_state().controlling_thread for lib in libs} + if len(holders) != 1: + return holders + return next(iter(holders)) + + +def _restore_openblas_thread_counts(libs): + originals = [(lib, lib.get_num_threads()) for lib in libs] + + def restore(): + OpenBLASController.clear_controlling_thread(force=True) + for lib, orig in originals: + if orig is not None: + lib.set_num_threads(orig) + OpenBLASController.clear_controlling_thread(force=True) + + return restore + +def test_openblas_controlling_thread_cleared_on_limiter_context_manager_exit(): with threadpool_limits(limits=1): assert _openblas_controlling_thread() is current_thread() with threadpool_limits(limits=1): @@ -849,7 +875,6 @@ def test_openblas_controlling_thread_cleared_on_limiter_context_manager_exit(): def test_openblas_controlling_thread_cleared_on_decorator_exit(): - @threadpool_limits.wrap(limits=1) def inner(): assert _openblas_controlling_thread() is current_thread() @@ -865,19 +890,190 @@ def outer(): def test_openblas_controlling_thread_manual_increment_decrement_api(request): - limiter = ThreadpoolController().limit() - request.addfinalizer(limiter.restore_original_limits) + # The manual counter has to be incremented before set_num_threads, one + # count per library. That makes the claim scoped, so the last decrement + # releases it. + libs = _non_openmp_openblas() + lib = libs[0] + orig = lib.get_num_threads() + if orig is None or orig < 2: + pytest.skip("Test requires a non-OpenMP OpenBLAS with at least 2 threads") + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + key = lib._state_key() + CONTROLLING_THREAD_OWNERSHIP.increment(key) + CONTROLLING_THREAD_OWNERSHIP.increment(key) + lib.set_num_threads(1) + assert lib._runtime_state().controlling_thread is current_thread() + assert not lib._runtime_state().sticky + CONTROLLING_THREAD_OWNERSHIP.decrement(key) + assert lib._runtime_state().controlling_thread is current_thread() + CONTROLLING_THREAD_OWNERSHIP.decrement(key) + assert lib._runtime_state().controlling_thread is None + + +def _require_resizable_openblas(): + libs = _non_openmp_openblas() + lib = libs[0] + orig = lib.get_num_threads() + if orig is None or orig < 2: + pytest.skip("Test requires a non-OpenMP OpenBLAS with at least 2 threads") + return libs, lib, orig + + +def test_plain_threadpool_limits_claim_survives_later_scopes(request): + # A plain call holds the controlling thread. A later with-block, including + # one that does not touch OpenBLAS, must not drop that claim. + libs, lib, orig = _require_resizable_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) threadpool_limits(limits=1) - CONTROLLING_THREAD_OWNERSHIP.increment() - assert _openblas_controlling_thread() is current_thread() - threadpool_limits(limits=1) - CONTROLLING_THREAD_OWNERSHIP.increment() - assert _openblas_controlling_thread() is current_thread() - CONTROLLING_THREAD_OWNERSHIP.decrement() - assert _openblas_controlling_thread() is current_thread() - CONTROLLING_THREAD_OWNERSHIP.decrement() - assert _openblas_controlling_thread() is None + assert lib._runtime_state().controlling_thread is current_thread() + assert lib.get_num_threads() == 1 + + with threadpool_limits(limits=1, user_api="openmp"): + assert lib.get_num_threads() == 1 + with threadpool_limits(limits=None): + assert lib.get_num_threads() == 1 + with threadpool_limits(limits=orig): + assert lib._runtime_state().controlling_thread is current_thread() + + assert lib.get_num_threads() == 1 + assert lib._runtime_state().controlling_thread is current_thread() + + observed = {} + + def other_thread(): + lib.set_num_threads(orig) + observed["threads"] = lib.get_num_threads() + observed["holder"] = lib._runtime_state().controlling_thread + + run_in_thread(other_thread) + assert observed["threads"] == 1 + assert observed["holder"] is current_thread() + + # An explicit clear drops the plain call. The next with-block is then the + # sole scope and releases the claim on the way out. + OpenBLASController.clear_controlling_thread() + assert lib._runtime_state().controlling_thread is None + with threadpool_limits(limits=1): + assert lib._runtime_state().controlling_thread is current_thread() + assert lib._runtime_state().controlling_thread is None + + +def test_direct_set_num_threads_is_not_released_by_limiter_scopes(request): + libs, lib, orig = _require_resizable_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + lib.set_num_threads(1) + state = lib._runtime_state() + assert state.controlling_thread is current_thread() + assert state.sticky + + with threadpool_limits(limits=1, user_api="openmp"): + pass + with threadpool_limits(limits=None): + pass + with threadpool_limits(limits=orig): + assert state.controlling_thread is current_thread() + + assert lib.get_num_threads() == 1 + assert state.sticky + assert state.controlling_thread is current_thread() + + observed = {} + + def other_thread(): + lib.set_num_threads(orig) + observed["threads"] = lib.get_num_threads() + + run_in_thread(other_thread) + assert observed["threads"] == 1 + + +def test_restore_original_limits_keeps_openblas_controlling_thread(request): + libs, lib, orig = _require_resizable_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + limiter = threadpool_limits(limits=1) + assert lib.get_num_threads() == 1 + assert lib._runtime_state().controlling_thread is current_thread() + + limiter.restore_original_limits() + assert lib.get_num_threads() == orig + assert lib._runtime_state().controlling_thread is current_thread() + + OpenBLASController.clear_controlling_thread() + assert lib._runtime_state().controlling_thread is None + + +def test_openblas_controlling_thread_is_per_library(request): + libs = _non_openmp_openblas() + if len(libs) < 2: + pytest.skip("need two non-OpenMP OpenBLAS libraries") + first, second = libs[0], libs[1] + if any( + n is None or n < 2 for n in (first.get_num_threads(), second.get_num_threads()) + ): + pytest.skip("Test requires non-OpenMP OpenBLAS with at least 2 threads") + request.addfinalizer(_restore_openblas_thread_counts([first, second])) + + first.set_num_threads(1) + observed = {} + + def other_thread(): + second.set_num_threads(1) + observed["holder"] = second._runtime_state().controlling_thread + observed["threads"] = second.get_num_threads() + + run_in_thread(other_thread) + assert first.get_num_threads() == 1 + assert first._runtime_state().controlling_thread is current_thread() + assert observed["threads"] == 1 + assert observed["holder"] is not None + assert observed["holder"] is not current_thread() + + +def test_openblas_scoped_exit_releases_only_that_library(): + # Bookkeeping for the ownership leak, including where no pthreads OpenBLAS + # is installed. A scope ending on one library must not drop a sticky claim + # or a still-held claim on another library. + key_a = "\0threadpoolctl-test-openblas-a" + key_b = "\0threadpoolctl-test-openblas-b" + state_a = threadpoolctl._OpenBLASLimitState() + state_b = threadpoolctl._OpenBLASLimitState() + OpenBLASController._states[key_a] = state_a + OpenBLASController._states[key_b] = state_b + try: + state_a.controlling_thread = current_thread() + state_a.sticky = True + + CONTROLLING_THREAD_OWNERSHIP.increment(key_b) + state_b.controlling_thread = current_thread() + state_b.sticky = CONTROLLING_THREAD_OWNERSHIP.count_for(key_b) == 0 + assert state_b.sticky is False + + CONTROLLING_THREAD_OWNERSHIP.decrement(key_b) + assert state_b.controlling_thread is None + assert state_a.controlling_thread is current_thread() + assert state_a.sticky is True + + # A plain call is non-sticky and is held by a leftover count. Ending a + # scope on B must not release A. + state_a.sticky = False + CONTROLLING_THREAD_OWNERSHIP.increment(key_a) + CONTROLLING_THREAD_OWNERSHIP.increment(key_b) + state_b.controlling_thread = current_thread() + state_b.sticky = False + CONTROLLING_THREAD_OWNERSHIP.decrement(key_b) + assert state_b.controlling_thread is None + assert state_a.controlling_thread is current_thread() + assert state_a.sticky is False + assert CONTROLLING_THREAD_OWNERSHIP.count_for(key_a) == 1 + finally: + OpenBLASController._states.pop(key_a, None) + OpenBLASController._states.pop(key_b, None) + CONTROLLING_THREAD_OWNERSHIP.reset() # skip test if not run in a azure pipelines job since it relies on a specific flexiblas diff --git a/threadpoolctl.py b/threadpoolctl.py index 756168cb..96c92e97 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -462,6 +462,16 @@ def wrapper(self): _CDLL_CACHE = _CDLLCache() +@dataclass +class _OpenBLASLimitState: + """Who may change one OpenBLAS library's thread pool.""" + + controlling_thread: Thread | None = None + # Set when control is taken outside a limiter scope (direct set_num_threads). + sticky: bool = False + lock: Lock = field(default_factory=Lock) + + class OpenBLASController(LibController): """Controller class for OpenBLAS. @@ -470,22 +480,29 @@ class OpenBLASController(LibController): set by ``threadpoolctl`` is per-thread. Each Python thread can therefore set its own limit without consideration for other Python threads. - In constrast, when the OpenMP threading layer is not used, there is a + In contrast, when the OpenMP threading layer is not used, there is a process-wide shared thread pool, so having multiple Python threads change - the limit at the same time makes no sense. Therefore: + the limit at the same time makes no sense. Each loaded OpenBLAS library + has its own pool and its own controlling thread. Therefore: 1. Thread pool size can only be set by a single thread, the first - thread to call ``set_num_threads()``. In practice this will usually - be the main thread. - - 2. Other threads silently do nothing. - - 3. ``clear_controlling_thread()`` can be used to reset the thread that - can meaningfully use the limiting API. - - ``with threadpool_limits()`` and the function decorator - ``threadpool_limits.wrap()`` will automatically clear the controlling - thread when they're done, including when nested. + thread to call ``set_num_threads()`` on that library. In practice + this will usually be the main thread. + + 2. Other threads silently do nothing. Once the controlling thread has + finished, the next ``set_num_threads()`` can take the claim. + + 3. ``clear_controlling_thread()`` resets the controlling thread. By + default only that thread can do so. ``force=True`` lets any thread + clear every OpenBLAS library. + + A plain ``threadpool_limits(...)`` call, and a direct ``set_num_threads()``, + keep control until ``clear_controlling_thread()``. ``with + threadpool_limits()`` and ``threadpool_limits.wrap()`` release control when + the outermost scope that took the claim exits, including when those scopes + are nested. A scope that does not touch that OpenBLAS library, and a scope + entered while a plain call or a direct ``set_num_threads()`` is still + holding the claim, does not release it. """ user_api = "blas" @@ -529,25 +546,76 @@ def get_num_threads(self): return get_num_threads_func() return None - class _SHARED_STATE: - """Shared state across instances.""" + # One record per loaded library. Separate OpenBLAS copies (for example + # NumPy's and SciPy's) do not share a thread pool, so they do not share a + # controlling thread either. + _states: dict[str, _OpenBLASLimitState] = {} # noqa: RUF012 + _states_lock = Lock() + _claim_generation = 0 + + def _state_key(self) -> str: + if self.filepath is not None: + return self.filepath + return str(id(self.dynlib)) + + def _runtime_state(self) -> _OpenBLASLimitState: + """Return the process-wide claim for this OpenBLAS copy.""" + key = self._state_key() + with self._states_lock: + state = self._states.get(key) + if state is None: + state = _OpenBLASLimitState() + self._states[key] = state + return state - controlling_thread: Thread | None = None - lock = Lock() + @classmethod + def _snapshot_states(cls) -> list[_OpenBLASLimitState]: + with cls._states_lock: + return list(cls._states.values()) + + @classmethod + def _release_scoped_claim(cls, key: str): + """Drop this thread's non-sticky claim on one library. + + A sticky claim was taken by a direct ``set_num_threads()`` while no + limiter scope was active for that library. Scoped exits must leave it + in place. + """ + with cls._states_lock: + state = cls._states.get(key) + if state is None: + return + this = current_thread() + with state.lock: + if state.controlling_thread is this and not state.sticky: + state.controlling_thread = None @classmethod def clear_controlling_thread(cls, force: bool = False): """Reset which thread can set limits (on non-OpenMP threading layer). - It's up the caller to ensure the current controlling thread is done. + Drops the claim on every OpenBLAS library this thread controls, and + resets this thread's ownership count so a later context manager can + release a claim of its own. It's up to the caller to ensure the current + controlling thread is done. - By default, only the controlling thread has an effect when calling this - method. You can have any thread clear the controlling thread by - setting ``force=True``. + By default, only the controlling thread has an effect. Any thread can + clear every library by setting ``force=True``. """ - with cls._SHARED_STATE.lock: - if force or cls._SHARED_STATE.controlling_thread is current_thread(): - cls._SHARED_STATE.controlling_thread = None + this = current_thread() + cleared_this_thread = False + for state in cls._snapshot_states(): + with state.lock: + if force or state.controlling_thread is this: + if state.controlling_thread is this: + cleared_this_thread = True + state.controlling_thread = None + state.sticky = False + if force: + with cls._states_lock: + cls._claim_generation += 1 + if force or cleared_this_thread: + CONTROLLING_THREAD_OWNERSHIP.reset() def set_num_threads(self, num_threads): if self.threading_layer == "openmp": @@ -556,26 +624,31 @@ def set_num_threads(self, num_threads): # A shared process-wide thread pool, so don't want different Python # threads fighting over limits: + CONTROLLING_THREAD_OWNERSHIP.sync() this_thread = current_thread() - controlling_thread = self._SHARED_STATE.controlling_thread + state = self._runtime_state() + controlling_thread = state.controlling_thread if controlling_thread is None or not controlling_thread.is_alive(): - # No thread controls OpenBLAS, so try to take control: - with self._SHARED_STATE.lock: - # Another thread might have changed it, so ensure that's not - # the case: - if ( - self._SHARED_STATE.controlling_thread is None - 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: + # No live thread controls this OpenBLAS, so try to take control. + # A direct call (ownership count still 0) is sticky: a later + # context manager must not drop it. A call made while a limiter + # scope is active is released with that scope. A finished thread + # no longer holds the claim. + with state.lock: + holder = state.controlling_thread + if holder is None or not holder.is_alive(): + state.controlling_thread = this_thread + state.sticky = ( + CONTROLLING_THREAD_OWNERSHIP.count_for(self._state_key()) == 0 + ) + # Controlling thread should now be set, to either this thread + # or another, so try again: return self.set_num_threads(num_threads) elif controlling_thread is this_thread: - # This thread controls OpenBLAS: + # This thread controls this OpenBLAS: return self._real_set_num_threads(num_threads) else: - # Some other thread controlls OpenBLAS: do nothing. + # Some other thread controls this OpenBLAS: return def _real_set_num_threads(self, num_threads): @@ -638,25 +711,63 @@ def _get_architecture(self): class _ControllingThreadOwnership(thread_local): - """Count how many paired limit-unlimits have modified OpenBLAS. + """Per-library count of OpenBLAS limiter scopes on this thread. - Clear controlling thread if it hits zero. + A count is kept separately for each OpenBLAS copy, keyed by filepath. + Hitting zero for a library releases that library's non-sticky controlling + thread. A plain ``threadpool_limits(...)`` call increments and never + decrements, so it holds its claim until + ``OpenBLASController.clear_controlling_thread()``. - In the future this could potentialy also apply to non-OpenBLAS libraries - that have a shared process pool. + The counts live on the thread-local instance. A class-level dict would be + shared by every thread. """ - count = 0 - - def increment(self): - """Increment ownership.""" - self.count += 1 - - def decrement(self): - """Decrement ownership.""" - self.count -= 1 - if self.count <= 0: - OpenBLASController.clear_controlling_thread() + generation = 0 + + def _counts(self) -> dict[str, int]: + counts = getattr(self, "counts", None) + if counts is None: + counts = {} + self.counts = counts + return counts + + def sync(self): + """Drop this thread's counts when another thread force-clears.""" + generation = OpenBLASController._claim_generation + if self.generation != generation: + self.generation = generation + self.counts = {} + + def reset(self): + """Forget every scope on this thread.""" + self.generation = OpenBLASController._claim_generation + self.counts = {} + + def count_for(self, key: str) -> int: + """Active scopes for one OpenBLAS library on this thread.""" + self.sync() + return self._counts().get(key, 0) + + def increment(self, key: str): + """Increment ownership of one library.""" + self.sync() + counts = self._counts() + counts[key] = counts.get(key, 0) + 1 + + def decrement(self, key: str): + """Decrement ownership of one library and release a scoped claim at zero.""" + self.sync() + counts = self._counts() + current = counts.get(key, 0) + if current <= 0: + return + current -= 1 + if current == 0: + counts.pop(key, None) + OpenBLASController._release_scoped_claim(key) + else: + counts[key] = current CONTROLLING_THREAD_OWNERSHIP = _ControllingThreadOwnership() @@ -1047,15 +1158,20 @@ def __init__(self, controller, *, limits=None, user_api=None): self._controller = controller self._limits, self._user_api = self._check_params(limits, user_api) self._original_info = self._controller.info() - self._set_threadpool_limits() + self._openblas_claim_keys: list[str] = [] + # A plain call never reaches __exit__, so this increment is what makes + # the claim stick until clear_controlling_thread(). A `with` block + # decrements the same increment on the way out. + self._apply_limits_with_openblas_claim() def __enter__(self): - CONTROLLING_THREAD_OWNERSHIP.increment() return self def __exit__(self, type, value, traceback): - self.restore_original_limits() - CONTROLLING_THREAD_OWNERSHIP.decrement() + try: + self.restore_original_limits() + finally: + self._release_openblas_claim() @classmethod def wrap(cls, controller, *, limits=None, user_api=None): @@ -1166,24 +1282,94 @@ def _check_params(self, limits, user_api): return limits, user_api + def _num_threads_for(self, lib_controller): + """Limit that applies to ``lib_controller``, or None to leave it alone. + + ``self._limits`` is a dict ``{key: num_threads}`` where key is either a + filepath, a prefix, or a user_api. Libraries are matched in that order. + """ + if self._limits is None: + return None + if lib_controller.filepath in self._limits: + num_threads = self._limits[lib_controller.filepath] + elif lib_controller.prefix in self._limits: + num_threads = self._limits[lib_controller.prefix] + elif lib_controller.user_api in self._limits: + num_threads = self._limits[lib_controller.user_api] + else: + return None + return num_threads + + def _iter_targeted_openblas(self): + """Non-OpenMP OpenBLAS libraries this limiter would actually set.""" + for lib in self._controller.lib_controllers: + if ( + isinstance(lib, OpenBLASController) + and lib.threading_layer != "openmp" + and self._num_threads_for(lib) is not None + ): + yield lib + + def _acquire_openblas_claim(self): + """Count a scope per library before setting, so the claim is not sticky. + + Only libraries this thread can control are counted. OpenMP-only limits + and ``limits=None`` do not take a claim. Libraries are counted + separately so a scope that limits one copy does not pin another. + """ + self._openblas_claim_keys = [] + this = current_thread() + seen = set() + for lib in self._iter_targeted_openblas(): + key = lib._state_key() + if key in seen: + continue + seen.add(key) + holder = lib._runtime_state().controlling_thread + if holder is not None and holder is not this: + continue + CONTROLLING_THREAD_OWNERSHIP.increment(key) + self._openblas_claim_keys.append(key) + + def _release_openblas_claim(self): + keys = self._openblas_claim_keys + self._openblas_claim_keys = [] + for key in keys: + CONTROLLING_THREAD_OWNERSHIP.decrement(key) + + def _drop_unheld_openblas_claim(self): + """Drop counts for libraries another thread won the race for.""" + this = current_thread() + kept = [] + for key in self._openblas_claim_keys: + with OpenBLASController._states_lock: + state = OpenBLASController._states.get(key) + held = False + if state is not None: + with state.lock: + held = state.controlling_thread is this + if held: + kept.append(key) + else: + CONTROLLING_THREAD_OWNERSHIP.decrement(key) + self._openblas_claim_keys = kept + + def _apply_limits_with_openblas_claim(self): + self._acquire_openblas_claim() + try: + self._set_threadpool_limits() + except Exception: + self._release_openblas_claim() + raise + self._drop_unheld_openblas_claim() + def _set_threadpool_limits(self): """Change the maximal number of threads in selected thread pools.""" if self._limits is None: return for lib_controller in self._controller.lib_controllers: - # self._limits is a dict {key: num_threads} where key is either a - # filepath, a prefix, or a user_api. Libraries are matched in that - # order. - if lib_controller.filepath in self._limits: - num_threads = self._limits[lib_controller.filepath] - elif lib_controller.prefix in self._limits: - num_threads = self._limits[lib_controller.prefix] - elif lib_controller.user_api in self._limits: - num_threads = self._limits[lib_controller.user_api] - else: - continue - + num_threads = self._num_threads_for(lib_controller) if num_threads is not None: lib_controller.set_num_threads(num_threads) @@ -1194,14 +1380,14 @@ class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter, ContextDecorator): def __init__(self, controller, *, limits=None, user_api=None): self._limits, self._user_api = self._check_params(limits, user_api) self._controller = controller + self._openblas_claim_keys = [] def __enter__(self): - # we need to set the limits here and not in the __init__ because we want the - # limits to be set when calling the decorated function, not when creating the - # decorator. + # Set the limits here and not in __init__ so they apply when the + # decorated function is called, not when the decorator is created. + # Increment before setting so this scope's claim is not sticky. self._original_info = self._controller.info() - self._set_threadpool_limits() - CONTROLLING_THREAD_OWNERSHIP.increment() + self._apply_limits_with_openblas_claim() return self From 4a7d12610bcb5ea10530cdebff231086a6b2e26a Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 12:05:56 +0100 Subject: [PATCH 02/21] MNT Keep shared-pool claims scoped to the controller class that owns them. Co-authored-by: Cursor --- tests/test_threadpoolctl.py | 125 +++++++++----- threadpoolctl.py | 323 ++++++++++++++++++++++-------------- 2 files changed, 286 insertions(+), 162 deletions(-) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 8c3c08f6..5d50cd8c 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -16,7 +16,7 @@ from threadpoolctl import ( _ALL_PREFIXES, _ALL_USER_APIS, - CONTROLLING_THREAD_OWNERSHIP, + SHARED_POOL_CONTROL_OWNERSHIP, LibController, OpenBLASController, ThreadpoolController, @@ -846,7 +846,7 @@ def _non_openmp_openblas(): def _openblas_controlling_thread(): libs = _non_openmp_openblas() - holders = {lib._runtime_state().controlling_thread for lib in libs} + holders = {lib._shared_pool_state().controlling_thread for lib in libs} if len(holders) != 1: return holders return next(iter(holders)) @@ -900,16 +900,16 @@ def test_openblas_controlling_thread_manual_increment_decrement_api(request): pytest.skip("Test requires a non-OpenMP OpenBLAS with at least 2 threads") request.addfinalizer(_restore_openblas_thread_counts(libs)) - key = lib._state_key() - CONTROLLING_THREAD_OWNERSHIP.increment(key) - CONTROLLING_THREAD_OWNERSHIP.increment(key) + key = lib._shared_pool_key() + SHARED_POOL_CONTROL_OWNERSHIP.increment(key) + SHARED_POOL_CONTROL_OWNERSHIP.increment(key) lib.set_num_threads(1) - assert lib._runtime_state().controlling_thread is current_thread() - assert not lib._runtime_state().sticky - CONTROLLING_THREAD_OWNERSHIP.decrement(key) - assert lib._runtime_state().controlling_thread is current_thread() - CONTROLLING_THREAD_OWNERSHIP.decrement(key) - assert lib._runtime_state().controlling_thread is None + assert lib._shared_pool_state().controlling_thread is current_thread() + assert not lib._shared_pool_state().sticky + SHARED_POOL_CONTROL_OWNERSHIP.decrement(key) + assert lib._shared_pool_state().controlling_thread is current_thread() + SHARED_POOL_CONTROL_OWNERSHIP.decrement(key) + assert lib._shared_pool_state().controlling_thread is None def _require_resizable_openblas(): @@ -928,7 +928,7 @@ def test_plain_threadpool_limits_claim_survives_later_scopes(request): request.addfinalizer(_restore_openblas_thread_counts(libs)) threadpool_limits(limits=1) - assert lib._runtime_state().controlling_thread is current_thread() + assert lib._shared_pool_state().controlling_thread is current_thread() assert lib.get_num_threads() == 1 with threadpool_limits(limits=1, user_api="openmp"): @@ -936,17 +936,17 @@ def test_plain_threadpool_limits_claim_survives_later_scopes(request): with threadpool_limits(limits=None): assert lib.get_num_threads() == 1 with threadpool_limits(limits=orig): - assert lib._runtime_state().controlling_thread is current_thread() + assert lib._shared_pool_state().controlling_thread is current_thread() assert lib.get_num_threads() == 1 - assert lib._runtime_state().controlling_thread is current_thread() + assert lib._shared_pool_state().controlling_thread is current_thread() observed = {} def other_thread(): lib.set_num_threads(orig) observed["threads"] = lib.get_num_threads() - observed["holder"] = lib._runtime_state().controlling_thread + observed["holder"] = lib._shared_pool_state().controlling_thread run_in_thread(other_thread) assert observed["threads"] == 1 @@ -955,10 +955,10 @@ def other_thread(): # An explicit clear drops the plain call. The next with-block is then the # sole scope and releases the claim on the way out. OpenBLASController.clear_controlling_thread() - assert lib._runtime_state().controlling_thread is None + assert lib._shared_pool_state().controlling_thread is None with threadpool_limits(limits=1): - assert lib._runtime_state().controlling_thread is current_thread() - assert lib._runtime_state().controlling_thread is None + assert lib._shared_pool_state().controlling_thread is current_thread() + assert lib._shared_pool_state().controlling_thread is None def test_direct_set_num_threads_is_not_released_by_limiter_scopes(request): @@ -966,7 +966,7 @@ def test_direct_set_num_threads_is_not_released_by_limiter_scopes(request): request.addfinalizer(_restore_openblas_thread_counts(libs)) lib.set_num_threads(1) - state = lib._runtime_state() + state = lib._shared_pool_state() assert state.controlling_thread is current_thread() assert state.sticky @@ -997,14 +997,14 @@ def test_restore_original_limits_keeps_openblas_controlling_thread(request): limiter = threadpool_limits(limits=1) assert lib.get_num_threads() == 1 - assert lib._runtime_state().controlling_thread is current_thread() + assert lib._shared_pool_state().controlling_thread is current_thread() limiter.restore_original_limits() assert lib.get_num_threads() == orig - assert lib._runtime_state().controlling_thread is current_thread() + assert lib._shared_pool_state().controlling_thread is current_thread() OpenBLASController.clear_controlling_thread() - assert lib._runtime_state().controlling_thread is None + assert lib._shared_pool_state().controlling_thread is None def test_openblas_controlling_thread_is_per_library(request): @@ -1023,12 +1023,12 @@ def test_openblas_controlling_thread_is_per_library(request): def other_thread(): second.set_num_threads(1) - observed["holder"] = second._runtime_state().controlling_thread + observed["holder"] = second._shared_pool_state().controlling_thread observed["threads"] = second.get_num_threads() run_in_thread(other_thread) assert first.get_num_threads() == 1 - assert first._runtime_state().controlling_thread is current_thread() + assert first._shared_pool_state().controlling_thread is current_thread() assert observed["threads"] == 1 assert observed["holder"] is not None assert observed["holder"] is not current_thread() @@ -1040,20 +1040,20 @@ def test_openblas_scoped_exit_releases_only_that_library(): # or a still-held claim on another library. key_a = "\0threadpoolctl-test-openblas-a" key_b = "\0threadpoolctl-test-openblas-b" - state_a = threadpoolctl._OpenBLASLimitState() - state_b = threadpoolctl._OpenBLASLimitState() - OpenBLASController._states[key_a] = state_a - OpenBLASController._states[key_b] = state_b + state_a = threadpoolctl._SharedPoolState() + state_b = threadpoolctl._SharedPoolState() + threadpoolctl._SHARED_POOL_STATES[key_a] = state_a + threadpoolctl._SHARED_POOL_STATES[key_b] = state_b try: state_a.controlling_thread = current_thread() state_a.sticky = True - CONTROLLING_THREAD_OWNERSHIP.increment(key_b) + SHARED_POOL_CONTROL_OWNERSHIP.increment(key_b) state_b.controlling_thread = current_thread() - state_b.sticky = CONTROLLING_THREAD_OWNERSHIP.count_for(key_b) == 0 + state_b.sticky = SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_b) == 0 assert state_b.sticky is False - CONTROLLING_THREAD_OWNERSHIP.decrement(key_b) + SHARED_POOL_CONTROL_OWNERSHIP.decrement(key_b) assert state_b.controlling_thread is None assert state_a.controlling_thread is current_thread() assert state_a.sticky is True @@ -1061,19 +1061,68 @@ def test_openblas_scoped_exit_releases_only_that_library(): # A plain call is non-sticky and is held by a leftover count. Ending a # scope on B must not release A. state_a.sticky = False - CONTROLLING_THREAD_OWNERSHIP.increment(key_a) - CONTROLLING_THREAD_OWNERSHIP.increment(key_b) + SHARED_POOL_CONTROL_OWNERSHIP.increment(key_a) + SHARED_POOL_CONTROL_OWNERSHIP.increment(key_b) state_b.controlling_thread = current_thread() state_b.sticky = False - CONTROLLING_THREAD_OWNERSHIP.decrement(key_b) + SHARED_POOL_CONTROL_OWNERSHIP.decrement(key_b) assert state_b.controlling_thread is None assert state_a.controlling_thread is current_thread() assert state_a.sticky is False - assert CONTROLLING_THREAD_OWNERSHIP.count_for(key_a) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_a) == 1 finally: - OpenBLASController._states.pop(key_a, None) - OpenBLASController._states.pop(key_b, None) - CONTROLLING_THREAD_OWNERSHIP.reset() + threadpoolctl._SHARED_POOL_STATES.pop(key_a, None) + threadpoolctl._SHARED_POOL_STATES.pop(key_b, None) + SHARED_POOL_CONTROL_OWNERSHIP.reset() + + +class _OtherSharedPool: + """Stand-in for a future library with its own shared pool.""" + + +def test_clear_controlling_thread_leaves_other_libraries(): + # OpenBLASController.clear_controlling_thread, including force=True, must + # not drop a claim or a count owned by another controller class. + key_openblas = "\0threadpoolctl-test-openblas-clear" + key_other = "\0threadpoolctl-test-other-clear" + state_openblas = threadpoolctl._SharedPoolState(owner=OpenBLASController) + state_other = threadpoolctl._SharedPoolState(owner=_OtherSharedPool) + threadpoolctl._SHARED_POOL_STATES[key_openblas] = state_openblas + threadpoolctl._SHARED_POOL_STATES[key_other] = state_other + try: + state_openblas.controlling_thread = current_thread() + state_openblas.sticky = True + state_other.controlling_thread = current_thread() + state_other.sticky = True + SHARED_POOL_CONTROL_OWNERSHIP.increment(key_openblas) + SHARED_POOL_CONTROL_OWNERSHIP.increment(key_other) + + OpenBLASController.clear_controlling_thread() + assert state_openblas.controlling_thread is None + assert state_openblas.sticky is False + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_openblas) == 0 + assert state_other.controlling_thread is current_thread() + assert state_other.sticky is True + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_other) == 1 + + state_openblas.controlling_thread = current_thread() + state_openblas.sticky = True + SHARED_POOL_CONTROL_OWNERSHIP.increment(key_openblas) + + def force_clear(): + OpenBLASController.clear_controlling_thread(force=True) + + run_in_thread(force_clear) + assert state_openblas.controlling_thread is None + assert state_openblas.sticky is False + assert state_other.controlling_thread is current_thread() + assert state_other.sticky is True + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_openblas) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_other) == 1 + finally: + threadpoolctl._SHARED_POOL_STATES.pop(key_openblas, None) + threadpoolctl._SHARED_POOL_STATES.pop(key_other, None) + SHARED_POOL_CONTROL_OWNERSHIP.reset() # skip test if not run in a azure pipelines job since it relies on a specific flexiblas diff --git a/threadpoolctl.py b/threadpoolctl.py index 96c92e97..e9b1b156 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -308,6 +308,23 @@ def _get_symbol(self, name): self.dynlib, f"{self._symbol_prefix}{name}{self._symbol_suffix}", None ) + def _has_shared_pool(self): + """Whether this library's thread pool is shared by all Python threads. + + When True, only one thread at a time can change the pool size, and + ``threadpool_limits`` scopes claim control of it on that library. Override + in subclasses whose pool is shared process-wide. + """ + return False + + def _shared_pool_key(self) -> str: + """Identify this library's shared pool, one per loaded copy.""" + return self.filepath + + def _shared_pool_state(self): + """Return the process-wide claim record for this library's shared pool.""" + return _get_shared_pool_state(self._shared_pool_key(), type(self)) + # Internal marker for missing cache lookups: _MISSING = object() @@ -462,14 +479,87 @@ def wrapper(self): _CDLL_CACHE = _CDLLCache() +# Process-wide claims on shared thread pools, keyed by library. A library whose +# pool is shared between Python threads reports it through +# LibController._has_shared_pool(). Separate copies of a library, for example +# NumPy's and SciPy's OpenBLAS, have separate pools, so each copy has its own +# record and its own controlling thread. +_SHARED_POOL_STATES: dict[str, _SharedPoolState] = {} +_SHARED_POOL_STATES_LOCK = Lock() +# Bumped per controller class by a forced clear. Other threads then drop that +# class's counts and keep every other library. +_shared_pool_generations: dict[type, int] = {} + + +def _get_shared_pool_state(key: str, owner: type) -> _SharedPoolState: + with _SHARED_POOL_STATES_LOCK: + state = _SHARED_POOL_STATES.get(key) + if state is None: + state = _SharedPoolState(owner=owner) + _SHARED_POOL_STATES[key] = state + elif state.owner is None: + state.owner = owner + return state + + +def _snapshot_shared_pool_states() -> list[_SharedPoolState]: + with _SHARED_POOL_STATES_LOCK: + return list(_SHARED_POOL_STATES.values()) + + +def _release_shared_pool_claim(key: str): + """Drop this thread's non-sticky claim on one shared pool. + + A sticky claim was taken by a direct ``set_num_threads()`` while no limiter + scope was active for that library. Scoped exits must leave it in place. + """ + with _SHARED_POOL_STATES_LOCK: + state = _SHARED_POOL_STATES.get(key) + if state is None: + return + this = current_thread() + with state.lock: + if state.controlling_thread is this and not state.sticky: + state.controlling_thread = None + + +def _clear_shared_pool_claims(owner: type, force: bool = False): + """Drop claims on shared pools owned by one controller class. + + ``owner`` is that class, for example ``OpenBLASController``. Pools of + other classes are left untouched, including when ``force`` is true. A + forced clear also drops that class's counts on other threads the next + time they sync. + """ + this = current_thread() + cleared_this_thread = False + for state in _snapshot_shared_pool_states(): + if state.owner is not owner: + continue + with state.lock: + if force or state.controlling_thread is this: + if state.controlling_thread is this: + cleared_this_thread = True + state.controlling_thread = None + state.sticky = False + if force: + with _SHARED_POOL_STATES_LOCK: + _shared_pool_generations[owner] = _shared_pool_generations.get(owner, 0) + 1 + if force or cleared_this_thread: + SHARED_POOL_CONTROL_OWNERSHIP.reset(owner) + + @dataclass -class _OpenBLASLimitState: - """Who may change one OpenBLAS library's thread pool.""" +class _SharedPoolState: + """Who may change one shared thread pool.""" controlling_thread: Thread | None = None # Set when control is taken outside a limiter scope (direct set_num_threads). sticky: bool = False lock: Lock = field(default_factory=Lock) + # Controller class that owns this pool, so one class can clear its claims + # without releasing another library. + owner: type | None = None class OpenBLASController(LibController): @@ -546,76 +636,23 @@ def get_num_threads(self): return get_num_threads_func() return None - # One record per loaded library. Separate OpenBLAS copies (for example - # NumPy's and SciPy's) do not share a thread pool, so they do not share a - # controlling thread either. - _states: dict[str, _OpenBLASLimitState] = {} # noqa: RUF012 - _states_lock = Lock() - _claim_generation = 0 - - def _state_key(self) -> str: - if self.filepath is not None: - return self.filepath - return str(id(self.dynlib)) - - def _runtime_state(self) -> _OpenBLASLimitState: - """Return the process-wide claim for this OpenBLAS copy.""" - key = self._state_key() - with self._states_lock: - state = self._states.get(key) - if state is None: - state = _OpenBLASLimitState() - self._states[key] = state - return state - - @classmethod - def _snapshot_states(cls) -> list[_OpenBLASLimitState]: - with cls._states_lock: - return list(cls._states.values()) - - @classmethod - def _release_scoped_claim(cls, key: str): - """Drop this thread's non-sticky claim on one library. - - A sticky claim was taken by a direct ``set_num_threads()`` while no - limiter scope was active for that library. Scoped exits must leave it - in place. - """ - with cls._states_lock: - state = cls._states.get(key) - if state is None: - return - this = current_thread() - with state.lock: - if state.controlling_thread is this and not state.sticky: - state.controlling_thread = None + def _has_shared_pool(self): + return self.threading_layer != "openmp" @classmethod def clear_controlling_thread(cls, force: bool = False): """Reset which thread can set limits (on non-OpenMP threading layer). Drops the claim on every OpenBLAS library this thread controls, and - resets this thread's ownership count so a later context manager can - release a claim of its own. It's up to the caller to ensure the current + resets this thread's OpenBLAS ownership counts so a later context + manager can release a claim of its own. Shared pools of other libraries + are left in place. It's up to the caller to ensure the current controlling thread is done. By default, only the controlling thread has an effect. Any thread can - clear every library by setting ``force=True``. + clear every OpenBLAS library by setting ``force=True``. """ - this = current_thread() - cleared_this_thread = False - for state in cls._snapshot_states(): - with state.lock: - if force or state.controlling_thread is this: - if state.controlling_thread is this: - cleared_this_thread = True - state.controlling_thread = None - state.sticky = False - if force: - with cls._states_lock: - cls._claim_generation += 1 - if force or cleared_this_thread: - CONTROLLING_THREAD_OWNERSHIP.reset() + _clear_shared_pool_claims(cls, force) def set_num_threads(self, num_threads): if self.threading_layer == "openmp": @@ -624,9 +661,9 @@ def set_num_threads(self, num_threads): # A shared process-wide thread pool, so don't want different Python # threads fighting over limits: - CONTROLLING_THREAD_OWNERSHIP.sync() + SHARED_POOL_CONTROL_OWNERSHIP.sync() this_thread = current_thread() - state = self._runtime_state() + state = self._shared_pool_state() controlling_thread = state.controlling_thread if controlling_thread is None or not controlling_thread.is_alive(): # No live thread controls this OpenBLAS, so try to take control. @@ -634,13 +671,15 @@ def set_num_threads(self, num_threads): # context manager must not drop it. A call made while a limiter # scope is active is released with that scope. A finished thread # no longer holds the claim. + # Read the count before taking the pool lock. Syncing ownership + # takes the states lock, and releasing a claim takes that lock + # before the pool lock. + count = SHARED_POOL_CONTROL_OWNERSHIP.count_for(self._shared_pool_key()) with state.lock: holder = state.controlling_thread if holder is None or not holder.is_alive(): state.controlling_thread = this_thread - state.sticky = ( - CONTROLLING_THREAD_OWNERSHIP.count_for(self._state_key()) == 0 - ) + state.sticky = count == 0 # Controlling thread should now be set, to either this thread # or another, so try again: return self.set_num_threads(num_threads) @@ -710,21 +749,19 @@ def _get_architecture(self): return None -class _ControllingThreadOwnership(thread_local): - """Per-library count of OpenBLAS limiter scopes on this thread. +class _SharedPoolControlOwnership(thread_local): + """Per-pool count of limiter scopes on this thread. - A count is kept separately for each OpenBLAS copy, keyed by filepath. - Hitting zero for a library releases that library's non-sticky controlling - thread. A plain ``threadpool_limits(...)`` call increments and never - decrements, so it holds its claim until - ``OpenBLASController.clear_controlling_thread()``. + A count is kept separately for each shared pool, keyed by library filepath. + Hitting zero for a pool releases that pool's non-sticky controlling thread. + A plain ``threadpool_limits(...)`` call increments and never decrements, so + it holds its claim until ``OpenBLASController.clear_controlling_thread()``. The counts live on the thread-local instance. A class-level dict would be - shared by every thread. + shared by every thread. A forced clear drops counts for that controller + class only. """ - generation = 0 - def _counts(self) -> dict[str, int]: counts = getattr(self, "counts", None) if counts is None: @@ -732,31 +769,73 @@ def _counts(self) -> dict[str, int]: self.counts = counts return counts + def _generations(self) -> dict[type, int]: + generations = getattr(self, "generations", None) + if generations is None: + generations = {} + self.generations = generations + return generations + def sync(self): - """Drop this thread's counts when another thread force-clears.""" - generation = OpenBLASController._claim_generation - if self.generation != generation: - self.generation = generation - self.counts = {} + """Drop counts for libraries another thread has force-cleared.""" + with _SHARED_POOL_STATES_LOCK: + generations = dict(_shared_pool_generations) + owned_keys = [ + (key, state.owner) + for key, state in _SHARED_POOL_STATES.items() + if state.owner is not None + ] + local = self._generations() + stale = { + owner + for owner, generation in generations.items() + if local.get(owner, 0) != generation + } + if not stale: + return + counts = self._counts() + for key, owner in owned_keys: + if owner in stale: + counts.pop(key, None) + for owner in stale: + local[owner] = generations[owner] - def reset(self): - """Forget every scope on this thread.""" - self.generation = OpenBLASController._claim_generation - self.counts = {} + def reset(self, owner: type | None = None): + """Forget scopes on this thread. + + With ``owner``, only that controller class. Without, every library. + """ + if owner is None: + with _SHARED_POOL_STATES_LOCK: + generations = dict(_shared_pool_generations) + self.counts = {} + self.generations = generations + return + with _SHARED_POOL_STATES_LOCK: + generation = _shared_pool_generations.get(owner, 0) + keys = [ + key + for key, state in _SHARED_POOL_STATES.items() + if state.owner is owner + ] + counts = self._counts() + for key in keys: + counts.pop(key, None) + self._generations()[owner] = generation def count_for(self, key: str) -> int: - """Active scopes for one OpenBLAS library on this thread.""" + """Active scopes for one shared pool on this thread.""" self.sync() return self._counts().get(key, 0) def increment(self, key: str): - """Increment ownership of one library.""" + """Increment ownership of one shared pool.""" self.sync() counts = self._counts() counts[key] = counts.get(key, 0) + 1 def decrement(self, key: str): - """Decrement ownership of one library and release a scoped claim at zero.""" + """Decrement ownership of one shared pool, releasing its scoped claim at zero.""" self.sync() counts = self._counts() current = counts.get(key, 0) @@ -765,12 +844,12 @@ def decrement(self, key: str): current -= 1 if current == 0: counts.pop(key, None) - OpenBLASController._release_scoped_claim(key) + _release_shared_pool_claim(key) else: counts[key] = current -CONTROLLING_THREAD_OWNERSHIP = _ControllingThreadOwnership() +SHARED_POOL_CONTROL_OWNERSHIP = _SharedPoolControlOwnership() class BLISController(LibController): @@ -1158,11 +1237,11 @@ def __init__(self, controller, *, limits=None, user_api=None): self._controller = controller self._limits, self._user_api = self._check_params(limits, user_api) self._original_info = self._controller.info() - self._openblas_claim_keys: list[str] = [] + self._shared_pool_claim_keys: list[str] = [] # A plain call never reaches __exit__, so this increment is what makes # the claim stick until clear_controlling_thread(). A `with` block # decrements the same increment on the way out. - self._apply_limits_with_openblas_claim() + self._apply_limits_with_shared_pool_claims() def __enter__(self): return self @@ -1171,7 +1250,7 @@ def __exit__(self, type, value, traceback): try: self.restore_original_limits() finally: - self._release_openblas_claim() + self._release_shared_pool_claims() @classmethod def wrap(cls, controller, *, limits=None, user_api=None): @@ -1300,50 +1379,46 @@ def _num_threads_for(self, lib_controller): return None return num_threads - def _iter_targeted_openblas(self): - """Non-OpenMP OpenBLAS libraries this limiter would actually set.""" + def _iter_shared_pool_targets(self): + """Libraries with a shared pool that this limiter would actually set.""" for lib in self._controller.lib_controllers: - if ( - isinstance(lib, OpenBLASController) - and lib.threading_layer != "openmp" - and self._num_threads_for(lib) is not None - ): + if lib._has_shared_pool() and self._num_threads_for(lib) is not None: yield lib - def _acquire_openblas_claim(self): + def _acquire_shared_pool_claims(self): """Count a scope per library before setting, so the claim is not sticky. Only libraries this thread can control are counted. OpenMP-only limits and ``limits=None`` do not take a claim. Libraries are counted separately so a scope that limits one copy does not pin another. """ - self._openblas_claim_keys = [] + self._shared_pool_claim_keys = [] this = current_thread() seen = set() - for lib in self._iter_targeted_openblas(): - key = lib._state_key() + for lib in self._iter_shared_pool_targets(): + key = lib._shared_pool_key() if key in seen: continue seen.add(key) - holder = lib._runtime_state().controlling_thread + holder = lib._shared_pool_state().controlling_thread if holder is not None and holder is not this: continue - CONTROLLING_THREAD_OWNERSHIP.increment(key) - self._openblas_claim_keys.append(key) + SHARED_POOL_CONTROL_OWNERSHIP.increment(key) + self._shared_pool_claim_keys.append(key) - def _release_openblas_claim(self): - keys = self._openblas_claim_keys - self._openblas_claim_keys = [] + def _release_shared_pool_claims(self): + keys = self._shared_pool_claim_keys + self._shared_pool_claim_keys = [] for key in keys: - CONTROLLING_THREAD_OWNERSHIP.decrement(key) + SHARED_POOL_CONTROL_OWNERSHIP.decrement(key) - def _drop_unheld_openblas_claim(self): + def _drop_unheld_shared_pool_claims(self): """Drop counts for libraries another thread won the race for.""" this = current_thread() kept = [] - for key in self._openblas_claim_keys: - with OpenBLASController._states_lock: - state = OpenBLASController._states.get(key) + for key in self._shared_pool_claim_keys: + with _SHARED_POOL_STATES_LOCK: + state = _SHARED_POOL_STATES.get(key) held = False if state is not None: with state.lock: @@ -1351,17 +1426,17 @@ def _drop_unheld_openblas_claim(self): if held: kept.append(key) else: - CONTROLLING_THREAD_OWNERSHIP.decrement(key) - self._openblas_claim_keys = kept + SHARED_POOL_CONTROL_OWNERSHIP.decrement(key) + self._shared_pool_claim_keys = kept - def _apply_limits_with_openblas_claim(self): - self._acquire_openblas_claim() + def _apply_limits_with_shared_pool_claims(self): + self._acquire_shared_pool_claims() try: self._set_threadpool_limits() except Exception: - self._release_openblas_claim() + self._release_shared_pool_claims() raise - self._drop_unheld_openblas_claim() + self._drop_unheld_shared_pool_claims() def _set_threadpool_limits(self): """Change the maximal number of threads in selected thread pools.""" @@ -1380,14 +1455,14 @@ class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter, ContextDecorator): def __init__(self, controller, *, limits=None, user_api=None): self._limits, self._user_api = self._check_params(limits, user_api) self._controller = controller - self._openblas_claim_keys = [] + self._shared_pool_claim_keys = [] def __enter__(self): # Set the limits here and not in __init__ so they apply when the # decorated function is called, not when the decorator is created. # Increment before setting so this scope's claim is not sticky. self._original_info = self._controller.info() - self._apply_limits_with_openblas_claim() + self._apply_limits_with_shared_pool_claims() return self From 24572832060e6320f5b46f534a8b9ca0a357455e Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 12:27:43 +0100 Subject: [PATCH 03/21] MNT Move shared-pool claim handling onto a reusable mixin. Co-authored-by: Cursor --- threadpoolctl.py | 166 +++++++++++++++++++++++++---------------------- 1 file changed, 90 insertions(+), 76 deletions(-) diff --git a/threadpoolctl.py b/threadpoolctl.py index e9b1b156..e88b5b1c 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -312,19 +312,12 @@ def _has_shared_pool(self): """Whether this library's thread pool is shared by all Python threads. When True, only one thread at a time can change the pool size, and - ``threadpool_limits`` scopes claim control of it on that library. Override - in subclasses whose pool is shared process-wide. + ``threadpool_limits`` scopes claim control of it on that library. + Controllers with a process-wide pool inherit ``_SharedPoolControllerMixin``, + which reports True. Override when only some builds share a pool. """ return False - def _shared_pool_key(self) -> str: - """Identify this library's shared pool, one per loaded copy.""" - return self.filepath - - def _shared_pool_state(self): - """Return the process-wide claim record for this library's shared pool.""" - return _get_shared_pool_state(self._shared_pool_key(), type(self)) - # Internal marker for missing cache lookups: _MISSING = object() @@ -479,15 +472,14 @@ def wrapper(self): _CDLL_CACHE = _CDLLCache() -# Process-wide claims on shared thread pools, keyed by library. A library whose -# pool is shared between Python threads reports it through -# LibController._has_shared_pool(). Separate copies of a library, for example -# NumPy's and SciPy's OpenBLAS, have separate pools, so each copy has its own -# record and its own controlling thread. +# Process-wide claims on shared thread pools, keyed by library. A controller opts in by +# inheriting _SharedPoolControllerMixin. Separate copies of a library, for example +# NumPy's and SciPy's OpenBLAS, have separate pools, so each copy has its own record and +# its own controlling thread. _SHARED_POOL_STATES: dict[str, _SharedPoolState] = {} _SHARED_POOL_STATES_LOCK = Lock() -# Bumped per controller class by a forced clear. Other threads then drop that -# class's counts and keep every other library. +# Bumped per controller class by a forced clear. Other threads then drop that class's +# counts and keep every other library. _shared_pool_generations: dict[type, int] = {} @@ -526,10 +518,9 @@ def _release_shared_pool_claim(key: str): def _clear_shared_pool_claims(owner: type, force: bool = False): """Drop claims on shared pools owned by one controller class. - ``owner`` is that class, for example ``OpenBLASController``. Pools of - other classes are left untouched, including when ``force`` is true. A - forced clear also drops that class's counts on other threads the next - time they sync. + Pools of other classes are left untouched, including when ``force`` is + true. A forced clear also drops that class's counts on other threads the + next time they sync. """ this = current_thread() cleared_this_thread = False @@ -549,6 +540,74 @@ def _clear_shared_pool_claims(owner: type, force: bool = False): SHARED_POOL_CONTROL_OWNERSHIP.reset(owner) +class _SharedPoolControllerMixin: + """Controlling-thread claims for a library with a process-wide thread pool. + + Inherit this next to ``LibController`` and implement ``_real_set_num_threads``. + ``set_num_threads`` should call ``_set_shared_pool_num_threads`` when the + loaded library shares one pool across Python threads. Override + ``_has_shared_pool`` when that depends on how the library was built. + ``clear_controlling_thread`` drops claims of this controller class only. + """ + + def _has_shared_pool(self): + """Whether this loaded library currently uses a process-wide pool.""" + return True + + def _shared_pool_key(self) -> str: + """Identify this library's shared pool, one per loaded copy.""" + return self.filepath + + def _shared_pool_state(self) -> _SharedPoolState: + """Return the process-wide claim record for this library's shared pool.""" + return _get_shared_pool_state(self._shared_pool_key(), type(self)) + + @classmethod + def clear_controlling_thread(cls, force: bool = False): + """Reset which thread can limit this controller's shared pools. + + Drops the claim on every library of this controller class that the + current thread controls, and resets this thread's ownership counts for + that class so a later context manager can release a claim of its own. + Shared pools of other controller classes are left in place. It's up to + the caller to ensure the current controlling thread is done. + + By default, only the controlling thread has an effect. Any thread can + clear every library of this class by setting ``force=True``. + """ + _clear_shared_pool_claims(cls, force) + + def _set_shared_pool_num_threads(self, num_threads): + """Set the pool size if this thread holds or can take the claim. + + Another live thread's call does nothing. A finished controlling thread + no longer holds the claim. A direct call, while this thread's ownership + count is still 0, is sticky: a later context manager must not drop it. + A call made while a limiter scope is active is released with that scope. + """ + SHARED_POOL_CONTROL_OWNERSHIP.sync() + this_thread = current_thread() + state = self._shared_pool_state() + controlling_thread = state.controlling_thread + if controlling_thread is None or not controlling_thread.is_alive(): + # Read the count before taking the pool lock. Syncing ownership takes the + # states lock, and releasing a claim takes that lock before the pool lock. + count = SHARED_POOL_CONTROL_OWNERSHIP.count_for(self._shared_pool_key()) + with state.lock: + holder = state.controlling_thread + if holder is None or not holder.is_alive(): + state.controlling_thread = this_thread + state.sticky = count == 0 + # Retry now that the claim attempt is done. If this thread won, the + # next call reaches _real_set_num_threads. If another live thread won, + # the next call does nothing. + return self._set_shared_pool_num_threads(num_threads) + elif controlling_thread is this_thread: + return self._real_set_num_threads(num_threads) + else: + return + + @dataclass class _SharedPoolState: """Who may change one shared thread pool.""" @@ -557,12 +616,12 @@ class _SharedPoolState: # Set when control is taken outside a limiter scope (direct set_num_threads). sticky: bool = False lock: Lock = field(default_factory=Lock) - # Controller class that owns this pool, so one class can clear its claims - # without releasing another library. + # Controller class that owns this pool, so one class can clear its claims without + # releasing another library. owner: type | None = None -class OpenBLASController(LibController): +class OpenBLASController(_SharedPoolControllerMixin, LibController): """Controller class for OpenBLAS. When OpenBLAS uses the OpenMP threading layer, each thread that calls into @@ -639,56 +698,11 @@ def get_num_threads(self): def _has_shared_pool(self): return self.threading_layer != "openmp" - @classmethod - def clear_controlling_thread(cls, force: bool = False): - """Reset which thread can set limits (on non-OpenMP threading layer). - - Drops the claim on every OpenBLAS library this thread controls, and - resets this thread's OpenBLAS ownership counts so a later context - manager can release a claim of its own. Shared pools of other libraries - are left in place. It's up to the caller to ensure the current - controlling thread is done. - - By default, only the controlling thread has an effect. Any thread can - clear every OpenBLAS library by setting ``force=True``. - """ - _clear_shared_pool_claims(cls, force) - def set_num_threads(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) - - # A shared process-wide thread pool, so don't want different Python - # threads fighting over limits: - SHARED_POOL_CONTROL_OWNERSHIP.sync() - this_thread = current_thread() - state = self._shared_pool_state() - controlling_thread = state.controlling_thread - if controlling_thread is None or not controlling_thread.is_alive(): - # No live thread controls this OpenBLAS, so try to take control. - # A direct call (ownership count still 0) is sticky: a later - # context manager must not drop it. A call made while a limiter - # scope is active is released with that scope. A finished thread - # no longer holds the claim. - # Read the count before taking the pool lock. Syncing ownership - # takes the states lock, and releasing a claim takes that lock - # before the pool lock. - count = SHARED_POOL_CONTROL_OWNERSHIP.count_for(self._shared_pool_key()) - with state.lock: - holder = state.controlling_thread - if holder is None or not holder.is_alive(): - state.controlling_thread = this_thread - state.sticky = count == 0 - # Controlling thread should now be set, to either this thread - # or another, so try again: - return self.set_num_threads(num_threads) - elif controlling_thread is this_thread: - # This thread controls this OpenBLAS: - return self._real_set_num_threads(num_threads) - else: - # Some other thread controls this OpenBLAS: - return + return self._set_shared_pool_num_threads(num_threads) def _real_set_num_threads(self, num_threads): # The OpenBLAS limiting API is process-wide, and we want current thread @@ -1238,9 +1252,9 @@ def __init__(self, controller, *, limits=None, user_api=None): self._limits, self._user_api = self._check_params(limits, user_api) self._original_info = self._controller.info() self._shared_pool_claim_keys: list[str] = [] - # A plain call never reaches __exit__, so this increment is what makes - # the claim stick until clear_controlling_thread(). A `with` block - # decrements the same increment on the way out. + # A plain call never reaches __exit__, so this increment is what makes the claim + # stick until clear_controlling_thread(). A `with` block decrements the same + # increment on the way out. self._apply_limits_with_shared_pool_claims() def __enter__(self): @@ -1458,9 +1472,9 @@ def __init__(self, controller, *, limits=None, user_api=None): self._shared_pool_claim_keys = [] def __enter__(self): - # Set the limits here and not in __init__ so they apply when the - # decorated function is called, not when the decorator is created. - # Increment before setting so this scope's claim is not sticky. + # Set the limits here and not in __init__ so they apply when the decorated + # function is called, not when the decorator is created. Increment before + # setting so this scope's claim is not sticky. self._original_info = self._controller.info() self._apply_limits_with_shared_pool_claims() return self From 7e9832a59b57a69ef01d0576ad9ecf967caf4767 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 14:29:15 +0100 Subject: [PATCH 04/21] FIX Let a scoped limit replace a finished controlling thread without sticking. Also add SHARED_POOL_CONTROL_OWNERSHIP.increment() and decrement() without a key to act on every shared pool this thread can control, and rename the ownership helpers so they are not OpenBLAS-specific. Co-authored-by: Cursor --- CHANGES.md | 4 ++ tests/test_threadpoolctl.py | 120 ++++++++++++++++++++++++++++++++++-- threadpoolctl.py | 68 +++++++++++++++++--- 3 files changed, 178 insertions(+), 14 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 72fd9bd2..0f4f2cb9 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -16,8 +16,12 @@ later scope does not drop a claim it did not take, for example an OpenMP-only limit or ``limits=None``. A thread that has finished no longer holds the claim; the next ``set_num_threads()`` can take it. + ``SHARED_POOL_CONTROL_OWNERSHIP.increment()`` and ``decrement()`` with no key + open and close that scope for every shared pool the current thread can + control. https://github.com/joblib/threadpoolctl/pull/275 https://github.com/joblib/threadpoolctl/pull/289 + https://github.com/joblib/threadpoolctl/pull/XXX - Add a new API to get a cached `ThreadpoolController` instance (`get_cached_controller()`), as an optimization for cases where limits are set diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 5d50cd8c..de23707e 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -912,6 +912,73 @@ def test_openblas_controlling_thread_manual_increment_decrement_api(request): assert lib._shared_pool_state().controlling_thread is None +class _TestSharedPoolController(threadpoolctl._SharedPoolControllerMixin): + """Stand-in shared pool addressed by filepath.""" + + def __init__(self, key): + self.filepath = key + self.prefix = key + self.user_api = "blas" + self.applied = None + + def set_num_threads(self, num_threads): + self._set_shared_pool_num_threads(num_threads) + + def _real_set_num_threads(self, num_threads): + self.applied = num_threads + + +def test_bare_increment_decrement_covers_every_controllable_pool(): + # No key: every loaded shared pool this thread can control, and not a pool + # another live thread already holds. + free_a = _TestSharedPoolController("threadpoolctl-test-bare-a") + free_b = _TestSharedPoolController("threadpoolctl-test-bare-b") + held = _TestSharedPoolController("threadpoolctl-test-bare-held") + started, finish = Lock(), Lock() + started.acquire() + finish.acquire() + + def hold(): + held._shared_pool_state().controlling_thread = current_thread() + started.release() + finish.acquire() + + holder = Thread(target=hold) + controller = get_cached_controller() + controller.lib_controllers.extend([free_a, free_b, held]) + try: + holder.start() + started.acquire() + SHARED_POOL_CONTROL_OWNERSHIP.increment() + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(free_a.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(free_b.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(held.filepath) == 0 + + free_a.set_num_threads(1) + free_b.set_num_threads(1) + assert free_a._shared_pool_state().controlling_thread is current_thread() + assert free_b._shared_pool_state().controlling_thread is current_thread() + assert free_a._shared_pool_state().sticky is False + assert free_b._shared_pool_state().sticky is False + + SHARED_POOL_CONTROL_OWNERSHIP.increment() + SHARED_POOL_CONTROL_OWNERSHIP.decrement() + assert free_a._shared_pool_state().controlling_thread is current_thread() + SHARED_POOL_CONTROL_OWNERSHIP.decrement() + assert free_a._shared_pool_state().controlling_thread is None + assert free_b._shared_pool_state().controlling_thread is None + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(free_a.filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(free_b.filepath) == 0 + finally: + finish.release() + if holder.ident is not None: + holder.join() + for lib in (free_a, free_b, held): + controller.lib_controllers.remove(lib) + threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) + SHARED_POOL_CONTROL_OWNERSHIP.reset() + + def _require_resizable_openblas(): libs = _non_openmp_openblas() lib = libs[0] @@ -1038,8 +1105,8 @@ def test_openblas_scoped_exit_releases_only_that_library(): # Bookkeeping for the ownership leak, including where no pthreads OpenBLAS # is installed. A scope ending on one library must not drop a sticky claim # or a still-held claim on another library. - key_a = "\0threadpoolctl-test-openblas-a" - key_b = "\0threadpoolctl-test-openblas-b" + key_a = "threadpoolctl-test-openblas-a" + key_b = "threadpoolctl-test-openblas-b" state_a = threadpoolctl._SharedPoolState() state_b = threadpoolctl._SharedPoolState() threadpoolctl._SHARED_POOL_STATES[key_a] = state_a @@ -1076,6 +1143,51 @@ def test_openblas_scoped_exit_releases_only_that_library(): SHARED_POOL_CONTROL_OWNERSHIP.reset() +class _DeadHolderLib(threadpoolctl._SharedPoolControllerMixin): + """Stand-in library whose pool is already held by a finished thread.""" + + def __init__(self, key): + self.filepath = key + self.prefix = "dead-holder" + self.user_api = "blas" + self.applied = None + + def set_num_threads(self, num_threads): + self._set_shared_pool_num_threads(num_threads) + + def _real_set_num_threads(self, num_threads): + self.applied = num_threads + + +def test_scoped_limit_replacing_dead_holder_is_not_sticky(): + # A with-block that takes over from a finished controlling thread must + # release that claim on the way out. + key = "threadpoolctl-test-dead-holder" + dead = Thread(target=lambda: None) + dead.start() + dead.join() + lib = _DeadHolderLib(key) + state = lib._shared_pool_state() + state.controlling_thread = dead + state.sticky = True + limiter = object.__new__(threadpoolctl._ThreadpoolLimiter) + limiter._controller = type("C", (), {"lib_controllers": [lib]})() + limiter._limits = {key: 1} + limiter._shared_pool_claim_keys = [] + try: + limiter._apply_limits_with_shared_pool_claims() + assert lib.applied == 1 + assert state.controlling_thread is current_thread() + assert state.sticky is False + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key) == 1 + limiter._release_shared_pool_claims() + assert state.controlling_thread is None + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key) == 0 + finally: + threadpoolctl._SHARED_POOL_STATES.pop(key, None) + SHARED_POOL_CONTROL_OWNERSHIP.reset() + + class _OtherSharedPool: """Stand-in for a future library with its own shared pool.""" @@ -1083,8 +1195,8 @@ class _OtherSharedPool: def test_clear_controlling_thread_leaves_other_libraries(): # OpenBLASController.clear_controlling_thread, including force=True, must # not drop a claim or a count owned by another controller class. - key_openblas = "\0threadpoolctl-test-openblas-clear" - key_other = "\0threadpoolctl-test-other-clear" + key_openblas = "threadpoolctl-test-openblas-clear" + key_other = "threadpoolctl-test-other-clear" state_openblas = threadpoolctl._SharedPoolState(owner=OpenBLASController) state_other = threadpoolctl._SharedPoolState(owner=_OtherSharedPool) threadpoolctl._SHARED_POOL_STATES[key_openblas] = state_openblas diff --git a/threadpoolctl.py b/threadpoolctl.py index e88b5b1c..d17bd3e8 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -59,6 +59,7 @@ __version__ = "3.8.0.dev0" __all__ = [ + "SHARED_POOL_CONTROL_OWNERSHIP", "LibController", "ThreadpoolController", "get_cached_controller", @@ -499,6 +500,11 @@ def _snapshot_shared_pool_states() -> list[_SharedPoolState]: return list(_SHARED_POOL_STATES.values()) +def _other_live_thread_holds(holder, this) -> bool: + """Whether ``holder`` is a different thread that still owns the pool.""" + return holder is not None and holder is not this and holder.is_alive() + + def _release_shared_pool_claim(key: str): """Drop this thread's non-sticky claim on one shared pool. @@ -769,11 +775,14 @@ class _SharedPoolControlOwnership(thread_local): A count is kept separately for each shared pool, keyed by library filepath. Hitting zero for a pool releases that pool's non-sticky controlling thread. A plain ``threadpool_limits(...)`` call increments and never decrements, so - it holds its claim until ``OpenBLASController.clear_controlling_thread()``. - - The counts live on the thread-local instance. A class-level dict would be - shared by every thread. A forced clear drops counts for that controller - class only. + it holds its claim until ``clear_controlling_thread()``. + + ``increment()`` and ``decrement()`` with no key apply to every loaded + shared pool this thread can control. Call ``increment()`` before + ``set_num_threads`` and ``decrement()`` when that scope ends to manage the + claim without a context manager. The counts live on the thread-local + instance. A class-level dict would be shared by every thread. A forced + clear drops counts for that controller class only. """ def _counts(self) -> dict[str, int]: @@ -837,19 +846,56 @@ def reset(self, owner: type | None = None): counts.pop(key, None) self._generations()[owner] = generation + def _claimable_shared_pool_keys(self) -> list[str]: + """Loaded shared pools this thread can control, one key per copy.""" + this = current_thread() + keys = [] + seen = set() + for lib in get_cached_controller().lib_controllers: + if not lib._has_shared_pool(): + continue + key = lib._shared_pool_key() + if key in seen: + continue + seen.add(key) + holder = lib._shared_pool_state().controlling_thread + if _other_live_thread_holds(holder, this): + continue + keys.append(key) + return keys + def count_for(self, key: str) -> int: """Active scopes for one shared pool on this thread.""" self.sync() return self._counts().get(key, 0) - def increment(self, key: str): - """Increment ownership of one shared pool.""" + def increment(self, key: str | None = None): + """Increment ownership of one shared pool, or of every controllable one. + + With no ``key``, each loaded library with a process-wide pool that + this thread can control is incremented once. Call this before + ``set_num_threads`` so the claim stays scoped. + """ + if key is None: + for pool_key in self._claimable_shared_pool_keys(): + self.increment(pool_key) + return self.sync() counts = self._counts() counts[key] = counts.get(key, 0) + 1 - def decrement(self, key: str): - """Decrement ownership of one shared pool, releasing its scoped claim at zero.""" + def decrement(self, key: str | None = None): + """Decrement ownership, releasing a non-sticky claim when it hits zero. + + With no ``key``, every pool this thread currently counts is + decremented once. That closes a matching ``increment()`` and also any + pool incremented on its own. + """ + if key is None: + self.sync() + for pool_key in list(self._counts()): + self.decrement(pool_key) + return self.sync() counts = self._counts() current = counts.get(key, 0) @@ -1415,7 +1461,9 @@ def _acquire_shared_pool_claims(self): continue seen.add(key) holder = lib._shared_pool_state().controlling_thread - if holder is not None and holder is not this: + # A finished holder does not block this scope. Counting first keeps + # the replacement claim non-sticky, so the scope can release it. + if _other_live_thread_holds(holder, this): continue SHARED_POOL_CONTROL_OWNERSHIP.increment(key) self._shared_pool_claim_keys.append(key) From a9ab4771253c1a1f5d74193b0f6b91763b5e1c44 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 14:42:57 +0100 Subject: [PATCH 05/21] MNT Share claim walk and rename owner to controller_class. Put _SharedPoolState above the registry and document the lock order. Use one helper to find claimable shared pools for both the bare increment() and the limiter. Drop a duplicate test stand-in. Co-authored-by: Cursor --- tests/test_threadpoolctl.py | 22 +---- threadpoolctl.py | 159 ++++++++++++++++++------------------ 2 files changed, 83 insertions(+), 98 deletions(-) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index de23707e..9097709e 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -1143,22 +1143,6 @@ def test_openblas_scoped_exit_releases_only_that_library(): SHARED_POOL_CONTROL_OWNERSHIP.reset() -class _DeadHolderLib(threadpoolctl._SharedPoolControllerMixin): - """Stand-in library whose pool is already held by a finished thread.""" - - def __init__(self, key): - self.filepath = key - self.prefix = "dead-holder" - self.user_api = "blas" - self.applied = None - - def set_num_threads(self, num_threads): - self._set_shared_pool_num_threads(num_threads) - - def _real_set_num_threads(self, num_threads): - self.applied = num_threads - - def test_scoped_limit_replacing_dead_holder_is_not_sticky(): # A with-block that takes over from a finished controlling thread must # release that claim on the way out. @@ -1166,7 +1150,7 @@ def test_scoped_limit_replacing_dead_holder_is_not_sticky(): dead = Thread(target=lambda: None) dead.start() dead.join() - lib = _DeadHolderLib(key) + lib = _TestSharedPoolController(key) state = lib._shared_pool_state() state.controlling_thread = dead state.sticky = True @@ -1197,8 +1181,8 @@ def test_clear_controlling_thread_leaves_other_libraries(): # not drop a claim or a count owned by another controller class. key_openblas = "threadpoolctl-test-openblas-clear" key_other = "threadpoolctl-test-other-clear" - state_openblas = threadpoolctl._SharedPoolState(owner=OpenBLASController) - state_other = threadpoolctl._SharedPoolState(owner=_OtherSharedPool) + state_openblas = threadpoolctl._SharedPoolState(controller_class=OpenBLASController) + state_other = threadpoolctl._SharedPoolState(controller_class=_OtherSharedPool) threadpoolctl._SHARED_POOL_STATES[key_openblas] = state_openblas threadpoolctl._SHARED_POOL_STATES[key_other] = state_other try: diff --git a/threadpoolctl.py b/threadpoolctl.py index d17bd3e8..fc070c64 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -473,25 +473,41 @@ def wrapper(self): _CDLL_CACHE = _CDLLCache() +@dataclass +class _SharedPoolState: + """Who may change one shared thread pool.""" + + controlling_thread: Thread | None = None + # Set when control is taken outside a limiter scope (direct set_num_threads). + sticky: bool = False + lock: Lock = field(default_factory=Lock) + # Controller class for this pool, so a clear of one class leaves the others. + controller_class: type | None = None + + # Process-wide claims on shared thread pools, keyed by library. A controller opts in by # inheriting _SharedPoolControllerMixin. Separate copies of a library, for example # NumPy's and SciPy's OpenBLAS, have separate pools, so each copy has its own record and # its own controlling thread. +# +# Lock order: take _SHARED_POOL_STATES_LOCK before state.lock. Do not call sync() or +# count_for() while holding state.lock. Both take the states lock, and releasing a +# claim takes the states lock before the pool lock. _SHARED_POOL_STATES: dict[str, _SharedPoolState] = {} _SHARED_POOL_STATES_LOCK = Lock() -# Bumped per controller class by a forced clear. Other threads then drop that class's -# counts and keep every other library. +# Bumped per controller class by a forced clear. A thread drops a class's counts when +# its stored generation for that class does not match. _shared_pool_generations: dict[type, int] = {} -def _get_shared_pool_state(key: str, owner: type) -> _SharedPoolState: +def _get_shared_pool_state(key: str, controller_class: type) -> _SharedPoolState: with _SHARED_POOL_STATES_LOCK: state = _SHARED_POOL_STATES.get(key) if state is None: - state = _SharedPoolState(owner=owner) + state = _SharedPoolState(controller_class=controller_class) _SHARED_POOL_STATES[key] = state - elif state.owner is None: - state.owner = owner + elif state.controller_class is None: + state.controller_class = controller_class return state @@ -505,6 +521,26 @@ def _other_live_thread_holds(holder, this) -> bool: return holder is not None and holder is not this and holder.is_alive() +def _iter_claimable_shared_pools(controllers): + """Yield ``(key, lib)`` for shared pools in ``controllers`` this thread can claim. + + One entry per pool. A pool another live thread already holds is skipped. + """ + this = current_thread() + seen = set() + for lib in controllers: + if not lib._has_shared_pool(): + continue + key = lib._shared_pool_key() + if key in seen: + continue + seen.add(key) + holder = lib._shared_pool_state().controlling_thread + if _other_live_thread_holds(holder, this): + continue + yield key, lib + + def _release_shared_pool_claim(key: str): """Drop this thread's non-sticky claim on one shared pool. @@ -521,8 +557,8 @@ def _release_shared_pool_claim(key: str): state.controlling_thread = None -def _clear_shared_pool_claims(owner: type, force: bool = False): - """Drop claims on shared pools owned by one controller class. +def _clear_shared_pool_claims(controller_class: type, force: bool = False): + """Drop claims on shared pools of one controller class. Pools of other classes are left untouched, including when ``force`` is true. A forced clear also drops that class's counts on other threads the @@ -531,7 +567,7 @@ def _clear_shared_pool_claims(owner: type, force: bool = False): this = current_thread() cleared_this_thread = False for state in _snapshot_shared_pool_states(): - if state.owner is not owner: + if state.controller_class is not controller_class: continue with state.lock: if force or state.controlling_thread is this: @@ -541,9 +577,10 @@ def _clear_shared_pool_claims(owner: type, force: bool = False): state.sticky = False if force: with _SHARED_POOL_STATES_LOCK: - _shared_pool_generations[owner] = _shared_pool_generations.get(owner, 0) + 1 + generation = _shared_pool_generations.get(controller_class, 0) + 1 + _shared_pool_generations[controller_class] = generation if force or cleared_this_thread: - SHARED_POOL_CONTROL_OWNERSHIP.reset(owner) + SHARED_POOL_CONTROL_OWNERSHIP.reset(controller_class) class _SharedPoolControllerMixin: @@ -596,8 +633,7 @@ def _set_shared_pool_num_threads(self, num_threads): state = self._shared_pool_state() controlling_thread = state.controlling_thread if controlling_thread is None or not controlling_thread.is_alive(): - # Read the count before taking the pool lock. Syncing ownership takes the - # states lock, and releasing a claim takes that lock before the pool lock. + # Read the count before state.lock. count_for() takes the states lock. count = SHARED_POOL_CONTROL_OWNERSHIP.count_for(self._shared_pool_key()) with state.lock: holder = state.controlling_thread @@ -614,19 +650,6 @@ def _set_shared_pool_num_threads(self, num_threads): return -@dataclass -class _SharedPoolState: - """Who may change one shared thread pool.""" - - controlling_thread: Thread | None = None - # Set when control is taken outside a limiter scope (direct set_num_threads). - sticky: bool = False - lock: Lock = field(default_factory=Lock) - # Controller class that owns this pool, so one class can clear its claims without - # releasing another library. - owner: type | None = None - - class OpenBLASController(_SharedPoolControllerMixin, LibController): """Controller class for OpenBLAS. @@ -804,65 +827,52 @@ def sync(self): with _SHARED_POOL_STATES_LOCK: generations = dict(_shared_pool_generations) owned_keys = [ - (key, state.owner) + (key, state.controller_class) for key, state in _SHARED_POOL_STATES.items() - if state.owner is not None + if state.controller_class is not None ] local = self._generations() stale = { - owner - for owner, generation in generations.items() - if local.get(owner, 0) != generation + controller_class + for controller_class, generation in generations.items() + if local.get(controller_class, 0) != generation } if not stale: return counts = self._counts() - for key, owner in owned_keys: - if owner in stale: + for key, controller_class in owned_keys: + if controller_class in stale: counts.pop(key, None) - for owner in stale: - local[owner] = generations[owner] + for controller_class in stale: + local[controller_class] = generations[controller_class] - def reset(self, owner: type | None = None): + def reset(self, controller_class: type | None = None): """Forget scopes on this thread. - With ``owner``, only that controller class. Without, every library. + With ``controller_class``, only that class. Without, every library. """ - if owner is None: + if controller_class is None: with _SHARED_POOL_STATES_LOCK: generations = dict(_shared_pool_generations) self.counts = {} self.generations = generations return with _SHARED_POOL_STATES_LOCK: - generation = _shared_pool_generations.get(owner, 0) + generation = _shared_pool_generations.get(controller_class, 0) keys = [ key for key, state in _SHARED_POOL_STATES.items() - if state.owner is owner + if state.controller_class is controller_class ] counts = self._counts() for key in keys: counts.pop(key, None) - self._generations()[owner] = generation + self._generations()[controller_class] = generation def _claimable_shared_pool_keys(self) -> list[str]: - """Loaded shared pools this thread can control, one key per copy.""" - this = current_thread() - keys = [] - seen = set() - for lib in get_cached_controller().lib_controllers: - if not lib._has_shared_pool(): - continue - key = lib._shared_pool_key() - if key in seen: - continue - seen.add(key) - holder = lib._shared_pool_state().controlling_thread - if _other_live_thread_holds(holder, this): - continue - keys.append(key) - return keys + """Pools this thread can control, from the cached controller.""" + libs = get_cached_controller().lib_controllers + return [key for key, _lib in _iter_claimable_shared_pools(libs)] def count_for(self, key: str) -> int: """Active scopes for one shared pool on this thread.""" @@ -873,7 +883,8 @@ def increment(self, key: str | None = None): """Increment ownership of one shared pool, or of every controllable one. With no ``key``, each loaded library with a process-wide pool that - this thread can control is incremented once. Call this before + this thread can control is incremented once. Those libraries are the + ones known to ``get_cached_controller()``. Call this before ``set_num_threads`` so the claim stays scoped. """ if key is None: @@ -1439,32 +1450,22 @@ def _num_threads_for(self, lib_controller): return None return num_threads - def _iter_shared_pool_targets(self): - """Libraries with a shared pool that this limiter would actually set.""" - for lib in self._controller.lib_controllers: - if lib._has_shared_pool() and self._num_threads_for(lib) is not None: - yield lib - def _acquire_shared_pool_claims(self): """Count a scope per library before setting, so the claim is not sticky. - Only libraries this thread can control are counted. OpenMP-only limits - and ``limits=None`` do not take a claim. Libraries are counted - separately so a scope that limits one copy does not pin another. + Only libraries this limiter would set, and that this thread can + control, are counted. OpenMP-only limits and ``limits=None`` do not + take a claim. Libraries are counted separately so a scope that limits + one copy does not pin another. A finished holder does not block this + scope: counting first keeps the replacement claim non-sticky. """ self._shared_pool_claim_keys = [] - this = current_thread() - seen = set() - for lib in self._iter_shared_pool_targets(): - key = lib._shared_pool_key() - if key in seen: - continue - seen.add(key) - holder = lib._shared_pool_state().controlling_thread - # A finished holder does not block this scope. Counting first keeps - # the replacement claim non-sticky, so the scope can release it. - if _other_live_thread_holds(holder, this): - continue + targets = [ + lib + for lib in self._controller.lib_controllers + if self._num_threads_for(lib) is not None + ] + for key, _lib in _iter_claimable_shared_pools(targets): SHARED_POOL_CONTROL_OWNERSHIP.increment(key) self._shared_pool_claim_keys.append(key) From c2a1c16cc4ab3aee3d0a9ebac1949f6ec09d5cb8 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 15:07:30 +0100 Subject: [PATCH 06/21] MNT Key shared-pool state by filepath and take an optional controller. Drop _shared_pool_key(): the library filepath already identifies the pool. SHARED_POOL_CONTROL_OWNERSHIP.increment() and decrement() take an optional controller and default to get_cached_controller(). Rename variables and tests to match what they now describe. Co-authored-by: Cursor --- CHANGES.md | 6 +- tests/test_threadpoolctl.py | 130 ++++++++++++++++++---------- threadpoolctl.py | 166 ++++++++++++++++++------------------ 3 files changed, 171 insertions(+), 131 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 0f4f2cb9..1deccfae 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -16,9 +16,9 @@ later scope does not drop a claim it did not take, for example an OpenMP-only limit or ``limits=None``. A thread that has finished no longer holds the claim; the next ``set_num_threads()`` can take it. - ``SHARED_POOL_CONTROL_OWNERSHIP.increment()`` and ``decrement()`` with no key - open and close that scope for every shared pool the current thread can - control. + ``SHARED_POOL_CONTROL_OWNERSHIP.increment(controller)`` and + ``decrement(controller)`` open and close that scope for the shared pools of + ``controller``, which defaults to ``get_cached_controller()``. https://github.com/joblib/threadpoolctl/pull/275 https://github.com/joblib/threadpoolctl/pull/289 https://github.com/joblib/threadpoolctl/pull/XXX diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 9097709e..652a43a7 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -889,10 +889,9 @@ def outer(): assert _openblas_controlling_thread() is None -def test_openblas_controlling_thread_manual_increment_decrement_api(request): - # The manual counter has to be incremented before set_num_threads, one - # count per library. That makes the claim scoped, so the last decrement - # releases it. +def test_openblas_controlling_thread_per_pool_counter(request): + # The per-pool counter has to be incremented before set_num_threads. That + # makes the claim scoped, so the last decrement releases it. libs = _non_openmp_openblas() lib = libs[0] orig = lib.get_num_threads() @@ -900,24 +899,24 @@ def test_openblas_controlling_thread_manual_increment_decrement_api(request): pytest.skip("Test requires a non-OpenMP OpenBLAS with at least 2 threads") request.addfinalizer(_restore_openblas_thread_counts(libs)) - key = lib._shared_pool_key() - SHARED_POOL_CONTROL_OWNERSHIP.increment(key) - SHARED_POOL_CONTROL_OWNERSHIP.increment(key) + filepath = lib.filepath + SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath) + SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath) lib.set_num_threads(1) assert lib._shared_pool_state().controlling_thread is current_thread() assert not lib._shared_pool_state().sticky - SHARED_POOL_CONTROL_OWNERSHIP.decrement(key) + SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath) assert lib._shared_pool_state().controlling_thread is current_thread() - SHARED_POOL_CONTROL_OWNERSHIP.decrement(key) + SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath) assert lib._shared_pool_state().controlling_thread is None class _TestSharedPoolController(threadpoolctl._SharedPoolControllerMixin): """Stand-in shared pool addressed by filepath.""" - def __init__(self, key): - self.filepath = key - self.prefix = key + def __init__(self, filepath): + self.filepath = filepath + self.prefix = filepath self.user_api = "blas" self.applied = None @@ -929,8 +928,8 @@ def _real_set_num_threads(self, num_threads): def test_bare_increment_decrement_covers_every_controllable_pool(): - # No key: every loaded shared pool this thread can control, and not a pool - # another live thread already holds. + # Every loaded shared pool this thread can control, and not a pool another + # live thread already holds. free_a = _TestSharedPoolController("threadpoolctl-test-bare-a") free_b = _TestSharedPoolController("threadpoolctl-test-bare-b") held = _TestSharedPoolController("threadpoolctl-test-bare-held") @@ -958,6 +957,9 @@ def hold(): free_b.set_num_threads(1) assert free_a._shared_pool_state().controlling_thread is current_thread() assert free_b._shared_pool_state().controlling_thread is current_thread() + # increment() ran before set_num_threads, so each claim is taken with a + # count above zero and is not sticky. A sticky claim would survive the + # decrements below; these are released when the count hits zero. assert free_a._shared_pool_state().sticky is False assert free_b._shared_pool_state().sticky is False @@ -979,6 +981,40 @@ def hold(): SHARED_POOL_CONTROL_OWNERSHIP.reset() +def test_increment_decrement_uses_the_given_controller(): + # Passing a controller selects its pools. Omitting it uses the cached + # controller. + chosen = _TestSharedPoolController("threadpoolctl-test-controller-chosen") + other = _TestSharedPoolController("threadpoolctl-test-controller-other") + controller = type("C", (), {"lib_controllers": [chosen]})() + cached = get_cached_controller() + cached.lib_controllers.append(other) + try: + SHARED_POOL_CONTROL_OWNERSHIP.increment(controller) + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(chosen.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(other.filepath) == 0 + + chosen.set_num_threads(1) + assert chosen._shared_pool_state().sticky is False + + SHARED_POOL_CONTROL_OWNERSHIP.increment() + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(chosen.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(other.filepath) == 1 + + SHARED_POOL_CONTROL_OWNERSHIP.decrement(controller) + assert chosen._shared_pool_state().controlling_thread is None + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(chosen.filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(other.filepath) == 1 + + SHARED_POOL_CONTROL_OWNERSHIP.decrement() + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(other.filepath) == 0 + finally: + cached.lib_controllers.remove(other) + for lib in (chosen, other): + threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) + SHARED_POOL_CONTROL_OWNERSHIP.reset() + + def _require_resizable_openblas(): libs = _non_openmp_openblas() lib = libs[0] @@ -1105,22 +1141,22 @@ def test_openblas_scoped_exit_releases_only_that_library(): # Bookkeeping for the ownership leak, including where no pthreads OpenBLAS # is installed. A scope ending on one library must not drop a sticky claim # or a still-held claim on another library. - key_a = "threadpoolctl-test-openblas-a" - key_b = "threadpoolctl-test-openblas-b" + filepath_a = "threadpoolctl-test-openblas-a" + filepath_b = "threadpoolctl-test-openblas-b" state_a = threadpoolctl._SharedPoolState() state_b = threadpoolctl._SharedPoolState() - threadpoolctl._SHARED_POOL_STATES[key_a] = state_a - threadpoolctl._SHARED_POOL_STATES[key_b] = state_b + threadpoolctl._SHARED_POOL_STATES[filepath_a] = state_a + threadpoolctl._SHARED_POOL_STATES[filepath_b] = state_b try: state_a.controlling_thread = current_thread() state_a.sticky = True - SHARED_POOL_CONTROL_OWNERSHIP.increment(key_b) + SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_b) state_b.controlling_thread = current_thread() - state_b.sticky = SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_b) == 0 + state_b.sticky = SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_b) == 0 assert state_b.sticky is False - SHARED_POOL_CONTROL_OWNERSHIP.decrement(key_b) + SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath_b) assert state_b.controlling_thread is None assert state_a.controlling_thread is current_thread() assert state_a.sticky is True @@ -1128,47 +1164,47 @@ def test_openblas_scoped_exit_releases_only_that_library(): # A plain call is non-sticky and is held by a leftover count. Ending a # scope on B must not release A. state_a.sticky = False - SHARED_POOL_CONTROL_OWNERSHIP.increment(key_a) - SHARED_POOL_CONTROL_OWNERSHIP.increment(key_b) + SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_a) + SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_b) state_b.controlling_thread = current_thread() state_b.sticky = False - SHARED_POOL_CONTROL_OWNERSHIP.decrement(key_b) + SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath_b) assert state_b.controlling_thread is None assert state_a.controlling_thread is current_thread() assert state_a.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_a) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_a) == 1 finally: - threadpoolctl._SHARED_POOL_STATES.pop(key_a, None) - threadpoolctl._SHARED_POOL_STATES.pop(key_b, None) + threadpoolctl._SHARED_POOL_STATES.pop(filepath_a, None) + threadpoolctl._SHARED_POOL_STATES.pop(filepath_b, None) SHARED_POOL_CONTROL_OWNERSHIP.reset() def test_scoped_limit_replacing_dead_holder_is_not_sticky(): # A with-block that takes over from a finished controlling thread must # release that claim on the way out. - key = "threadpoolctl-test-dead-holder" + filepath = "threadpoolctl-test-dead-holder" dead = Thread(target=lambda: None) dead.start() dead.join() - lib = _TestSharedPoolController(key) + lib = _TestSharedPoolController(filepath) state = lib._shared_pool_state() state.controlling_thread = dead state.sticky = True limiter = object.__new__(threadpoolctl._ThreadpoolLimiter) limiter._controller = type("C", (), {"lib_controllers": [lib]})() - limiter._limits = {key: 1} - limiter._shared_pool_claim_keys = [] + limiter._limits = {filepath: 1} + limiter._shared_pool_filepaths = [] try: limiter._apply_limits_with_shared_pool_claims() assert lib.applied == 1 assert state.controlling_thread is current_thread() assert state.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath) == 1 limiter._release_shared_pool_claims() assert state.controlling_thread is None - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath) == 0 finally: - threadpoolctl._SHARED_POOL_STATES.pop(key, None) + threadpoolctl._SHARED_POOL_STATES.pop(filepath, None) SHARED_POOL_CONTROL_OWNERSHIP.reset() @@ -1179,31 +1215,31 @@ class _OtherSharedPool: def test_clear_controlling_thread_leaves_other_libraries(): # OpenBLASController.clear_controlling_thread, including force=True, must # not drop a claim or a count owned by another controller class. - key_openblas = "threadpoolctl-test-openblas-clear" - key_other = "threadpoolctl-test-other-clear" + filepath_openblas = "threadpoolctl-test-openblas-clear" + filepath_other = "threadpoolctl-test-other-clear" state_openblas = threadpoolctl._SharedPoolState(controller_class=OpenBLASController) state_other = threadpoolctl._SharedPoolState(controller_class=_OtherSharedPool) - threadpoolctl._SHARED_POOL_STATES[key_openblas] = state_openblas - threadpoolctl._SHARED_POOL_STATES[key_other] = state_other + threadpoolctl._SHARED_POOL_STATES[filepath_openblas] = state_openblas + threadpoolctl._SHARED_POOL_STATES[filepath_other] = state_other try: state_openblas.controlling_thread = current_thread() state_openblas.sticky = True state_other.controlling_thread = current_thread() state_other.sticky = True - SHARED_POOL_CONTROL_OWNERSHIP.increment(key_openblas) - SHARED_POOL_CONTROL_OWNERSHIP.increment(key_other) + SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_openblas) + SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_other) OpenBLASController.clear_controlling_thread() assert state_openblas.controlling_thread is None assert state_openblas.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_openblas) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_openblas) == 0 assert state_other.controlling_thread is current_thread() assert state_other.sticky is True - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_other) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_other) == 1 state_openblas.controlling_thread = current_thread() state_openblas.sticky = True - SHARED_POOL_CONTROL_OWNERSHIP.increment(key_openblas) + SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_openblas) def force_clear(): OpenBLASController.clear_controlling_thread(force=True) @@ -1213,11 +1249,11 @@ def force_clear(): assert state_openblas.sticky is False assert state_other.controlling_thread is current_thread() assert state_other.sticky is True - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_openblas) == 0 - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(key_other) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_openblas) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_other) == 1 finally: - threadpoolctl._SHARED_POOL_STATES.pop(key_openblas, None) - threadpoolctl._SHARED_POOL_STATES.pop(key_other, None) + threadpoolctl._SHARED_POOL_STATES.pop(filepath_openblas, None) + threadpoolctl._SHARED_POOL_STATES.pop(filepath_other, None) SHARED_POOL_CONTROL_OWNERSHIP.reset() diff --git a/threadpoolctl.py b/threadpoolctl.py index fc070c64..bb55beb7 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -485,7 +485,7 @@ class _SharedPoolState: controller_class: type | None = None -# Process-wide claims on shared thread pools, keyed by library. A controller opts in by +# Process-wide claims on shared thread pools, keyed by filepath. A controller opts in by # inheriting _SharedPoolControllerMixin. Separate copies of a library, for example # NumPy's and SciPy's OpenBLAS, have separate pools, so each copy has its own record and # its own controlling thread. @@ -500,12 +500,12 @@ class _SharedPoolState: _shared_pool_generations: dict[type, int] = {} -def _get_shared_pool_state(key: str, controller_class: type) -> _SharedPoolState: +def _get_shared_pool_state(filepath: str, controller_class: type) -> _SharedPoolState: with _SHARED_POOL_STATES_LOCK: - state = _SHARED_POOL_STATES.get(key) + state = _SHARED_POOL_STATES.get(filepath) if state is None: state = _SharedPoolState(controller_class=controller_class) - _SHARED_POOL_STATES[key] = state + _SHARED_POOL_STATES[filepath] = state elif state.controller_class is None: state.controller_class = controller_class return state @@ -522,33 +522,32 @@ def _other_live_thread_holds(holder, this) -> bool: def _iter_claimable_shared_pools(controllers): - """Yield ``(key, lib)`` for shared pools in ``controllers`` this thread can claim. + """Yield shared-pool libraries in ``controllers`` this thread can claim. - One entry per pool. A pool another live thread already holds is skipped. + One library per filepath. A pool another live thread already holds is skipped. """ this = current_thread() seen = set() for lib in controllers: if not lib._has_shared_pool(): continue - key = lib._shared_pool_key() - if key in seen: + if lib.filepath in seen: continue - seen.add(key) + seen.add(lib.filepath) holder = lib._shared_pool_state().controlling_thread if _other_live_thread_holds(holder, this): continue - yield key, lib + yield lib -def _release_shared_pool_claim(key: str): +def _release_shared_pool_claim(filepath: str): """Drop this thread's non-sticky claim on one shared pool. A sticky claim was taken by a direct ``set_num_threads()`` while no limiter scope was active for that library. Scoped exits must leave it in place. """ with _SHARED_POOL_STATES_LOCK: - state = _SHARED_POOL_STATES.get(key) + state = _SHARED_POOL_STATES.get(filepath) if state is None: return this = current_thread() @@ -597,13 +596,9 @@ def _has_shared_pool(self): """Whether this loaded library currently uses a process-wide pool.""" return True - def _shared_pool_key(self) -> str: - """Identify this library's shared pool, one per loaded copy.""" - return self.filepath - def _shared_pool_state(self) -> _SharedPoolState: """Return the process-wide claim record for this library's shared pool.""" - return _get_shared_pool_state(self._shared_pool_key(), type(self)) + return _get_shared_pool_state(self.filepath, type(self)) @classmethod def clear_controlling_thread(cls, force: bool = False): @@ -634,7 +629,7 @@ def _set_shared_pool_num_threads(self, num_threads): controlling_thread = state.controlling_thread if controlling_thread is None or not controlling_thread.is_alive(): # Read the count before state.lock. count_for() takes the states lock. - count = SHARED_POOL_CONTROL_OWNERSHIP.count_for(self._shared_pool_key()) + count = SHARED_POOL_CONTROL_OWNERSHIP.count_for(self.filepath) with state.lock: holder = state.controlling_thread if holder is None or not holder.is_alive(): @@ -800,12 +795,12 @@ class _SharedPoolControlOwnership(thread_local): A plain ``threadpool_limits(...)`` call increments and never decrements, so it holds its claim until ``clear_controlling_thread()``. - ``increment()`` and ``decrement()`` with no key apply to every loaded - shared pool this thread can control. Call ``increment()`` before - ``set_num_threads`` and ``decrement()`` when that scope ends to manage the - claim without a context manager. The counts live on the thread-local - instance. A class-level dict would be shared by every thread. A forced - clear drops counts for that controller class only. + ``increment(controller)`` counts every shared pool of ``controller`` that + this thread can control. ``decrement(controller)`` counts those pools back + down. ``controller`` defaults to ``get_cached_controller()``. Call + ``increment()`` before ``set_num_threads`` and ``decrement()`` when that + scope ends to manage the claim without a context manager. A forced clear + drops counts for that controller class only. """ def _counts(self) -> dict[str, int]: @@ -826,9 +821,9 @@ def sync(self): """Drop counts for libraries another thread has force-cleared.""" with _SHARED_POOL_STATES_LOCK: generations = dict(_shared_pool_generations) - owned_keys = [ - (key, state.controller_class) - for key, state in _SHARED_POOL_STATES.items() + owned_filepaths = [ + (filepath, state.controller_class) + for filepath, state in _SHARED_POOL_STATES.items() if state.controller_class is not None ] local = self._generations() @@ -840,9 +835,9 @@ def sync(self): if not stale: return counts = self._counts() - for key, controller_class in owned_keys: + for filepath, controller_class in owned_filepaths: if controller_class in stale: - counts.pop(key, None) + counts.pop(filepath, None) for controller_class in stale: local[controller_class] = generations[controller_class] @@ -859,65 +854,74 @@ def reset(self, controller_class: type | None = None): return with _SHARED_POOL_STATES_LOCK: generation = _shared_pool_generations.get(controller_class, 0) - keys = [ - key - for key, state in _SHARED_POOL_STATES.items() + filepaths = [ + filepath + for filepath, state in _SHARED_POOL_STATES.items() if state.controller_class is controller_class ] counts = self._counts() - for key in keys: - counts.pop(key, None) + for filepath in filepaths: + counts.pop(filepath, None) self._generations()[controller_class] = generation - def _claimable_shared_pool_keys(self) -> list[str]: - """Pools this thread can control, from the cached controller.""" - libs = get_cached_controller().lib_controllers - return [key for key, _lib in _iter_claimable_shared_pools(libs)] - - def count_for(self, key: str) -> int: + def count_for(self, filepath: str) -> int: """Active scopes for one shared pool on this thread.""" self.sync() - return self._counts().get(key, 0) + return self._counts().get(filepath, 0) - def increment(self, key: str | None = None): - """Increment ownership of one shared pool, or of every controllable one. + def increment(self, controller: ThreadpoolController | None = None): + """Increment every shared pool of ``controller`` this thread can control. - With no ``key``, each loaded library with a process-wide pool that - this thread can control is incremented once. Those libraries are the - ones known to ``get_cached_controller()``. Call this before - ``set_num_threads`` so the claim stays scoped. + ``controller`` defaults to ``get_cached_controller()``. Call this + before ``set_num_threads`` so the claim stays scoped. """ - if key is None: - for pool_key in self._claimable_shared_pool_keys(): - self.increment(pool_key) - return + if controller is None: + controller = get_cached_controller() + for lib in _iter_claimable_shared_pools(controller.lib_controllers): + self._increment(lib.filepath) + + def _increment(self, filepath: str): + """Count one more scope for a single shared pool.""" self.sync() counts = self._counts() - counts[key] = counts.get(key, 0) + 1 + counts[filepath] = counts.get(filepath, 0) + 1 - def decrement(self, key: str | None = None): - """Decrement ownership, releasing a non-sticky claim when it hits zero. + def decrement(self, controller: ThreadpoolController | None = None): + """Decrement pools of ``controller`` that this thread currently counts. - With no ``key``, every pool this thread currently counts is - decremented once. That closes a matching ``increment()`` and also any - pool incremented on its own. + ``controller`` defaults to ``get_cached_controller()``. A pool whose + count hits zero releases its non-sticky claim. This closes a matching + ``increment()`` for that controller. + """ + if controller is None: + controller = get_cached_controller() + self.sync() + counts = self._counts() + seen = set() + for lib in controller.lib_controllers: + if not lib._has_shared_pool(): + continue + if lib.filepath in seen or lib.filepath not in counts: + continue + seen.add(lib.filepath) + self._decrement(lib.filepath) + + def _decrement(self, filepath: str): + """Count one fewer scope for a single shared pool. + + Reaching zero releases that pool's non-sticky claim. """ - if key is None: - self.sync() - for pool_key in list(self._counts()): - self.decrement(pool_key) - return self.sync() counts = self._counts() - current = counts.get(key, 0) + current = counts.get(filepath, 0) if current <= 0: return current -= 1 if current == 0: - counts.pop(key, None) - _release_shared_pool_claim(key) + counts.pop(filepath, None) + _release_shared_pool_claim(filepath) else: - counts[key] = current + counts[filepath] = current SHARED_POOL_CONTROL_OWNERSHIP = _SharedPoolControlOwnership() @@ -1308,7 +1312,7 @@ def __init__(self, controller, *, limits=None, user_api=None): self._controller = controller self._limits, self._user_api = self._check_params(limits, user_api) self._original_info = self._controller.info() - self._shared_pool_claim_keys: list[str] = [] + self._shared_pool_filepaths: list[str] = [] # A plain call never reaches __exit__, so this increment is what makes the claim # stick until clear_controlling_thread(). A `with` block decrements the same # increment on the way out. @@ -1459,38 +1463,38 @@ def _acquire_shared_pool_claims(self): one copy does not pin another. A finished holder does not block this scope: counting first keeps the replacement claim non-sticky. """ - self._shared_pool_claim_keys = [] + self._shared_pool_filepaths = [] targets = [ lib for lib in self._controller.lib_controllers if self._num_threads_for(lib) is not None ] - for key, _lib in _iter_claimable_shared_pools(targets): - SHARED_POOL_CONTROL_OWNERSHIP.increment(key) - self._shared_pool_claim_keys.append(key) + for lib in _iter_claimable_shared_pools(targets): + SHARED_POOL_CONTROL_OWNERSHIP._increment(lib.filepath) + self._shared_pool_filepaths.append(lib.filepath) def _release_shared_pool_claims(self): - keys = self._shared_pool_claim_keys - self._shared_pool_claim_keys = [] - for key in keys: - SHARED_POOL_CONTROL_OWNERSHIP.decrement(key) + filepaths = self._shared_pool_filepaths + self._shared_pool_filepaths = [] + for filepath in filepaths: + SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath) def _drop_unheld_shared_pool_claims(self): """Drop counts for libraries another thread won the race for.""" this = current_thread() kept = [] - for key in self._shared_pool_claim_keys: + for filepath in self._shared_pool_filepaths: with _SHARED_POOL_STATES_LOCK: - state = _SHARED_POOL_STATES.get(key) + state = _SHARED_POOL_STATES.get(filepath) held = False if state is not None: with state.lock: held = state.controlling_thread is this if held: - kept.append(key) + kept.append(filepath) else: - SHARED_POOL_CONTROL_OWNERSHIP.decrement(key) - self._shared_pool_claim_keys = kept + SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath) + self._shared_pool_filepaths = kept def _apply_limits_with_shared_pool_claims(self): self._acquire_shared_pool_claims() @@ -1518,7 +1522,7 @@ class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter, ContextDecorator): def __init__(self, controller, *, limits=None, user_api=None): self._limits, self._user_api = self._check_params(limits, user_api) self._controller = controller - self._shared_pool_claim_keys = [] + self._shared_pool_filepaths = [] def __enter__(self): # Set the limits here and not in __init__ so they apply when the decorated From 35c6ac2b7e65db271e94c3d8fb1cf56d31880717 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 15:12:13 +0100 Subject: [PATCH 07/21] MNT Make ownership counter helpers private and document per-copy control. Only increment() and decrement() are part of SHARED_POOL_CONTROL_OWNERSHIP. Update the README note on OpenBLAS pthreads/Windows to describe one controlling thread per loaded copy. Co-authored-by: Cursor --- README.md | 15 +++++------ tests/test_threadpoolctl.py | 50 ++++++++++++++++++------------------- threadpoolctl.py | 26 +++++++++---------- 3 files changed, 46 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index 989abbb2..858c2a38 100644 --- a/README.md +++ b/README.md @@ -274,13 +274,14 @@ OpenBLAS+pthreads) and in each worker thread. > **Why you need both:** When setting a limit on OpenBLAS with pthreads or the > Windows threading layer, setting the limit affects the whole process, because -> it uses a process-wide shared thread pool. (This in contrast to some other -> BLAS/OpenMP libraries, or OpenBLAS backed by OpenMP, that allow setting the -> limit per-thread.) Having multiple Python threads set limits in this case -> makes no sense, since they would effectively be arguing over a globally shared -> resource. ``threadpoolctl`` therefore constrains the limiting API in the -> OpenBLAS pthreads/Windows case to only work from the first thread that uses -> it, the so-called controlling thread. +> each loaded copy uses one process-wide shared thread pool. (This in contrast +> to some other BLAS/OpenMP libraries, or OpenBLAS backed by OpenMP, that allow +> setting the limit per-thread.) Having multiple Python threads set limits on +> the same copy makes no sense, since they would be arguing over that copy's +> pool. Each copy therefore has its own controlling thread, the first thread to +> call ``set_num_threads()`` on it. Other threads do nothing. Once that thread +> has finished, or after ``OpenBLASController.clear_controlling_thread()``, the +> next call can take the claim. > > For many other libraries, for example OpenMP, the limit is set per thread. As > a result, you want to set limits both ways in generic code that doesn't know diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 652a43a7..5207cd1b 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -949,9 +949,9 @@ def hold(): holder.start() started.acquire() SHARED_POOL_CONTROL_OWNERSHIP.increment() - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(free_a.filepath) == 1 - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(free_b.filepath) == 1 - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(held.filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(free_a.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(free_b.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(held.filepath) == 0 free_a.set_num_threads(1) free_b.set_num_threads(1) @@ -969,8 +969,8 @@ def hold(): SHARED_POOL_CONTROL_OWNERSHIP.decrement() assert free_a._shared_pool_state().controlling_thread is None assert free_b._shared_pool_state().controlling_thread is None - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(free_a.filepath) == 0 - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(free_b.filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(free_a.filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(free_b.filepath) == 0 finally: finish.release() if holder.ident is not None: @@ -978,7 +978,7 @@ def hold(): for lib in (free_a, free_b, held): controller.lib_controllers.remove(lib) threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) - SHARED_POOL_CONTROL_OWNERSHIP.reset() + SHARED_POOL_CONTROL_OWNERSHIP._reset() def test_increment_decrement_uses_the_given_controller(): @@ -991,28 +991,28 @@ def test_increment_decrement_uses_the_given_controller(): cached.lib_controllers.append(other) try: SHARED_POOL_CONTROL_OWNERSHIP.increment(controller) - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(chosen.filepath) == 1 - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(other.filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(chosen.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(other.filepath) == 0 chosen.set_num_threads(1) assert chosen._shared_pool_state().sticky is False SHARED_POOL_CONTROL_OWNERSHIP.increment() - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(chosen.filepath) == 1 - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(other.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(chosen.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(other.filepath) == 1 SHARED_POOL_CONTROL_OWNERSHIP.decrement(controller) assert chosen._shared_pool_state().controlling_thread is None - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(chosen.filepath) == 0 - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(other.filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(chosen.filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(other.filepath) == 1 SHARED_POOL_CONTROL_OWNERSHIP.decrement() - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(other.filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(other.filepath) == 0 finally: cached.lib_controllers.remove(other) for lib in (chosen, other): threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) - SHARED_POOL_CONTROL_OWNERSHIP.reset() + SHARED_POOL_CONTROL_OWNERSHIP._reset() def _require_resizable_openblas(): @@ -1153,7 +1153,7 @@ def test_openblas_scoped_exit_releases_only_that_library(): SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_b) state_b.controlling_thread = current_thread() - state_b.sticky = SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_b) == 0 + state_b.sticky = SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_b) == 0 assert state_b.sticky is False SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath_b) @@ -1172,11 +1172,11 @@ def test_openblas_scoped_exit_releases_only_that_library(): assert state_b.controlling_thread is None assert state_a.controlling_thread is current_thread() assert state_a.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_a) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_a) == 1 finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath_a, None) threadpoolctl._SHARED_POOL_STATES.pop(filepath_b, None) - SHARED_POOL_CONTROL_OWNERSHIP.reset() + SHARED_POOL_CONTROL_OWNERSHIP._reset() def test_scoped_limit_replacing_dead_holder_is_not_sticky(): @@ -1199,13 +1199,13 @@ def test_scoped_limit_replacing_dead_holder_is_not_sticky(): assert lib.applied == 1 assert state.controlling_thread is current_thread() assert state.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath) == 1 limiter._release_shared_pool_claims() assert state.controlling_thread is None - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath) == 0 finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath, None) - SHARED_POOL_CONTROL_OWNERSHIP.reset() + SHARED_POOL_CONTROL_OWNERSHIP._reset() class _OtherSharedPool: @@ -1232,10 +1232,10 @@ def test_clear_controlling_thread_leaves_other_libraries(): OpenBLASController.clear_controlling_thread() assert state_openblas.controlling_thread is None assert state_openblas.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_openblas) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_openblas) == 0 assert state_other.controlling_thread is current_thread() assert state_other.sticky is True - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_other) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_other) == 1 state_openblas.controlling_thread = current_thread() state_openblas.sticky = True @@ -1249,12 +1249,12 @@ def force_clear(): assert state_openblas.sticky is False assert state_other.controlling_thread is current_thread() assert state_other.sticky is True - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_openblas) == 0 - assert SHARED_POOL_CONTROL_OWNERSHIP.count_for(filepath_other) == 1 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_openblas) == 0 + assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_other) == 1 finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath_openblas, None) threadpoolctl._SHARED_POOL_STATES.pop(filepath_other, None) - SHARED_POOL_CONTROL_OWNERSHIP.reset() + SHARED_POOL_CONTROL_OWNERSHIP._reset() # skip test if not run in a azure pipelines job since it relies on a specific flexiblas diff --git a/threadpoolctl.py b/threadpoolctl.py index bb55beb7..27c2169b 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -490,8 +490,8 @@ class _SharedPoolState: # NumPy's and SciPy's OpenBLAS, have separate pools, so each copy has its own record and # its own controlling thread. # -# Lock order: take _SHARED_POOL_STATES_LOCK before state.lock. Do not call sync() or -# count_for() while holding state.lock. Both take the states lock, and releasing a +# Lock order: take _SHARED_POOL_STATES_LOCK before state.lock. Do not call _sync() or +# _count_for() while holding state.lock. Both take the states lock, and releasing a # claim takes the states lock before the pool lock. _SHARED_POOL_STATES: dict[str, _SharedPoolState] = {} _SHARED_POOL_STATES_LOCK = Lock() @@ -579,7 +579,7 @@ def _clear_shared_pool_claims(controller_class: type, force: bool = False): generation = _shared_pool_generations.get(controller_class, 0) + 1 _shared_pool_generations[controller_class] = generation if force or cleared_this_thread: - SHARED_POOL_CONTROL_OWNERSHIP.reset(controller_class) + SHARED_POOL_CONTROL_OWNERSHIP._reset(controller_class) class _SharedPoolControllerMixin: @@ -623,13 +623,13 @@ def _set_shared_pool_num_threads(self, num_threads): count is still 0, is sticky: a later context manager must not drop it. A call made while a limiter scope is active is released with that scope. """ - SHARED_POOL_CONTROL_OWNERSHIP.sync() + SHARED_POOL_CONTROL_OWNERSHIP._sync() this_thread = current_thread() state = self._shared_pool_state() controlling_thread = state.controlling_thread if controlling_thread is None or not controlling_thread.is_alive(): - # Read the count before state.lock. count_for() takes the states lock. - count = SHARED_POOL_CONTROL_OWNERSHIP.count_for(self.filepath) + # Read the count before state.lock. _count_for() takes the states lock. + count = SHARED_POOL_CONTROL_OWNERSHIP._count_for(self.filepath) with state.lock: holder = state.controlling_thread if holder is None or not holder.is_alive(): @@ -817,7 +817,7 @@ def _generations(self) -> dict[type, int]: self.generations = generations return generations - def sync(self): + def _sync(self): """Drop counts for libraries another thread has force-cleared.""" with _SHARED_POOL_STATES_LOCK: generations = dict(_shared_pool_generations) @@ -841,7 +841,7 @@ def sync(self): for controller_class in stale: local[controller_class] = generations[controller_class] - def reset(self, controller_class: type | None = None): + def _reset(self, controller_class: type | None = None): """Forget scopes on this thread. With ``controller_class``, only that class. Without, every library. @@ -864,9 +864,9 @@ def reset(self, controller_class: type | None = None): counts.pop(filepath, None) self._generations()[controller_class] = generation - def count_for(self, filepath: str) -> int: + def _count_for(self, filepath: str) -> int: """Active scopes for one shared pool on this thread.""" - self.sync() + self._sync() return self._counts().get(filepath, 0) def increment(self, controller: ThreadpoolController | None = None): @@ -882,7 +882,7 @@ def increment(self, controller: ThreadpoolController | None = None): def _increment(self, filepath: str): """Count one more scope for a single shared pool.""" - self.sync() + self._sync() counts = self._counts() counts[filepath] = counts.get(filepath, 0) + 1 @@ -895,7 +895,7 @@ def decrement(self, controller: ThreadpoolController | None = None): """ if controller is None: controller = get_cached_controller() - self.sync() + self._sync() counts = self._counts() seen = set() for lib in controller.lib_controllers: @@ -911,7 +911,7 @@ def _decrement(self, filepath: str): Reaching zero releases that pool's non-sticky claim. """ - self.sync() + self._sync() counts = self._counts() current = counts.get(filepath, 0) if current <= 0: From 38a6ebfe215c6ad1ac9aeca28395ea4c6bf2ee77 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 15:14:09 +0100 Subject: [PATCH 08/21] DOC Describe the manual shared-pool scope API in the changelog. Co-authored-by: Cursor --- CHANGES.md | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 1deccfae..8d2df355 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -15,10 +15,9 @@ decorators release it when the outermost scope that took the claim exits. A later scope does not drop a claim it did not take, for example an OpenMP-only limit or ``limits=None``. A thread that has finished no longer holds the - claim; the next ``set_num_threads()`` can take it. - ``SHARED_POOL_CONTROL_OWNERSHIP.increment(controller)`` and - ``decrement(controller)`` open and close that scope for the shared pools of - ``controller``, which defaults to ``get_cached_controller()``. + claim; the next ``set_num_threads()`` can take it. The + ``SHARED_POOL_CONTROL_OWNERSHIP`` singleton can be used to manually manage the + scope of ``set_num_threads`` calls. https://github.com/joblib/threadpoolctl/pull/275 https://github.com/joblib/threadpoolctl/pull/289 https://github.com/joblib/threadpoolctl/pull/XXX From 9874e94e8589fc0a038ed77da56c6951e8efb5cb Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 16:27:38 +0100 Subject: [PATCH 09/21] MNT Pin unset non-OpenMP OpenBLAS thread count in the test fixture OpenBLAS reports 0 threads until a BLAS call has picked a count, and set_num_threads(0) cannot return to that state. The autouse fixture now sets the processor count when the count is unset, before and after each test, and clears the controlling thread. Co-authored-by: Cursor --- conftest.py | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/conftest.py b/conftest.py index 101040c7..84cbaa67 100644 --- a/conftest.py +++ b/conftest.py @@ -9,11 +9,32 @@ ] +def _ensure_openblas_thread_count_is_restorable(): + """Replace an unset non-OpenMP OpenBLAS count with the processor count. + + ``get_num_threads() == 0`` means OpenBLAS has not chosen an active count. + ``set_num_threads(0)`` does not return to that state: a non-positive value + means the number of worker threads already created, and that number only + grows. A positive count can be saved and restored. + """ + from threadpoolctl import OpenBLASController, ThreadpoolController + + controller = ThreadpoolController().select(internal_api="openblas") + if controller: + for lib in controller.lib_controllers: + if lib.threading_layer == "openmp" or lib.get_num_threads() != 0: + continue + get_num_procs = lib._get_symbol("openblas_get_num_procs") + nprocs = get_num_procs() if get_num_procs is not None else 1 + if nprocs >= 1: + lib.set_num_threads(nprocs) + OpenBLASController.clear_controlling_thread(force=True) + + @pytest.fixture(autouse=True) def reset_openblas_controlling_thread(): - from threadpoolctl import OpenBLASController - + _ensure_openblas_thread_count_is_restorable() try: yield finally: - OpenBLASController.clear_controlling_thread(force=True) + _ensure_openblas_thread_count_is_restorable() From 67d7278df8861637834c890b9cfb87d2f0610bf1 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 16:27:41 +0100 Subject: [PATCH 10/21] MNT Rename the shared-pool ownership singleton to CONTROLLING_THREAD_OWNERSHIP Restore the name used before the per-library ownership change. Co-authored-by: Cursor --- CHANGES.md | 2 +- tests/test_threadpoolctl.py | 92 ++++++++++++++++++------------------- threadpoolctl.py | 16 +++---- 3 files changed, 55 insertions(+), 55 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 8d2df355..701f3937 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -16,7 +16,7 @@ later scope does not drop a claim it did not take, for example an OpenMP-only limit or ``limits=None``. A thread that has finished no longer holds the claim; the next ``set_num_threads()`` can take it. The - ``SHARED_POOL_CONTROL_OWNERSHIP`` singleton can be used to manually manage the + ``CONTROLLING_THREAD_OWNERSHIP`` singleton can be used to manually manage the scope of ``set_num_threads`` calls. https://github.com/joblib/threadpoolctl/pull/275 https://github.com/joblib/threadpoolctl/pull/289 diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 5207cd1b..9a93d08b 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -16,7 +16,7 @@ from threadpoolctl import ( _ALL_PREFIXES, _ALL_USER_APIS, - SHARED_POOL_CONTROL_OWNERSHIP, + CONTROLLING_THREAD_OWNERSHIP, LibController, OpenBLASController, ThreadpoolController, @@ -900,14 +900,14 @@ def test_openblas_controlling_thread_per_pool_counter(request): request.addfinalizer(_restore_openblas_thread_counts(libs)) filepath = lib.filepath - SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath) - SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath) lib.set_num_threads(1) assert lib._shared_pool_state().controlling_thread is current_thread() assert not lib._shared_pool_state().sticky - SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath) + CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) assert lib._shared_pool_state().controlling_thread is current_thread() - SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath) + CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) assert lib._shared_pool_state().controlling_thread is None @@ -948,10 +948,10 @@ def hold(): try: holder.start() started.acquire() - SHARED_POOL_CONTROL_OWNERSHIP.increment() - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(free_a.filepath) == 1 - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(free_b.filepath) == 1 - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(held.filepath) == 0 + CONTROLLING_THREAD_OWNERSHIP.increment() + assert CONTROLLING_THREAD_OWNERSHIP._count_for(free_a.filepath) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(free_b.filepath) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(held.filepath) == 0 free_a.set_num_threads(1) free_b.set_num_threads(1) @@ -963,14 +963,14 @@ def hold(): assert free_a._shared_pool_state().sticky is False assert free_b._shared_pool_state().sticky is False - SHARED_POOL_CONTROL_OWNERSHIP.increment() - SHARED_POOL_CONTROL_OWNERSHIP.decrement() + CONTROLLING_THREAD_OWNERSHIP.increment() + CONTROLLING_THREAD_OWNERSHIP.decrement() assert free_a._shared_pool_state().controlling_thread is current_thread() - SHARED_POOL_CONTROL_OWNERSHIP.decrement() + CONTROLLING_THREAD_OWNERSHIP.decrement() assert free_a._shared_pool_state().controlling_thread is None assert free_b._shared_pool_state().controlling_thread is None - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(free_a.filepath) == 0 - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(free_b.filepath) == 0 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(free_a.filepath) == 0 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(free_b.filepath) == 0 finally: finish.release() if holder.ident is not None: @@ -978,7 +978,7 @@ def hold(): for lib in (free_a, free_b, held): controller.lib_controllers.remove(lib) threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) - SHARED_POOL_CONTROL_OWNERSHIP._reset() + CONTROLLING_THREAD_OWNERSHIP._reset() def test_increment_decrement_uses_the_given_controller(): @@ -990,29 +990,29 @@ def test_increment_decrement_uses_the_given_controller(): cached = get_cached_controller() cached.lib_controllers.append(other) try: - SHARED_POOL_CONTROL_OWNERSHIP.increment(controller) - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(chosen.filepath) == 1 - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(other.filepath) == 0 + CONTROLLING_THREAD_OWNERSHIP.increment(controller) + assert CONTROLLING_THREAD_OWNERSHIP._count_for(chosen.filepath) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(other.filepath) == 0 chosen.set_num_threads(1) assert chosen._shared_pool_state().sticky is False - SHARED_POOL_CONTROL_OWNERSHIP.increment() - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(chosen.filepath) == 1 - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(other.filepath) == 1 + CONTROLLING_THREAD_OWNERSHIP.increment() + assert CONTROLLING_THREAD_OWNERSHIP._count_for(chosen.filepath) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(other.filepath) == 1 - SHARED_POOL_CONTROL_OWNERSHIP.decrement(controller) + CONTROLLING_THREAD_OWNERSHIP.decrement(controller) assert chosen._shared_pool_state().controlling_thread is None - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(chosen.filepath) == 0 - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(other.filepath) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(chosen.filepath) == 0 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(other.filepath) == 1 - SHARED_POOL_CONTROL_OWNERSHIP.decrement() - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(other.filepath) == 0 + CONTROLLING_THREAD_OWNERSHIP.decrement() + assert CONTROLLING_THREAD_OWNERSHIP._count_for(other.filepath) == 0 finally: cached.lib_controllers.remove(other) for lib in (chosen, other): threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) - SHARED_POOL_CONTROL_OWNERSHIP._reset() + CONTROLLING_THREAD_OWNERSHIP._reset() def _require_resizable_openblas(): @@ -1151,12 +1151,12 @@ def test_openblas_scoped_exit_releases_only_that_library(): state_a.controlling_thread = current_thread() state_a.sticky = True - SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_b) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_b) state_b.controlling_thread = current_thread() - state_b.sticky = SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_b) == 0 + state_b.sticky = CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_b) == 0 assert state_b.sticky is False - SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath_b) + CONTROLLING_THREAD_OWNERSHIP._decrement(filepath_b) assert state_b.controlling_thread is None assert state_a.controlling_thread is current_thread() assert state_a.sticky is True @@ -1164,19 +1164,19 @@ def test_openblas_scoped_exit_releases_only_that_library(): # A plain call is non-sticky and is held by a leftover count. Ending a # scope on B must not release A. state_a.sticky = False - SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_a) - SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_b) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_a) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_b) state_b.controlling_thread = current_thread() state_b.sticky = False - SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath_b) + CONTROLLING_THREAD_OWNERSHIP._decrement(filepath_b) assert state_b.controlling_thread is None assert state_a.controlling_thread is current_thread() assert state_a.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_a) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_a) == 1 finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath_a, None) threadpoolctl._SHARED_POOL_STATES.pop(filepath_b, None) - SHARED_POOL_CONTROL_OWNERSHIP._reset() + CONTROLLING_THREAD_OWNERSHIP._reset() def test_scoped_limit_replacing_dead_holder_is_not_sticky(): @@ -1199,13 +1199,13 @@ def test_scoped_limit_replacing_dead_holder_is_not_sticky(): assert lib.applied == 1 assert state.controlling_thread is current_thread() assert state.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath) == 1 limiter._release_shared_pool_claims() assert state.controlling_thread is None - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath) == 0 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath) == 0 finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath, None) - SHARED_POOL_CONTROL_OWNERSHIP._reset() + CONTROLLING_THREAD_OWNERSHIP._reset() class _OtherSharedPool: @@ -1226,20 +1226,20 @@ def test_clear_controlling_thread_leaves_other_libraries(): state_openblas.sticky = True state_other.controlling_thread = current_thread() state_other.sticky = True - SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_openblas) - SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_other) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_openblas) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_other) OpenBLASController.clear_controlling_thread() assert state_openblas.controlling_thread is None assert state_openblas.sticky is False - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_openblas) == 0 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_openblas) == 0 assert state_other.controlling_thread is current_thread() assert state_other.sticky is True - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_other) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_other) == 1 state_openblas.controlling_thread = current_thread() state_openblas.sticky = True - SHARED_POOL_CONTROL_OWNERSHIP._increment(filepath_openblas) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_openblas) def force_clear(): OpenBLASController.clear_controlling_thread(force=True) @@ -1249,12 +1249,12 @@ def force_clear(): assert state_openblas.sticky is False assert state_other.controlling_thread is current_thread() assert state_other.sticky is True - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_openblas) == 0 - assert SHARED_POOL_CONTROL_OWNERSHIP._count_for(filepath_other) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_openblas) == 0 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_other) == 1 finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath_openblas, None) threadpoolctl._SHARED_POOL_STATES.pop(filepath_other, None) - SHARED_POOL_CONTROL_OWNERSHIP._reset() + CONTROLLING_THREAD_OWNERSHIP._reset() # skip test if not run in a azure pipelines job since it relies on a specific flexiblas diff --git a/threadpoolctl.py b/threadpoolctl.py index 27c2169b..b1a4c225 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -59,7 +59,7 @@ __version__ = "3.8.0.dev0" __all__ = [ - "SHARED_POOL_CONTROL_OWNERSHIP", + "CONTROLLING_THREAD_OWNERSHIP", "LibController", "ThreadpoolController", "get_cached_controller", @@ -579,7 +579,7 @@ def _clear_shared_pool_claims(controller_class: type, force: bool = False): generation = _shared_pool_generations.get(controller_class, 0) + 1 _shared_pool_generations[controller_class] = generation if force or cleared_this_thread: - SHARED_POOL_CONTROL_OWNERSHIP._reset(controller_class) + CONTROLLING_THREAD_OWNERSHIP._reset(controller_class) class _SharedPoolControllerMixin: @@ -623,13 +623,13 @@ def _set_shared_pool_num_threads(self, num_threads): count is still 0, is sticky: a later context manager must not drop it. A call made while a limiter scope is active is released with that scope. """ - SHARED_POOL_CONTROL_OWNERSHIP._sync() + CONTROLLING_THREAD_OWNERSHIP._sync() this_thread = current_thread() state = self._shared_pool_state() controlling_thread = state.controlling_thread if controlling_thread is None or not controlling_thread.is_alive(): # Read the count before state.lock. _count_for() takes the states lock. - count = SHARED_POOL_CONTROL_OWNERSHIP._count_for(self.filepath) + count = CONTROLLING_THREAD_OWNERSHIP._count_for(self.filepath) with state.lock: holder = state.controlling_thread if holder is None or not holder.is_alive(): @@ -924,7 +924,7 @@ def _decrement(self, filepath: str): counts[filepath] = current -SHARED_POOL_CONTROL_OWNERSHIP = _SharedPoolControlOwnership() +CONTROLLING_THREAD_OWNERSHIP = _SharedPoolControlOwnership() class BLISController(LibController): @@ -1470,14 +1470,14 @@ def _acquire_shared_pool_claims(self): if self._num_threads_for(lib) is not None ] for lib in _iter_claimable_shared_pools(targets): - SHARED_POOL_CONTROL_OWNERSHIP._increment(lib.filepath) + CONTROLLING_THREAD_OWNERSHIP._increment(lib.filepath) self._shared_pool_filepaths.append(lib.filepath) def _release_shared_pool_claims(self): filepaths = self._shared_pool_filepaths self._shared_pool_filepaths = [] for filepath in filepaths: - SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath) + CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) def _drop_unheld_shared_pool_claims(self): """Drop counts for libraries another thread won the race for.""" @@ -1493,7 +1493,7 @@ def _drop_unheld_shared_pool_claims(self): if held: kept.append(filepath) else: - SHARED_POOL_CONTROL_OWNERSHIP._decrement(filepath) + CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) self._shared_pool_filepaths = kept def _apply_limits_with_shared_pool_claims(self): From e2b0600f493c6220a2043c3e35650b3860430861 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 16:27:43 +0100 Subject: [PATCH 11/21] DOC Keep the OpenBLAS pthreads/Windows README section unchanged Revert the per-copy wording in that section and the clear_controlling_thread() call in the restore example, so the section stays focused on the limits. Co-authored-by: Cursor --- README.md | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 858c2a38..7fd25e45 100644 --- a/README.md +++ b/README.md @@ -323,7 +323,7 @@ with ( You can also operate without a context manager: ```python -from threadpoolctl import OpenBLASController, ThreadpoolController +from threadpoolctl import ThreadpoolController controller = ThreadpoolController() try: @@ -333,10 +333,6 @@ try: pool.map(somefunc, someargs) finally: limiter.restore_original_limits() - # restore puts the thread counts back. It leaves this thread in control - # of any non-OpenMP OpenBLAS it claimed. Drop that when other threads - # should be allowed to change those limits again. - OpenBLASController.clear_controlling_thread() ``` ### Switching Back And Forth Between Main Thread and Python Threads @@ -386,10 +382,8 @@ with CONTROLLER.limit(limits=2): ### OpenBLAS with pthreads/Windows As mentioned above, OpenBLAS with pthreads or Windows threading layer uses a -process-wide shared thread pool. Each loaded copy has its own pool and its own -controlling thread, so a limit on one copy does not decide who may limit -another. As a result, there are only two limits that make sense in most -circumstances: +process-wide shared thread pool. As a result, there are only two limits that +make sense in most circumstances: * A limit of `1`, which effectively means no OpenBLAS thread pool: all operations run inline in the calling thread. This is a good setting when you From 13f075808af82de0b7347e66926ae77ce9f1e3d9 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 16:31:12 +0100 Subject: [PATCH 12/21] MNT Rename _require_resizable_openblas to _require_resizable_non_openmp_openblas The helper only returns non-OpenMP OpenBLAS libraries, so the name says so. Co-authored-by: Cursor --- tests/test_threadpoolctl.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 9a93d08b..19f2643a 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -1015,7 +1015,7 @@ def test_increment_decrement_uses_the_given_controller(): CONTROLLING_THREAD_OWNERSHIP._reset() -def _require_resizable_openblas(): +def _require_resizable_non_openmp_openblas(): libs = _non_openmp_openblas() lib = libs[0] orig = lib.get_num_threads() @@ -1027,7 +1027,7 @@ def _require_resizable_openblas(): def test_plain_threadpool_limits_claim_survives_later_scopes(request): # A plain call holds the controlling thread. A later with-block, including # one that does not touch OpenBLAS, must not drop that claim. - libs, lib, orig = _require_resizable_openblas() + libs, lib, orig = _require_resizable_non_openmp_openblas() request.addfinalizer(_restore_openblas_thread_counts(libs)) threadpool_limits(limits=1) @@ -1065,7 +1065,7 @@ def other_thread(): def test_direct_set_num_threads_is_not_released_by_limiter_scopes(request): - libs, lib, orig = _require_resizable_openblas() + libs, lib, orig = _require_resizable_non_openmp_openblas() request.addfinalizer(_restore_openblas_thread_counts(libs)) lib.set_num_threads(1) @@ -1095,7 +1095,7 @@ def other_thread(): def test_restore_original_limits_keeps_openblas_controlling_thread(request): - libs, lib, orig = _require_resizable_openblas() + libs, lib, orig = _require_resizable_non_openmp_openblas() request.addfinalizer(_restore_openblas_thread_counts(libs)) limiter = threadpool_limits(limits=1) From e999cf90bfd4d92b35f573e65a5a1023b339e668 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 16:52:49 +0100 Subject: [PATCH 13/21] MNT Split OpenBLAS claim clearing into per-library and per-class methods clear_controlling_thread() is now an instance method that drops the claim on one loaded copy. The former class method is clear_all_controlling_threads() and still clears every copy of that controller class. Forced clears bump a generation per filepath, so clearing one copy does not drop another copy's counts. Co-authored-by: Cursor --- CHANGES.md | 11 ++-- README.md | 4 +- conftest.py | 2 +- tests/test_threadpoolctl.py | 62 ++++++++++++++++--- threadpoolctl.py | 120 +++++++++++++++++++++++++----------- 5 files changed, 145 insertions(+), 54 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 701f3937..f562fd8b 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -10,12 +10,11 @@ process-wide thread pool, having multiple Python threads fight over control wouldn't work well. Each loaded OpenBLAS library has its own pool and its own controlling thread. A plain ``threadpool_limits(...)`` call or a direct - ``set_num_threads()`` keeps that role until - ``OpenBLASController.clear_controlling_thread()``. Context managers and - decorators release it when the outermost scope that took the claim exits. A - later scope does not drop a claim it did not take, for example an OpenMP-only - limit or ``limits=None``. A thread that has finished no longer holds the - claim; the next ``set_num_threads()`` can take it. The + ``set_num_threads()`` keeps that role until ``clear_controlling_thread()`` on + that library. ``OpenBLASController.clear_all_controlling_threads()`` drops + every loaded copy. Context managers and decorators release it when the + outermost scope that took the claim exits. A thread that has finished no + longer holds the claim; the next ``set_num_threads()`` can take it. The ``CONTROLLING_THREAD_OWNERSHIP`` singleton can be used to manually manage the scope of ``set_num_threads`` calls. https://github.com/joblib/threadpoolctl/pull/275 diff --git a/README.md b/README.md index 7fd25e45..ec6d98b0 100644 --- a/README.md +++ b/README.md @@ -280,8 +280,8 @@ OpenBLAS+pthreads) and in each worker thread. > the same copy makes no sense, since they would be arguing over that copy's > pool. Each copy therefore has its own controlling thread, the first thread to > call ``set_num_threads()`` on it. Other threads do nothing. Once that thread -> has finished, or after ``OpenBLASController.clear_controlling_thread()``, the -> next call can take the claim. +> has finished, or after ``clear_controlling_thread()`` on that copy, the next +> call can take the claim. > > For many other libraries, for example OpenMP, the limit is set per thread. As > a result, you want to set limits both ways in generic code that doesn't know diff --git a/conftest.py b/conftest.py index 84cbaa67..d334ac42 100644 --- a/conftest.py +++ b/conftest.py @@ -28,7 +28,7 @@ def _ensure_openblas_thread_count_is_restorable(): nprocs = get_num_procs() if get_num_procs is not None else 1 if nprocs >= 1: lib.set_num_threads(nprocs) - OpenBLASController.clear_controlling_thread(force=True) + OpenBLASController.clear_all_controlling_threads(force=True) @pytest.fixture(autouse=True) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 19f2643a..7a7856f6 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -856,11 +856,11 @@ def _restore_openblas_thread_counts(libs): originals = [(lib, lib.get_num_threads()) for lib in libs] def restore(): - OpenBLASController.clear_controlling_thread(force=True) + OpenBLASController.clear_all_controlling_threads(force=True) for lib, orig in originals: if orig is not None: lib.set_num_threads(orig) - OpenBLASController.clear_controlling_thread(force=True) + OpenBLASController.clear_all_controlling_threads(force=True) return restore @@ -1057,7 +1057,7 @@ def other_thread(): # An explicit clear drops the plain call. The next with-block is then the # sole scope and releases the claim on the way out. - OpenBLASController.clear_controlling_thread() + OpenBLASController.clear_all_controlling_threads() assert lib._shared_pool_state().controlling_thread is None with threadpool_limits(limits=1): assert lib._shared_pool_state().controlling_thread is current_thread() @@ -1106,7 +1106,7 @@ def test_restore_original_limits_keeps_openblas_controlling_thread(request): assert lib.get_num_threads() == orig assert lib._shared_pool_state().controlling_thread is current_thread() - OpenBLASController.clear_controlling_thread() + OpenBLASController.clear_all_controlling_threads() assert lib._shared_pool_state().controlling_thread is None @@ -1212,9 +1212,9 @@ class _OtherSharedPool: """Stand-in for a future library with its own shared pool.""" -def test_clear_controlling_thread_leaves_other_libraries(): - # OpenBLASController.clear_controlling_thread, including force=True, must - # not drop a claim or a count owned by another controller class. +def test_clear_all_controlling_threads_leaves_other_libraries(): + # OpenBLASController.clear_all_controlling_threads, including force=True, + # must not drop a claim or a count owned by another controller class. filepath_openblas = "threadpoolctl-test-openblas-clear" filepath_other = "threadpoolctl-test-other-clear" state_openblas = threadpoolctl._SharedPoolState(controller_class=OpenBLASController) @@ -1229,7 +1229,7 @@ def test_clear_controlling_thread_leaves_other_libraries(): CONTROLLING_THREAD_OWNERSHIP._increment(filepath_openblas) CONTROLLING_THREAD_OWNERSHIP._increment(filepath_other) - OpenBLASController.clear_controlling_thread() + OpenBLASController.clear_all_controlling_threads() assert state_openblas.controlling_thread is None assert state_openblas.sticky is False assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_openblas) == 0 @@ -1242,7 +1242,7 @@ def test_clear_controlling_thread_leaves_other_libraries(): CONTROLLING_THREAD_OWNERSHIP._increment(filepath_openblas) def force_clear(): - OpenBLASController.clear_controlling_thread(force=True) + OpenBLASController.clear_all_controlling_threads(force=True) run_in_thread(force_clear) assert state_openblas.controlling_thread is None @@ -1257,6 +1257,50 @@ def force_clear(): CONTROLLING_THREAD_OWNERSHIP._reset() +def test_clear_controlling_thread_leaves_other_copies(): + # Clearing one library must leave another copy of the same class alone, + # including a forced clear from a thread that does not hold the claim. + lib_a = _TestSharedPoolController("threadpoolctl-test-clear-copy-a") + lib_b = _TestSharedPoolController("threadpoolctl-test-clear-copy-b") + state_a = lib_a._shared_pool_state() + state_b = lib_b._shared_pool_state() + try: + state_a.controlling_thread = current_thread() + state_a.sticky = True + state_b.controlling_thread = current_thread() + state_b.sticky = True + CONTROLLING_THREAD_OWNERSHIP._increment(lib_a.filepath) + CONTROLLING_THREAD_OWNERSHIP._increment(lib_b.filepath) + + lib_a.clear_controlling_thread() + assert state_a.controlling_thread is None + assert state_a.sticky is False + assert CONTROLLING_THREAD_OWNERSHIP._count_for(lib_a.filepath) == 0 + assert state_b.controlling_thread is current_thread() + assert state_b.sticky is True + assert CONTROLLING_THREAD_OWNERSHIP._count_for(lib_b.filepath) == 1 + + state_a.controlling_thread = current_thread() + state_a.sticky = True + CONTROLLING_THREAD_OWNERSHIP._increment(lib_a.filepath) + + def force_clear(): + lib_a.clear_controlling_thread(force=True) + + run_in_thread(force_clear) + assert state_a.controlling_thread is None + assert state_a.sticky is False + assert state_b.controlling_thread is current_thread() + assert state_b.sticky is True + assert CONTROLLING_THREAD_OWNERSHIP._count_for(lib_a.filepath) == 0 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(lib_b.filepath) == 1 + finally: + for filepath in (lib_a.filepath, lib_b.filepath): + threadpoolctl._SHARED_POOL_STATES.pop(filepath, None) + threadpoolctl._shared_pool_generations.pop(filepath, None) + CONTROLLING_THREAD_OWNERSHIP._reset() + + # 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 b1a4c225..20c94fab 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -173,7 +173,7 @@ def get_and_set() -> None: thread = Thread(target=get_and_set) thread.start() thread.join() - OpenBLASController.clear_controlling_thread(force=True) + OpenBLASController.clear_all_controlling_threads(force=True) # First, getting in the same thread as a set should always give same # number, if it's a number in a reasonable range. A possible exception @@ -199,7 +199,7 @@ def get_and_set() -> None: return "unknown" finally: set_n_threads(previous) - OpenBLASController.clear_controlling_thread(force=True) + OpenBLASController.clear_all_controlling_threads(force=True) class LibController(ABC): @@ -495,9 +495,10 @@ class _SharedPoolState: # claim takes the states lock before the pool lock. _SHARED_POOL_STATES: dict[str, _SharedPoolState] = {} _SHARED_POOL_STATES_LOCK = Lock() -# Bumped per controller class by a forced clear. A thread drops a class's counts when -# its stored generation for that class does not match. -_shared_pool_generations: dict[type, int] = {} +# Bumped per filepath by a forced clear. A thread drops a count when its stored +# generation for that filepath does not match. clear_all_controlling_threads() bumps +# every filepath of that controller class. +_shared_pool_generations: dict[str, int] = {} def _get_shared_pool_state(filepath: str, controller_class: type) -> _SharedPoolState: @@ -556,6 +557,14 @@ def _release_shared_pool_claim(filepath: str): state.controlling_thread = None +def _bump_shared_pool_generations(filepaths) -> None: + """Invalidate other threads' counts for these shared pools.""" + with _SHARED_POOL_STATES_LOCK: + for filepath in filepaths: + current = _shared_pool_generations.get(filepath, 0) + _shared_pool_generations[filepath] = current + 1 + + def _clear_shared_pool_claims(controller_class: type, force: bool = False): """Drop claims on shared pools of one controller class. @@ -565,9 +574,13 @@ def _clear_shared_pool_claims(controller_class: type, force: bool = False): """ this = current_thread() cleared_this_thread = False - for state in _snapshot_shared_pool_states(): - if state.controller_class is not controller_class: - continue + with _SHARED_POOL_STATES_LOCK: + pairs = [ + (filepath, state) + for filepath, state in _SHARED_POOL_STATES.items() + if state.controller_class is controller_class + ] + for _filepath, state in pairs: with state.lock: if force or state.controlling_thread is this: if state.controlling_thread is this: @@ -575,9 +588,7 @@ def _clear_shared_pool_claims(controller_class: type, force: bool = False): state.controlling_thread = None state.sticky = False if force: - with _SHARED_POOL_STATES_LOCK: - generation = _shared_pool_generations.get(controller_class, 0) + 1 - _shared_pool_generations[controller_class] = generation + _bump_shared_pool_generations(filepath for filepath, _state in pairs) if force or cleared_this_thread: CONTROLLING_THREAD_OWNERSHIP._reset(controller_class) @@ -589,7 +600,8 @@ class _SharedPoolControllerMixin: ``set_num_threads`` should call ``_set_shared_pool_num_threads`` when the loaded library shares one pool across Python threads. Override ``_has_shared_pool`` when that depends on how the library was built. - ``clear_controlling_thread`` drops claims of this controller class only. + ``clear_controlling_thread`` drops this library's claim. + ``clear_all_controlling_threads`` drops every library of this controller class. """ def _has_shared_pool(self): @@ -600,8 +612,36 @@ def _shared_pool_state(self) -> _SharedPoolState: """Return the process-wide claim record for this library's shared pool.""" return _get_shared_pool_state(self.filepath, type(self)) + def clear_controlling_thread(self, force: bool = False): + """Drop the claim on this library's shared pool. + + Resets this thread's ownership count for this library so a later context + manager can release a claim of its own. Other libraries stay in place, + including other copies of this controller class. It's up to the caller + to ensure the current controlling thread is done. + + By default, only the controlling thread has an effect. Any thread can + clear this library by setting ``force=True``. + """ + this = current_thread() + with _SHARED_POOL_STATES_LOCK: + state = _SHARED_POOL_STATES.get(self.filepath) + cleared_this_thread = False + if state is not None: + with state.lock: + if force or state.controlling_thread is this: + if state.controlling_thread is this: + cleared_this_thread = True + state.controlling_thread = None + state.sticky = False + if not (force or cleared_this_thread): + return + if force: + _bump_shared_pool_generations([self.filepath]) + CONTROLLING_THREAD_OWNERSHIP._forget_filepath(self.filepath) + @classmethod - def clear_controlling_thread(cls, force: bool = False): + def clear_all_controlling_threads(cls, force: bool = False): """Reset which thread can limit this controller's shared pools. Drops the claim on every library of this controller class that the @@ -665,12 +705,14 @@ class OpenBLASController(_SharedPoolControllerMixin, LibController): 2. Other threads silently do nothing. Once the controlling thread has finished, the next ``set_num_threads()`` can take the claim. - 3. ``clear_controlling_thread()`` resets the controlling thread. By - default only that thread can do so. ``force=True`` lets any thread - clear every OpenBLAS library. + 3. ``clear_controlling_thread()`` resets the controlling thread for that + library. By default only that thread can do so. ``force=True`` lets + any thread clear that library. ``clear_all_controlling_threads()`` + does the same for every loaded OpenBLAS library. A plain ``threadpool_limits(...)`` call, and a direct ``set_num_threads()``, - keep control until ``clear_controlling_thread()``. ``with + keep control until ``clear_controlling_thread()`` on that library, or + ``clear_all_controlling_threads()`` for every loaded copy. ``with threadpool_limits()`` and ``threadpool_limits.wrap()`` release control when the outermost scope that took the claim exits, including when those scopes are nested. A scope that does not touch that OpenBLAS library, and a scope @@ -793,14 +835,15 @@ class _SharedPoolControlOwnership(thread_local): A count is kept separately for each shared pool, keyed by library filepath. Hitting zero for a pool releases that pool's non-sticky controlling thread. A plain ``threadpool_limits(...)`` call increments and never decrements, so - it holds its claim until ``clear_controlling_thread()``. + it holds its claim until ``clear_controlling_thread()`` on that library. ``increment(controller)`` counts every shared pool of ``controller`` that this thread can control. ``decrement(controller)`` counts those pools back down. ``controller`` defaults to ``get_cached_controller()``. Call ``increment()`` before ``set_num_threads`` and ``decrement()`` when that - scope ends to manage the claim without a context manager. A forced clear - drops counts for that controller class only. + scope ends to manage the claim without a context manager. A forced clear of + one library drops that library's counts. ``clear_all_controlling_threads()`` + drops counts for that controller class. """ def _counts(self) -> dict[str, int]: @@ -810,7 +853,7 @@ def _counts(self) -> dict[str, int]: self.counts = counts return counts - def _generations(self) -> dict[type, int]: + def _generations(self) -> dict[str, int]: generations = getattr(self, "generations", None) if generations is None: generations = {} @@ -821,25 +864,26 @@ def _sync(self): """Drop counts for libraries another thread has force-cleared.""" with _SHARED_POOL_STATES_LOCK: generations = dict(_shared_pool_generations) - owned_filepaths = [ - (filepath, state.controller_class) - for filepath, state in _SHARED_POOL_STATES.items() - if state.controller_class is not None - ] local = self._generations() stale = { - controller_class - for controller_class, generation in generations.items() - if local.get(controller_class, 0) != generation + filepath + for filepath, generation in generations.items() + if local.get(filepath, 0) != generation } if not stale: return counts = self._counts() - for filepath, controller_class in owned_filepaths: - if controller_class in stale: - counts.pop(filepath, None) - for controller_class in stale: - local[controller_class] = generations[controller_class] + for filepath in stale: + counts.pop(filepath, None) + local[filepath] = generations[filepath] + + def _forget_filepath(self, filepath: str): + """Drop this thread's count for one pool and adopt its generation.""" + self._sync() + self._counts().pop(filepath, None) + with _SHARED_POOL_STATES_LOCK: + generation = _shared_pool_generations.get(filepath, 0) + self._generations()[filepath] = generation def _reset(self, controller_class: type | None = None): """Forget scopes on this thread. @@ -853,16 +897,20 @@ def _reset(self, controller_class: type | None = None): self.generations = generations return with _SHARED_POOL_STATES_LOCK: - generation = _shared_pool_generations.get(controller_class, 0) filepaths = [ filepath for filepath, state in _SHARED_POOL_STATES.items() if state.controller_class is controller_class ] + generations = { + filepath: _shared_pool_generations.get(filepath, 0) + for filepath in filepaths + } counts = self._counts() + local = self._generations() for filepath in filepaths: counts.pop(filepath, None) - self._generations()[controller_class] = generation + local[filepath] = generations[filepath] def _count_for(self, filepath: str) -> int: """Active scopes for one shared pool on this thread.""" From 513f001524399ab620d4d1fcfc29191ee024d4ed Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 17:26:22 +0100 Subject: [PATCH 14/21] MNT Remove unused shared-pool snapshot and share holder clearing Drop _snapshot_shared_pool_states(), which had no callers. Clearing a holder and its sticky flag now goes through _clear_shared_pool_holder() for both per-library and per-class clears. A per-library clear no longer syncs every pool before forgetting its own count. Co-authored-by: Cursor --- threadpoolctl.py | 50 ++++++++++++++++++++++++------------------------ 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/threadpoolctl.py b/threadpoolctl.py index 20c94fab..fff3156b 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -512,11 +512,6 @@ def _get_shared_pool_state(filepath: str, controller_class: type) -> _SharedPool return state -def _snapshot_shared_pool_states() -> list[_SharedPoolState]: - with _SHARED_POOL_STATES_LOCK: - return list(_SHARED_POOL_STATES.values()) - - def _other_live_thread_holds(holder, this) -> bool: """Whether ``holder`` is a different thread that still owns the pool.""" return holder is not None and holder is not this and holder.is_alive() @@ -565,6 +560,21 @@ def _bump_shared_pool_generations(filepaths) -> None: _shared_pool_generations[filepath] = current + 1 +def _clear_shared_pool_holder(state: _SharedPoolState, force: bool) -> bool: + """Drop the holder and sticky flag when ``force`` is set or this thread holds it. + + Returns whether this thread held the claim. + """ + this = current_thread() + with state.lock: + held = state.controlling_thread is this + if not (force or held): + return False + state.controlling_thread = None + state.sticky = False + return held + + def _clear_shared_pool_claims(controller_class: type, force: bool = False): """Drop claims on shared pools of one controller class. @@ -572,21 +582,16 @@ def _clear_shared_pool_claims(controller_class: type, force: bool = False): true. A forced clear also drops that class's counts on other threads the next time they sync. """ - this = current_thread() - cleared_this_thread = False with _SHARED_POOL_STATES_LOCK: pairs = [ (filepath, state) for filepath, state in _SHARED_POOL_STATES.items() if state.controller_class is controller_class ] + cleared_this_thread = False for _filepath, state in pairs: - with state.lock: - if force or state.controlling_thread is this: - if state.controlling_thread is this: - cleared_this_thread = True - state.controlling_thread = None - state.sticky = False + if _clear_shared_pool_holder(state, force): + cleared_this_thread = True if force: _bump_shared_pool_generations(filepath for filepath, _state in pairs) if force or cleared_this_thread: @@ -623,18 +628,10 @@ def clear_controlling_thread(self, force: bool = False): By default, only the controlling thread has an effect. Any thread can clear this library by setting ``force=True``. """ - this = current_thread() with _SHARED_POOL_STATES_LOCK: state = _SHARED_POOL_STATES.get(self.filepath) - cleared_this_thread = False - if state is not None: - with state.lock: - if force or state.controlling_thread is this: - if state.controlling_thread is this: - cleared_this_thread = True - state.controlling_thread = None - state.sticky = False - if not (force or cleared_this_thread): + held = _clear_shared_pool_holder(state, force) if state is not None else False + if not (force or held): return if force: _bump_shared_pool_generations([self.filepath]) @@ -878,8 +875,11 @@ def _sync(self): local[filepath] = generations[filepath] def _forget_filepath(self, filepath: str): - """Drop this thread's count for one pool and adopt its generation.""" - self._sync() + """Drop this thread's count for one pool and adopt its generation. + + The caller has already bumped the generation when this clear is forced. + Other pools stay stale until the next ``_sync``. + """ self._counts().pop(filepath, None) with _SHARED_POOL_STATES_LOCK: generation = _shared_pool_generations.get(filepath, 0) From 640427cf31853144c6e087b18a8aa112d773edd0 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 18:17:26 +0100 Subject: [PATCH 15/21] FIX Give each wrapped call its own limiter so recursion releases the claim A threadpool_limits.wrap() decorator reused one limiter for every call, so a recursive or concurrent call overwrote the claims of the outer call and left the OpenBLAS controlling thread set after the outermost call returned. Each call now gets a fresh limiter. Also rename the ownership helpers and class to match CONTROLLING_THREAD_OWNERSHIP, make the ownership counters per-thread state set in __init__, merge the duplicated holder-clearing code, and document that a forced clear during an active scope can make the next claim sticky. Co-authored-by: Cursor --- tests/test_threadpoolctl.py | 51 +++++++++--- threadpoolctl.py | 161 +++++++++++++++++------------------- 2 files changed, 116 insertions(+), 96 deletions(-) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 7a7856f6..c1123559 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -927,6 +927,39 @@ def _real_set_num_threads(self, num_threads): self.applied = num_threads +def test_recursive_wrap_releases_the_shared_pool_claim(): + # Each call of a wrapped function must release only the claim that call took. + lib = _TestSharedPoolController("threadpoolctl-test-recursive-wrap") + state = lib._shared_pool_state() + + class _Controller: + def __init__(self): + self.lib_controllers = [lib] + + def info(self): + return [{"num_threads": 4}] + + wrapped = threadpoolctl._ThreadpoolLimiterDecorator( + _Controller(), limits=1, user_api="blas" + ) + + @wrapped + def recurse(depth): + assert state.controlling_thread is current_thread() + if depth: + recurse(depth - 1) + assert state.controlling_thread is current_thread() + + try: + recurse(1) + assert state.controlling_thread is None + assert CONTROLLING_THREAD_OWNERSHIP._count_for(lib.filepath) == 0 + finally: + threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) + threadpoolctl._SHARED_POOL_GENERATIONS.pop(lib.filepath, None) + CONTROLLING_THREAD_OWNERSHIP._reset() + + def test_bare_increment_decrement_covers_every_controllable_pool(): # Every loaded shared pool this thread can control, and not a pool another # live thread already holds. @@ -1143,10 +1176,8 @@ def test_openblas_scoped_exit_releases_only_that_library(): # or a still-held claim on another library. filepath_a = "threadpoolctl-test-openblas-a" filepath_b = "threadpoolctl-test-openblas-b" - state_a = threadpoolctl._SharedPoolState() - state_b = threadpoolctl._SharedPoolState() - threadpoolctl._SHARED_POOL_STATES[filepath_a] = state_a - threadpoolctl._SHARED_POOL_STATES[filepath_b] = state_b + state_a = threadpoolctl._get_shared_pool_state(filepath_a, OpenBLASController) + state_b = threadpoolctl._get_shared_pool_state(filepath_b, OpenBLASController) try: state_a.controlling_thread = current_thread() state_a.sticky = True @@ -1200,7 +1231,7 @@ def test_scoped_limit_replacing_dead_holder_is_not_sticky(): assert state.controlling_thread is current_thread() assert state.sticky is False assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath) == 1 - limiter._release_shared_pool_claims() + limiter._exit_shared_pool_scopes() assert state.controlling_thread is None assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath) == 0 finally: @@ -1217,10 +1248,10 @@ def test_clear_all_controlling_threads_leaves_other_libraries(): # must not drop a claim or a count owned by another controller class. filepath_openblas = "threadpoolctl-test-openblas-clear" filepath_other = "threadpoolctl-test-other-clear" - state_openblas = threadpoolctl._SharedPoolState(controller_class=OpenBLASController) - state_other = threadpoolctl._SharedPoolState(controller_class=_OtherSharedPool) - threadpoolctl._SHARED_POOL_STATES[filepath_openblas] = state_openblas - threadpoolctl._SHARED_POOL_STATES[filepath_other] = state_other + state_openblas = threadpoolctl._get_shared_pool_state( + filepath_openblas, OpenBLASController + ) + state_other = threadpoolctl._get_shared_pool_state(filepath_other, _OtherSharedPool) try: state_openblas.controlling_thread = current_thread() state_openblas.sticky = True @@ -1297,7 +1328,7 @@ def force_clear(): finally: for filepath in (lib_a.filepath, lib_b.filepath): threadpoolctl._SHARED_POOL_STATES.pop(filepath, None) - threadpoolctl._shared_pool_generations.pop(filepath, None) + threadpoolctl._SHARED_POOL_GENERATIONS.pop(filepath, None) CONTROLLING_THREAD_OWNERSHIP._reset() diff --git a/threadpoolctl.py b/threadpoolctl.py index fff3156b..7bb123ce 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -477,12 +477,12 @@ def wrapper(self): class _SharedPoolState: """Who may change one shared thread pool.""" + # Controller class for this pool, so a clear of one class leaves the others. + controller_class: type controlling_thread: Thread | None = None # Set when control is taken outside a limiter scope (direct set_num_threads). sticky: bool = False lock: Lock = field(default_factory=Lock) - # Controller class for this pool, so a clear of one class leaves the others. - controller_class: type | None = None # Process-wide claims on shared thread pools, keyed by filepath. A controller opts in by @@ -498,7 +498,7 @@ class _SharedPoolState: # Bumped per filepath by a forced clear. A thread drops a count when its stored # generation for that filepath does not match. clear_all_controlling_threads() bumps # every filepath of that controller class. -_shared_pool_generations: dict[str, int] = {} +_SHARED_POOL_GENERATIONS: dict[str, int] = {} def _get_shared_pool_state(filepath: str, controller_class: type) -> _SharedPoolState: @@ -507,8 +507,6 @@ def _get_shared_pool_state(filepath: str, controller_class: type) -> _SharedPool if state is None: state = _SharedPoolState(controller_class=controller_class) _SHARED_POOL_STATES[filepath] = state - elif state.controller_class is None: - state.controller_class = controller_class return state @@ -556,8 +554,8 @@ def _bump_shared_pool_generations(filepaths) -> None: """Invalidate other threads' counts for these shared pools.""" with _SHARED_POOL_STATES_LOCK: for filepath in filepaths: - current = _shared_pool_generations.get(filepath, 0) - _shared_pool_generations[filepath] = current + 1 + current = _SHARED_POOL_GENERATIONS.get(filepath, 0) + _SHARED_POOL_GENERATIONS[filepath] = current + 1 def _clear_shared_pool_holder(state: _SharedPoolState, force: bool) -> bool: @@ -575,29 +573,6 @@ def _clear_shared_pool_holder(state: _SharedPoolState, force: bool) -> bool: return held -def _clear_shared_pool_claims(controller_class: type, force: bool = False): - """Drop claims on shared pools of one controller class. - - Pools of other classes are left untouched, including when ``force`` is - true. A forced clear also drops that class's counts on other threads the - next time they sync. - """ - with _SHARED_POOL_STATES_LOCK: - pairs = [ - (filepath, state) - for filepath, state in _SHARED_POOL_STATES.items() - if state.controller_class is controller_class - ] - cleared_this_thread = False - for _filepath, state in pairs: - if _clear_shared_pool_holder(state, force): - cleared_this_thread = True - if force: - _bump_shared_pool_generations(filepath for filepath, _state in pairs) - if force or cleared_this_thread: - CONTROLLING_THREAD_OWNERSHIP._reset(controller_class) - - class _SharedPoolControllerMixin: """Controlling-thread claims for a library with a process-wide thread pool. @@ -623,19 +598,20 @@ def clear_controlling_thread(self, force: bool = False): Resets this thread's ownership count for this library so a later context manager can release a claim of its own. Other libraries stay in place, including other copies of this controller class. It's up to the caller - to ensure the current controlling thread is done. + to ensure the current controlling thread is done. A forced clear while + that thread is still inside a scope can make the claim it takes next + sticky, so the scope does not release it. By default, only the controlling thread has an effect. Any thread can clear this library by setting ``force=True``. """ - with _SHARED_POOL_STATES_LOCK: - state = _SHARED_POOL_STATES.get(self.filepath) - held = _clear_shared_pool_holder(state, force) if state is not None else False + state = self._shared_pool_state() + held = _clear_shared_pool_holder(state, force) if not (force or held): return if force: _bump_shared_pool_generations([self.filepath]) - CONTROLLING_THREAD_OWNERSHIP._forget_filepath(self.filepath) + CONTROLLING_THREAD_OWNERSHIP._forget_pool(self.filepath) @classmethod def clear_all_controlling_threads(cls, force: bool = False): @@ -645,12 +621,27 @@ def clear_all_controlling_threads(cls, force: bool = False): current thread controls, and resets this thread's ownership counts for that class so a later context manager can release a claim of its own. Shared pools of other controller classes are left in place. It's up to - the caller to ensure the current controlling thread is done. + the caller to ensure the current controlling thread is done. A forced + clear while that thread is still inside a scope can make the claim it + takes next sticky, so the scope does not release it. By default, only the controlling thread has an effect. Any thread can clear every library of this class by setting ``force=True``. """ - _clear_shared_pool_claims(cls, force) + with _SHARED_POOL_STATES_LOCK: + class_states = [ + (filepath, state) + for filepath, state in _SHARED_POOL_STATES.items() + if state.controller_class is cls + ] + cleared_this_thread = False + for _filepath, state in class_states: + if _clear_shared_pool_holder(state, force): + cleared_this_thread = True + if force: + _bump_shared_pool_generations(filepath for filepath, _state in class_states) + if force or cleared_this_thread: + CONTROLLING_THREAD_OWNERSHIP._reset(cls) def _set_shared_pool_num_threads(self, num_threads): """Set the pool size if this thread holds or can take the claim. @@ -826,7 +817,7 @@ def _get_architecture(self): return None -class _SharedPoolControlOwnership(thread_local): +class _ControllingThreadOwnership(thread_local): """Per-pool count of limiter scopes on this thread. A count is kept separately for each shared pool, keyed by library filepath. @@ -843,25 +834,16 @@ class _SharedPoolControlOwnership(thread_local): drops counts for that controller class. """ - def _counts(self) -> dict[str, int]: - counts = getattr(self, "counts", None) - if counts is None: - counts = {} - self.counts = counts - return counts - - def _generations(self) -> dict[str, int]: - generations = getattr(self, "generations", None) - if generations is None: - generations = {} - self.generations = generations - return generations + def __init__(self): + # threading.local runs __init__ once per thread. + self.counts: dict[str, int] = {} + self.generations: dict[str, int] = {} def _sync(self): """Drop counts for libraries another thread has force-cleared.""" with _SHARED_POOL_STATES_LOCK: - generations = dict(_shared_pool_generations) - local = self._generations() + generations = dict(_SHARED_POOL_GENERATIONS) + local = self.generations stale = { filepath for filepath, generation in generations.items() @@ -869,21 +851,21 @@ def _sync(self): } if not stale: return - counts = self._counts() + counts = self.counts for filepath in stale: counts.pop(filepath, None) local[filepath] = generations[filepath] - def _forget_filepath(self, filepath: str): + def _forget_pool(self, filepath: str): """Drop this thread's count for one pool and adopt its generation. The caller has already bumped the generation when this clear is forced. Other pools stay stale until the next ``_sync``. """ - self._counts().pop(filepath, None) + self.counts.pop(filepath, None) with _SHARED_POOL_STATES_LOCK: - generation = _shared_pool_generations.get(filepath, 0) - self._generations()[filepath] = generation + generation = _SHARED_POOL_GENERATIONS.get(filepath, 0) + self.generations[filepath] = generation def _reset(self, controller_class: type | None = None): """Forget scopes on this thread. @@ -892,7 +874,7 @@ def _reset(self, controller_class: type | None = None): """ if controller_class is None: with _SHARED_POOL_STATES_LOCK: - generations = dict(_shared_pool_generations) + generations = dict(_SHARED_POOL_GENERATIONS) self.counts = {} self.generations = generations return @@ -903,11 +885,11 @@ def _reset(self, controller_class: type | None = None): if state.controller_class is controller_class ] generations = { - filepath: _shared_pool_generations.get(filepath, 0) + filepath: _SHARED_POOL_GENERATIONS.get(filepath, 0) for filepath in filepaths } - counts = self._counts() - local = self._generations() + counts = self.counts + local = self.generations for filepath in filepaths: counts.pop(filepath, None) local[filepath] = generations[filepath] @@ -915,7 +897,7 @@ def _reset(self, controller_class: type | None = None): def _count_for(self, filepath: str) -> int: """Active scopes for one shared pool on this thread.""" self._sync() - return self._counts().get(filepath, 0) + return self.counts.get(filepath, 0) def increment(self, controller: ThreadpoolController | None = None): """Increment every shared pool of ``controller`` this thread can control. @@ -931,7 +913,7 @@ def increment(self, controller: ThreadpoolController | None = None): def _increment(self, filepath: str): """Count one more scope for a single shared pool.""" self._sync() - counts = self._counts() + counts = self.counts counts[filepath] = counts.get(filepath, 0) + 1 def decrement(self, controller: ThreadpoolController | None = None): @@ -943,13 +925,9 @@ def decrement(self, controller: ThreadpoolController | None = None): """ if controller is None: controller = get_cached_controller() - self._sync() - counts = self._counts() seen = set() for lib in controller.lib_controllers: - if not lib._has_shared_pool(): - continue - if lib.filepath in seen or lib.filepath not in counts: + if lib.filepath in seen: continue seen.add(lib.filepath) self._decrement(lib.filepath) @@ -960,7 +938,7 @@ def _decrement(self, filepath: str): Reaching zero releases that pool's non-sticky claim. """ self._sync() - counts = self._counts() + counts = self.counts current = counts.get(filepath, 0) if current <= 0: return @@ -972,7 +950,7 @@ def _decrement(self, filepath: str): counts[filepath] = current -CONTROLLING_THREAD_OWNERSHIP = _SharedPoolControlOwnership() +CONTROLLING_THREAD_OWNERSHIP = _ControllingThreadOwnership() class BLISController(LibController): @@ -1373,7 +1351,7 @@ def __exit__(self, type, value, traceback): try: self.restore_original_limits() finally: - self._release_shared_pool_claims() + self._exit_shared_pool_scopes() @classmethod def wrap(cls, controller, *, limits=None, user_api=None): @@ -1493,16 +1471,14 @@ def _num_threads_for(self, lib_controller): if self._limits is None: return None if lib_controller.filepath in self._limits: - num_threads = self._limits[lib_controller.filepath] - elif lib_controller.prefix in self._limits: - num_threads = self._limits[lib_controller.prefix] - elif lib_controller.user_api in self._limits: - num_threads = self._limits[lib_controller.user_api] - else: - return None - return num_threads + return self._limits[lib_controller.filepath] + if lib_controller.prefix in self._limits: + return self._limits[lib_controller.prefix] + if lib_controller.user_api in self._limits: + return self._limits[lib_controller.user_api] + return None - def _acquire_shared_pool_claims(self): + def _enter_shared_pool_scopes(self): """Count a scope per library before setting, so the claim is not sticky. Only libraries this limiter would set, and that this thread can @@ -1521,13 +1497,13 @@ def _acquire_shared_pool_claims(self): CONTROLLING_THREAD_OWNERSHIP._increment(lib.filepath) self._shared_pool_filepaths.append(lib.filepath) - def _release_shared_pool_claims(self): + def _exit_shared_pool_scopes(self): filepaths = self._shared_pool_filepaths self._shared_pool_filepaths = [] for filepath in filepaths: CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) - def _drop_unheld_shared_pool_claims(self): + def _drop_unheld_shared_pool_scopes(self): """Drop counts for libraries another thread won the race for.""" this = current_thread() kept = [] @@ -1545,13 +1521,13 @@ def _drop_unheld_shared_pool_claims(self): self._shared_pool_filepaths = kept def _apply_limits_with_shared_pool_claims(self): - self._acquire_shared_pool_claims() + self._enter_shared_pool_scopes() try: self._set_threadpool_limits() except Exception: - self._release_shared_pool_claims() + self._exit_shared_pool_scopes() raise - self._drop_unheld_shared_pool_claims() + self._drop_unheld_shared_pool_scopes() def _set_threadpool_limits(self): """Change the maximal number of threads in selected thread pools.""" @@ -1580,6 +1556,19 @@ def __enter__(self): self._apply_limits_with_shared_pool_claims() return self + def _recreate_cm(self): + """Return a fresh limiter for this call. + + Recursive and concurrent calls of the same wrapped function must not + share ``_shared_pool_filepaths`` or ``_original_info``. + """ + recreated = object.__new__(type(self)) + recreated._controller = self._controller + recreated._limits = self._limits + recreated._user_api = self._user_api + recreated._shared_pool_filepaths = [] + return recreated + @_format_docstring( USER_APIS=", ".join(f'"{api}"' for api in _ALL_USER_APIS), From 89c5226cce6a9009792ad388bf174620c00b1623 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 18:32:40 +0100 Subject: [PATCH 16/21] MNT Skip redundant shared-pool syncs and stop locking to check the holder A generation epoch lets a thread skip the generation walk when nothing has been force-cleared. A limiter scope keeps the pool state it counted so dropping an unheld claim does not take the registry lock. The wrap decorator now copies itself per call instead of relying on ContextDecorator. Co-authored-by: Cursor --- tests/test_threadpoolctl.py | 29 +++++--- threadpoolctl.py | 138 +++++++++++++++++------------------- 2 files changed, 85 insertions(+), 82 deletions(-) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index c1123559..5485946f 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -39,6 +39,17 @@ ) +def _reset_ownership(): + """Drop this thread's scope counts and adopt the current generations.""" + ownership = CONTROLLING_THREAD_OWNERSHIP + with threadpoolctl._SHARED_POOL_STATES_LOCK: + generations = dict(threadpoolctl._SHARED_POOL_GENERATIONS) + epoch = threadpoolctl._SHARED_POOL_GENERATION_EPOCH + ownership._counts = {} + ownership._generations = generations + ownership._generation_epoch = epoch + + def is_old_openblas(lib_controller): # Possible bug in getting maximum number of threads with OpenBLAS < 0.2.16 # and OpenBLAS does not expose its version before 0.3.4. @@ -957,7 +968,7 @@ def recurse(depth): finally: threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) threadpoolctl._SHARED_POOL_GENERATIONS.pop(lib.filepath, None) - CONTROLLING_THREAD_OWNERSHIP._reset() + _reset_ownership() def test_bare_increment_decrement_covers_every_controllable_pool(): @@ -1011,7 +1022,7 @@ def hold(): for lib in (free_a, free_b, held): controller.lib_controllers.remove(lib) threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) - CONTROLLING_THREAD_OWNERSHIP._reset() + _reset_ownership() def test_increment_decrement_uses_the_given_controller(): @@ -1045,7 +1056,7 @@ def test_increment_decrement_uses_the_given_controller(): cached.lib_controllers.remove(other) for lib in (chosen, other): threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) - CONTROLLING_THREAD_OWNERSHIP._reset() + _reset_ownership() def _require_resizable_non_openmp_openblas(): @@ -1207,7 +1218,7 @@ def test_openblas_scoped_exit_releases_only_that_library(): finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath_a, None) threadpoolctl._SHARED_POOL_STATES.pop(filepath_b, None) - CONTROLLING_THREAD_OWNERSHIP._reset() + _reset_ownership() def test_scoped_limit_replacing_dead_holder_is_not_sticky(): @@ -1224,9 +1235,9 @@ def test_scoped_limit_replacing_dead_holder_is_not_sticky(): limiter = object.__new__(threadpoolctl._ThreadpoolLimiter) limiter._controller = type("C", (), {"lib_controllers": [lib]})() limiter._limits = {filepath: 1} - limiter._shared_pool_filepaths = [] + limiter._shared_pool_scopes = [] try: - limiter._apply_limits_with_shared_pool_claims() + limiter._apply_limits_in_shared_pool_scopes() assert lib.applied == 1 assert state.controlling_thread is current_thread() assert state.sticky is False @@ -1236,7 +1247,7 @@ def test_scoped_limit_replacing_dead_holder_is_not_sticky(): assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath) == 0 finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath, None) - CONTROLLING_THREAD_OWNERSHIP._reset() + _reset_ownership() class _OtherSharedPool: @@ -1285,7 +1296,7 @@ def force_clear(): finally: threadpoolctl._SHARED_POOL_STATES.pop(filepath_openblas, None) threadpoolctl._SHARED_POOL_STATES.pop(filepath_other, None) - CONTROLLING_THREAD_OWNERSHIP._reset() + _reset_ownership() def test_clear_controlling_thread_leaves_other_copies(): @@ -1329,7 +1340,7 @@ def force_clear(): for filepath in (lib_a.filepath, lib_b.filepath): threadpoolctl._SHARED_POOL_STATES.pop(filepath, None) threadpoolctl._SHARED_POOL_GENERATIONS.pop(filepath, None) - CONTROLLING_THREAD_OWNERSHIP._reset() + _reset_ownership() # skip test if not run in a azure pipelines job since it relies on a specific flexiblas diff --git a/threadpoolctl.py b/threadpoolctl.py index 7bb123ce..e3b84ee7 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -14,6 +14,7 @@ from __future__ import annotations +import copy import ctypes import itertools import os @@ -22,7 +23,6 @@ import warnings from abc import ABC, abstractmethod from collections.abc import Callable -from contextlib import ContextDecorator from dataclasses import dataclass, field from functools import lru_cache, wraps from threading import Lock, Thread, current_thread @@ -497,8 +497,10 @@ class _SharedPoolState: _SHARED_POOL_STATES_LOCK = Lock() # Bumped per filepath by a forced clear. A thread drops a count when its stored # generation for that filepath does not match. clear_all_controlling_threads() bumps -# every filepath of that controller class. +# every filepath of that controller class. The epoch lets a thread skip that work +# when no forced clear has happened since its last sync. _SHARED_POOL_GENERATIONS: dict[str, int] = {} +_SHARED_POOL_GENERATION_EPOCH = 0 def _get_shared_pool_state(filepath: str, controller_class: type) -> _SharedPoolState: @@ -534,7 +536,7 @@ def _iter_claimable_shared_pools(controllers): yield lib -def _release_shared_pool_claim(filepath: str): +def _release_scoped_claim(filepath: str): """Drop this thread's non-sticky claim on one shared pool. A sticky claim was taken by a direct ``set_num_threads()`` while no limiter @@ -552,7 +554,9 @@ def _release_shared_pool_claim(filepath: str): def _bump_shared_pool_generations(filepaths) -> None: """Invalidate other threads' counts for these shared pools.""" + global _SHARED_POOL_GENERATION_EPOCH with _SHARED_POOL_STATES_LOCK: + _SHARED_POOL_GENERATION_EPOCH += 1 for filepath in filepaths: current = _SHARED_POOL_GENERATIONS.get(filepath, 0) _SHARED_POOL_GENERATIONS[filepath] = current + 1 @@ -651,7 +655,6 @@ def _set_shared_pool_num_threads(self, num_threads): count is still 0, is sticky: a later context manager must not drop it. A call made while a limiter scope is active is released with that scope. """ - CONTROLLING_THREAD_OWNERSHIP._sync() this_thread = current_thread() state = self._shared_pool_state() controlling_thread = state.controlling_thread @@ -836,25 +839,24 @@ class _ControllingThreadOwnership(thread_local): def __init__(self): # threading.local runs __init__ once per thread. - self.counts: dict[str, int] = {} - self.generations: dict[str, int] = {} + self._counts: dict[str, int] = {} + self._generations: dict[str, int] = {} + self._generation_epoch = 0 def _sync(self): """Drop counts for libraries another thread has force-cleared.""" with _SHARED_POOL_STATES_LOCK: + epoch = _SHARED_POOL_GENERATION_EPOCH + if epoch == self._generation_epoch: + return generations = dict(_SHARED_POOL_GENERATIONS) - local = self.generations - stale = { - filepath - for filepath, generation in generations.items() - if local.get(filepath, 0) != generation - } - if not stale: - return - counts = self.counts - for filepath in stale: - counts.pop(filepath, None) - local[filepath] = generations[filepath] + self._generation_epoch = epoch + local = self._generations + counts = self._counts + for filepath, generation in generations.items(): + if local.get(filepath, 0) != generation: + counts.pop(filepath, None) + local[filepath] = generation def _forget_pool(self, filepath: str): """Drop this thread's count for one pool and adopt its generation. @@ -862,22 +864,13 @@ def _forget_pool(self, filepath: str): The caller has already bumped the generation when this clear is forced. Other pools stay stale until the next ``_sync``. """ - self.counts.pop(filepath, None) + self._counts.pop(filepath, None) with _SHARED_POOL_STATES_LOCK: generation = _SHARED_POOL_GENERATIONS.get(filepath, 0) - self.generations[filepath] = generation - - def _reset(self, controller_class: type | None = None): - """Forget scopes on this thread. + self._generations[filepath] = generation - With ``controller_class``, only that class. Without, every library. - """ - if controller_class is None: - with _SHARED_POOL_STATES_LOCK: - generations = dict(_SHARED_POOL_GENERATIONS) - self.counts = {} - self.generations = generations - return + def _reset(self, controller_class: type): + """Forget this thread's scopes for one controller class.""" with _SHARED_POOL_STATES_LOCK: filepaths = [ filepath @@ -888,8 +881,8 @@ def _reset(self, controller_class: type | None = None): filepath: _SHARED_POOL_GENERATIONS.get(filepath, 0) for filepath in filepaths } - counts = self.counts - local = self.generations + counts = self._counts + local = self._generations for filepath in filepaths: counts.pop(filepath, None) local[filepath] = generations[filepath] @@ -897,7 +890,7 @@ def _reset(self, controller_class: type | None = None): def _count_for(self, filepath: str) -> int: """Active scopes for one shared pool on this thread.""" self._sync() - return self.counts.get(filepath, 0) + return self._counts.get(filepath, 0) def increment(self, controller: ThreadpoolController | None = None): """Increment every shared pool of ``controller`` this thread can control. @@ -913,7 +906,7 @@ def increment(self, controller: ThreadpoolController | None = None): def _increment(self, filepath: str): """Count one more scope for a single shared pool.""" self._sync() - counts = self.counts + counts = self._counts counts[filepath] = counts.get(filepath, 0) + 1 def decrement(self, controller: ThreadpoolController | None = None): @@ -925,27 +918,29 @@ def decrement(self, controller: ThreadpoolController | None = None): """ if controller is None: controller = get_cached_controller() + self._sync() seen = set() for lib in controller.lib_controllers: - if lib.filepath in seen: + if lib.filepath in seen or lib.filepath not in self._counts: continue seen.add(lib.filepath) - self._decrement(lib.filepath) + self._decrement(lib.filepath, synced=True) - def _decrement(self, filepath: str): + def _decrement(self, filepath: str, *, synced: bool = False): """Count one fewer scope for a single shared pool. Reaching zero releases that pool's non-sticky claim. """ - self._sync() - counts = self.counts + if not synced: + self._sync() + counts = self._counts current = counts.get(filepath, 0) if current <= 0: return current -= 1 if current == 0: counts.pop(filepath, None) - _release_shared_pool_claim(filepath) + _release_scoped_claim(filepath) else: counts[filepath] = current @@ -1338,11 +1333,11 @@ def __init__(self, controller, *, limits=None, user_api=None): self._controller = controller self._limits, self._user_api = self._check_params(limits, user_api) self._original_info = self._controller.info() - self._shared_pool_filepaths: list[str] = [] + self._shared_pool_scopes: list[tuple[str, _SharedPoolState]] = [] # A plain call never reaches __exit__, so this increment is what makes the claim # stick until clear_controlling_thread(). A `with` block decrements the same # increment on the way out. - self._apply_limits_with_shared_pool_claims() + self._apply_limits_in_shared_pool_scopes() def __enter__(self): return self @@ -1487,7 +1482,7 @@ def _enter_shared_pool_scopes(self): one copy does not pin another. A finished holder does not block this scope: counting first keeps the replacement claim non-sticky. """ - self._shared_pool_filepaths = [] + self._shared_pool_scopes = [] targets = [ lib for lib in self._controller.lib_controllers @@ -1495,32 +1490,26 @@ def _enter_shared_pool_scopes(self): ] for lib in _iter_claimable_shared_pools(targets): CONTROLLING_THREAD_OWNERSHIP._increment(lib.filepath) - self._shared_pool_filepaths.append(lib.filepath) + self._shared_pool_scopes.append((lib.filepath, lib._shared_pool_state())) def _exit_shared_pool_scopes(self): - filepaths = self._shared_pool_filepaths - self._shared_pool_filepaths = [] - for filepath in filepaths: + scopes = self._shared_pool_scopes + self._shared_pool_scopes = [] + for filepath, _state in scopes: CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) def _drop_unheld_shared_pool_scopes(self): """Drop counts for libraries another thread won the race for.""" this = current_thread() kept = [] - for filepath in self._shared_pool_filepaths: - with _SHARED_POOL_STATES_LOCK: - state = _SHARED_POOL_STATES.get(filepath) - held = False - if state is not None: - with state.lock: - held = state.controlling_thread is this - if held: - kept.append(filepath) + for filepath, state in self._shared_pool_scopes: + if state.controlling_thread is this: + kept.append((filepath, state)) else: CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) - self._shared_pool_filepaths = kept + self._shared_pool_scopes = kept - def _apply_limits_with_shared_pool_claims(self): + def _apply_limits_in_shared_pool_scopes(self): self._enter_shared_pool_scopes() try: self._set_threadpool_limits() @@ -1540,34 +1529,37 @@ def _set_threadpool_limits(self): lib_controller.set_num_threads(num_threads) -class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter, ContextDecorator): +class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter): """Same as _ThreadpoolLimiter but to be used as a decorator""" def __init__(self, controller, *, limits=None, user_api=None): self._limits, self._user_api = self._check_params(limits, user_api) self._controller = controller - self._shared_pool_filepaths = [] + self._shared_pool_scopes: list[tuple[str, _SharedPoolState]] = [] def __enter__(self): # Set the limits here and not in __init__ so they apply when the decorated # function is called, not when the decorator is created. Increment before # setting so this scope's claim is not sticky. self._original_info = self._controller.info() - self._apply_limits_with_shared_pool_claims() + self._apply_limits_in_shared_pool_scopes() return self - def _recreate_cm(self): - """Return a fresh limiter for this call. + def _fresh_limiter(self): + """Return a limiter that does not share this call's scope list.""" + fresh = copy.copy(self) + fresh._shared_pool_scopes = [] + return fresh - Recursive and concurrent calls of the same wrapped function must not - share ``_shared_pool_filepaths`` or ``_original_info``. - """ - recreated = object.__new__(type(self)) - recreated._controller = self._controller - recreated._limits = self._limits - recreated._user_api = self._user_api - recreated._shared_pool_filepaths = [] - return recreated + def __call__(self, func): + """Decorate ``func`` so each call enters a fresh limiter.""" + + @wraps(func) + def inner(*args, **kwargs): + with self._fresh_limiter(): + return func(*args, **kwargs) + + return inner @_format_docstring( From ec5885f6651ef97b7ca8594c432f4cf6d7c4981c Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 18:38:41 +0100 Subject: [PATCH 17/21] FIX Read the controlling thread under its lock when dropping unheld scopes A limiter that lost the claim race dropped its count after an unlocked read of the winner. On free-threaded CPython that read could miss the winning write and keep a count for a pool it does not hold, making the next claim sticky. The read now takes the pool lock. The iterator also yields the pool state it already looked up, and the decrement no longer takes a flag, since the generation epoch makes the extra sync cheap. Co-authored-by: Cursor --- threadpoolctl.py | 34 +++++++++++++++++++++------------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/threadpoolctl.py b/threadpoolctl.py index e3b84ee7..3339f850 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -518,9 +518,11 @@ def _other_live_thread_holds(holder, this) -> bool: def _iter_claimable_shared_pools(controllers): - """Yield shared-pool libraries in ``controllers`` this thread can claim. + """Yield ``(library, state)`` for shared pools this thread can claim. One library per filepath. A pool another live thread already holds is skipped. + The state is the one ``library._shared_pool_state()`` just returned, so the + caller does not look it up again. """ this = current_thread() seen = set() @@ -530,10 +532,10 @@ def _iter_claimable_shared_pools(controllers): if lib.filepath in seen: continue seen.add(lib.filepath) - holder = lib._shared_pool_state().controlling_thread - if _other_live_thread_holds(holder, this): + state = lib._shared_pool_state() + if _other_live_thread_holds(state.controlling_thread, this): continue - yield lib + yield lib, state def _release_scoped_claim(filepath: str): @@ -900,7 +902,7 @@ def increment(self, controller: ThreadpoolController | None = None): """ if controller is None: controller = get_cached_controller() - for lib in _iter_claimable_shared_pools(controller.lib_controllers): + for lib, _state in _iter_claimable_shared_pools(controller.lib_controllers): self._increment(lib.filepath) def _increment(self, filepath: str): @@ -924,15 +926,14 @@ def decrement(self, controller: ThreadpoolController | None = None): if lib.filepath in seen or lib.filepath not in self._counts: continue seen.add(lib.filepath) - self._decrement(lib.filepath, synced=True) + self._decrement(lib.filepath) - def _decrement(self, filepath: str, *, synced: bool = False): + def _decrement(self, filepath: str): """Count one fewer scope for a single shared pool. Reaching zero releases that pool's non-sticky claim. """ - if not synced: - self._sync() + self._sync() counts = self._counts current = counts.get(filepath, 0) if current <= 0: @@ -1488,9 +1489,9 @@ def _enter_shared_pool_scopes(self): for lib in self._controller.lib_controllers if self._num_threads_for(lib) is not None ] - for lib in _iter_claimable_shared_pools(targets): + for lib, state in _iter_claimable_shared_pools(targets): CONTROLLING_THREAD_OWNERSHIP._increment(lib.filepath) - self._shared_pool_scopes.append((lib.filepath, lib._shared_pool_state())) + self._shared_pool_scopes.append((lib.filepath, state)) def _exit_shared_pool_scopes(self): scopes = self._shared_pool_scopes @@ -1499,11 +1500,18 @@ def _exit_shared_pool_scopes(self): CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) def _drop_unheld_shared_pool_scopes(self): - """Drop counts for libraries another thread won the race for.""" + """Drop counts for libraries another thread won the race for. + + Read ``controlling_thread`` under the pool lock. This thread's own write + is visible without it, but the winning write from another thread is not. + A stale empty holder would keep a count and make the next claim non-sticky. + """ this = current_thread() kept = [] for filepath, state in self._shared_pool_scopes: - if state.controlling_thread is this: + with state.lock: + held = state.controlling_thread is this + if held: kept.append((filepath, state)) else: CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) From 102ba246356d3eabae7b9e420e1acc70a0787366 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 19:03:51 +0100 Subject: [PATCH 18/21] FIX Ignore the libomp/libiomp warning while the test fixture scans libraries The autouse fixture builds a ThreadpoolController, which warns when LLVM and Intel OpenMP are both loaded. pytest runs with -W error, so that warning aborted test_libomp_libiomp_warning during setup, before the test could record it. Co-authored-by: Cursor --- conftest.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/conftest.py b/conftest.py index d334ac42..0fbac7b1 100644 --- a/conftest.py +++ b/conftest.py @@ -1,3 +1,5 @@ +import warnings + import pytest collect_ignore = [ @@ -19,7 +21,16 @@ def _ensure_openblas_thread_count_is_restorable(): """ from threadpoolctl import OpenBLASController, ThreadpoolController - controller = ThreadpoolController().select(internal_api="openblas") + # Building a controller warns when libomp and libiomp are both loaded. + # That warning is the subject of test_libomp_libiomp_warning, and pytest + # runs with -W error, so the fixture must not emit it. + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"\s*Found Intel OpenMP", + category=RuntimeWarning, + ) + controller = ThreadpoolController().select(internal_api="openblas") if controller: for lib in controller.lib_controllers: if lib.threading_layer == "openmp" or lib.get_num_threads() != 0: From 7eed92d392126cf1daaafcf7355ccc72fe764104 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 20:04:01 +0100 Subject: [PATCH 19/21] MNT Assert the shared-pool holder in OpenBLAS ownership tests Tests that hand the claim to another thread now record and assert which thread holds it, not only the thread count. The recursive-wrap test now runs through threadpool_limits.wrap on a real OpenBLAS library. Co-authored-by: Cursor --- tests/test_threadpoolctl.py | 80 ++++++++++++++++++++----------------- 1 file changed, 43 insertions(+), 37 deletions(-) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 5485946f..fd364bbf 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -728,6 +728,7 @@ def run_in_thread(f, *args, **kwargs): thread = Thread(target=lambda: f(*args, **kwargs)) thread.start() thread.join() + return thread def test_openblas_non_openmp_restricted_to_controlling_thread(): @@ -774,6 +775,7 @@ def check_openblas_non_openmp_restricted_to_controlling_thread(openblas_lib): # This thread can change num threads: openblas_lib.set_num_threads(1) assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._shared_pool_state().controlling_thread is current_thread() # Another thread's attempts are no-op: result = [] @@ -785,10 +787,12 @@ def another_thread(): run_in_thread(another_thread) assert result[-1] == 1 assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._shared_pool_state().controlling_thread is current_thread() # This thread can still make changes: openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == orig_num_threads + assert openblas_lib._shared_pool_state().controlling_thread is current_thread() # Another thread clearing the controlling thread has no effect: def clear_and_set(): @@ -797,15 +801,18 @@ def clear_and_set(): run_in_thread(clear_and_set) assert openblas_lib.get_num_threads() == orig_num_threads + assert openblas_lib._shared_pool_state().controlling_thread is current_thread() # The controlling thread can clear, and then another thread can set: openblas_lib.clear_controlling_thread() - run_in_thread(openblas_lib.set_num_threads, 1) + worker = run_in_thread(openblas_lib.set_num_threads, 1) assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._shared_pool_state().controlling_thread is worker # Dead threads lose control: openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == orig_num_threads + assert openblas_lib._shared_pool_state().controlling_thread is current_thread() # Release control: openblas_lib.clear_controlling_thread() @@ -824,15 +831,18 @@ def set_and_exit_after_lock_released(): t.start() was_set.acquire() assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._shared_pool_state().controlling_thread is t # This thread can't set since it's not the controller: openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._shared_pool_state().controlling_thread is t # But when forced clear happens, we can set: openblas_lib.clear_controlling_thread(force=True) openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == orig_num_threads + assert openblas_lib._shared_pool_state().controlling_thread is current_thread() # Clean up finish.release() @@ -900,6 +910,35 @@ def outer(): assert _openblas_controlling_thread() is None +def test_recursive_wrap_releases_the_shared_pool_claim(request): + # One wrapped function calling itself must keep the limit between the inner + # return and the outer return, then restore the original count and drop the + # controlling thread. Another thread can then set the count. + libs, lib, orig = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + @threadpool_limits.wrap(limits=1, user_api="blas") + def recurse(depth): + assert lib.get_num_threads() == 1 + if depth: + recurse(depth - 1) + assert lib.get_num_threads() == 1 + + recurse(1) + assert lib.get_num_threads() == orig + + observed = {} + + def other_thread(): + lib.set_num_threads(1) + observed["threads"] = lib.get_num_threads() + observed["holder"] = lib._shared_pool_state().controlling_thread + + worker = run_in_thread(other_thread) + assert observed["threads"] == 1 + assert observed["holder"] is worker + + def test_openblas_controlling_thread_per_pool_counter(request): # The per-pool counter has to be incremented before set_num_threads. That # makes the claim scoped, so the last decrement releases it. @@ -938,39 +977,6 @@ def _real_set_num_threads(self, num_threads): self.applied = num_threads -def test_recursive_wrap_releases_the_shared_pool_claim(): - # Each call of a wrapped function must release only the claim that call took. - lib = _TestSharedPoolController("threadpoolctl-test-recursive-wrap") - state = lib._shared_pool_state() - - class _Controller: - def __init__(self): - self.lib_controllers = [lib] - - def info(self): - return [{"num_threads": 4}] - - wrapped = threadpoolctl._ThreadpoolLimiterDecorator( - _Controller(), limits=1, user_api="blas" - ) - - @wrapped - def recurse(depth): - assert state.controlling_thread is current_thread() - if depth: - recurse(depth - 1) - assert state.controlling_thread is current_thread() - - try: - recurse(1) - assert state.controlling_thread is None - assert CONTROLLING_THREAD_OWNERSHIP._count_for(lib.filepath) == 0 - finally: - threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) - threadpoolctl._SHARED_POOL_GENERATIONS.pop(lib.filepath, None) - _reset_ownership() - - def test_bare_increment_decrement_covers_every_controllable_pool(): # Every loaded shared pool this thread can control, and not a pool another # live thread already holds. @@ -1136,6 +1142,7 @@ def other_thread(): run_in_thread(other_thread) assert observed["threads"] == 1 + assert state.controlling_thread is current_thread() def test_restore_original_limits_keeps_openblas_controlling_thread(request): @@ -1173,12 +1180,11 @@ def other_thread(): observed["holder"] = second._shared_pool_state().controlling_thread observed["threads"] = second.get_num_threads() - run_in_thread(other_thread) + worker = run_in_thread(other_thread) assert first.get_num_threads() == 1 assert first._shared_pool_state().controlling_thread is current_thread() assert observed["threads"] == 1 - assert observed["holder"] is not None - assert observed["holder"] is not current_thread() + assert observed["holder"] is worker def test_openblas_scoped_exit_releases_only_that_library(): From 4f3b7a3ab292556979fb1877b098870815ce7d5d Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 20:15:15 +0100 Subject: [PATCH 20/21] DOC Explain the OpenBLAS controlling thread in terms of loaded libraries Say loaded OpenBLAS library instead of copy, and note that NumPy and SciPy can each load their own. Describe the automatic release for context managers and decorators and point to the OpenBLASController docstring for details. Co-authored-by: Cursor --- CHANGES.md | 2 +- README.md | 19 +++++++++++-------- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index f562fd8b..a88cb259 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -12,7 +12,7 @@ controlling thread. A plain ``threadpool_limits(...)`` call or a direct ``set_num_threads()`` keeps that role until ``clear_controlling_thread()`` on that library. ``OpenBLASController.clear_all_controlling_threads()`` drops - every loaded copy. Context managers and decorators release it when the + every loaded OpenBLAS library. Context managers and decorators release it when the outermost scope that took the claim exits. A thread that has finished no longer holds the claim; the next ``set_num_threads()`` can take it. The ``CONTROLLING_THREAD_OWNERSHIP`` singleton can be used to manually manage the diff --git a/README.md b/README.md index ec6d98b0..e0c451fc 100644 --- a/README.md +++ b/README.md @@ -274,14 +274,17 @@ OpenBLAS+pthreads) and in each worker thread. > **Why you need both:** When setting a limit on OpenBLAS with pthreads or the > Windows threading layer, setting the limit affects the whole process, because -> each loaded copy uses one process-wide shared thread pool. (This in contrast -> to some other BLAS/OpenMP libraries, or OpenBLAS backed by OpenMP, that allow -> setting the limit per-thread.) Having multiple Python threads set limits on -> the same copy makes no sense, since they would be arguing over that copy's -> pool. Each copy therefore has its own controlling thread, the first thread to -> call ``set_num_threads()`` on it. Other threads do nothing. Once that thread -> has finished, or after ``clear_controlling_thread()`` on that copy, the next -> call can take the claim. +> each loaded OpenBLAS library uses one process-wide shared thread pool. (This in +> contrast to some other BLAS/OpenMP libraries, or OpenBLAS backed by OpenMP, +> that allow setting the limit per-thread.) A process can load more than one +> OpenBLAS library, for example one bundled with NumPy and another bundled with +> SciPy, and each has its own pool. Having multiple Python threads set limits on +> the same library makes no sense, since they would be arguing over its pool. +> So each loaded OpenBLAS library has its own controlling thread: the first +> thread to set thread limits on it. Other threads do nothing. The controlling +> thread is released automatically when the ``with`` block or decorated function +> that took it exits, and once that thread has finished. See the docstring of +> ``OpenBLASController`` for more details. > > For many other libraries, for example OpenMP, the limit is set per thread. As > a result, you want to set limits both ways in generic code that doesn't know From 52cac5c63203098c6db20f04652fa96570dfb1c5 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 30 Oct 2026 20:17:49 +0100 Subject: [PATCH 21/21] Apply suggestion from @ogrisel --- CHANGES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index a88cb259..6ae4eb21 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -19,7 +19,7 @@ scope of ``set_num_threads`` calls. https://github.com/joblib/threadpoolctl/pull/275 https://github.com/joblib/threadpoolctl/pull/289 - https://github.com/joblib/threadpoolctl/pull/XXX + https://github.com/joblib/threadpoolctl/pull/290 - Add a new API to get a cached `ThreadpoolController` instance (`get_cached_controller()`), as an optimization for cases where limits are set