diff --git a/CHANGES.md b/CHANGES.md index 7317c063..6ae4eb21 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -8,12 +8,18 @@ - 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 ``clear_controlling_thread()`` on + that library. ``OpenBLASController.clear_all_controlling_threads()`` drops + 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 + 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/290 - 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/README.md b/README.md index fcc26221..e0c451fc 100644 --- a/README.md +++ b/README.md @@ -274,13 +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 -> 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 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 diff --git a/conftest.py b/conftest.py index 101040c7..0fbac7b1 100644 --- a/conftest.py +++ b/conftest.py @@ -1,3 +1,5 @@ +import warnings + import pytest collect_ignore = [ @@ -9,11 +11,41 @@ ] +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 + + # 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: + 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_all_controlling_threads(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() 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..fd364bbf 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. @@ -717,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(): @@ -763,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 = [] @@ -774,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(): @@ -786,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() @@ -813,33 +831,62 @@ 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() 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._shared_pool_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_all_controlling_threads(force=True) + for lib, orig in originals: + if orig is not None: + lib.set_num_threads(orig) + OpenBLASController.clear_all_controlling_threads(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 +896,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() @@ -864,20 +910,443 @@ def outer(): assert _openblas_controlling_thread() is None -def test_openblas_controlling_thread_manual_increment_decrement_api(request): - limiter = ThreadpoolController().limit() - request.addfinalizer(limiter.restore_original_limits) +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. + 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)) + + filepath = lib.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 + CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) + assert lib._shared_pool_state().controlling_thread is current_thread() + CONTROLLING_THREAD_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, filepath): + self.filepath = filepath + self.prefix = filepath + 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(): + # 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() + 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) + 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 + + CONTROLLING_THREAD_OWNERSHIP.increment() + CONTROLLING_THREAD_OWNERSHIP.decrement() + assert free_a._shared_pool_state().controlling_thread is current_thread() + 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 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: + holder.join() + for lib in (free_a, free_b, held): + controller.lib_controllers.remove(lib) + threadpoolctl._SHARED_POOL_STATES.pop(lib.filepath, None) + _reset_ownership() + + +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: + 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 + + CONTROLLING_THREAD_OWNERSHIP.increment() + assert CONTROLLING_THREAD_OWNERSHIP._count_for(chosen.filepath) == 1 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(other.filepath) == 1 + + CONTROLLING_THREAD_OWNERSHIP.decrement(controller) + assert chosen._shared_pool_state().controlling_thread is None + assert CONTROLLING_THREAD_OWNERSHIP._count_for(chosen.filepath) == 0 + assert CONTROLLING_THREAD_OWNERSHIP._count_for(other.filepath) == 1 + + 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) + _reset_ownership() + + +def _require_resizable_non_openmp_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_non_openmp_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._shared_pool_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._shared_pool_state().controlling_thread is current_thread() + + assert lib.get_num_threads() == 1 + 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._shared_pool_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_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() + assert lib._shared_pool_state().controlling_thread is None + + +def test_direct_set_num_threads_is_not_released_by_limiter_scopes(request): + libs, lib, orig = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + lib.set_num_threads(1) + state = lib._shared_pool_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 + assert state.controlling_thread is current_thread() + + +def test_restore_original_limits_keeps_openblas_controlling_thread(request): + libs, lib, orig = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + limiter = threadpool_limits(limits=1) + assert lib.get_num_threads() == 1 + assert lib._shared_pool_state().controlling_thread is current_thread() + + limiter.restore_original_limits() + assert lib.get_num_threads() == orig + assert lib._shared_pool_state().controlling_thread is current_thread() + + OpenBLASController.clear_all_controlling_threads() + assert lib._shared_pool_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._shared_pool_state().controlling_thread + observed["threads"] = second.get_num_threads() + + 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 worker + + +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. + filepath_a = "threadpoolctl-test-openblas-a" + filepath_b = "threadpoolctl-test-openblas-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 + + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_b) + state_b.controlling_thread = current_thread() + state_b.sticky = CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_b) == 0 + assert state_b.sticky is False + + 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 + + # 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(filepath_a) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_b) + state_b.controlling_thread = current_thread() + state_b.sticky = False + 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 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) + _reset_ownership() + + +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. + filepath = "threadpoolctl-test-dead-holder" + dead = Thread(target=lambda: None) + dead.start() + dead.join() + 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 = {filepath: 1} + limiter._shared_pool_scopes = [] + try: + limiter._apply_limits_in_shared_pool_scopes() + assert lib.applied == 1 + assert state.controlling_thread is current_thread() + assert state.sticky is False + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath) == 1 + limiter._exit_shared_pool_scopes() + assert state.controlling_thread is None + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath) == 0 + finally: + threadpoolctl._SHARED_POOL_STATES.pop(filepath, None) + _reset_ownership() + + +class _OtherSharedPool: + """Stand-in for a future library with its own shared pool.""" + + +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._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 + state_other.controlling_thread = current_thread() + state_other.sticky = True + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_openblas) + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_other) + + 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 + assert state_other.controlling_thread is current_thread() + assert state_other.sticky is True + assert CONTROLLING_THREAD_OWNERSHIP._count_for(filepath_other) == 1 + + state_openblas.controlling_thread = current_thread() + state_openblas.sticky = True + CONTROLLING_THREAD_OWNERSHIP._increment(filepath_openblas) + + def force_clear(): + OpenBLASController.clear_all_controlling_threads(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 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) + _reset_ownership() + + +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) + _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 756168cb..3339f850 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 @@ -59,6 +59,7 @@ __version__ = "3.8.0.dev0" __all__ = [ + "CONTROLLING_THREAD_OWNERSHIP", "LibController", "ThreadpoolController", "get_cached_controller", @@ -172,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 @@ -198,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): @@ -308,6 +309,16 @@ 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. + Controllers with a process-wide pool inherit ``_SharedPoolControllerMixin``, + which reports True. Override when only some builds share a pool. + """ + return False + # Internal marker for missing cache lookups: _MISSING = object() @@ -462,7 +473,212 @@ def wrapper(self): _CDLL_CACHE = _CDLLCache() -class OpenBLASController(LibController): +@dataclass +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) + + +# 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. +# +# 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 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. 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: + with _SHARED_POOL_STATES_LOCK: + state = _SHARED_POOL_STATES.get(filepath) + if state is None: + state = _SharedPoolState(controller_class=controller_class) + _SHARED_POOL_STATES[filepath] = state + return state + + +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 _iter_claimable_shared_pools(controllers): + """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() + for lib in controllers: + if not lib._has_shared_pool(): + continue + if lib.filepath in seen: + continue + seen.add(lib.filepath) + state = lib._shared_pool_state() + if _other_live_thread_holds(state.controlling_thread, this): + continue + yield lib, state + + +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 + scope was active for that library. Scoped exits must leave it in place. + """ + with _SHARED_POOL_STATES_LOCK: + state = _SHARED_POOL_STATES.get(filepath) + 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 _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 + + +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 + + +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 this library's claim. + ``clear_all_controlling_threads`` drops every library of this controller class. + """ + + def _has_shared_pool(self): + """Whether this loaded library currently uses a process-wide pool.""" + return True + + 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. 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``. + """ + 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_pool(self.filepath) + + @classmethod + 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 + 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. 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``. + """ + 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. + + 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. + """ + 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 = CONTROLLING_THREAD_OWNERSHIP._count_for(self.filepath) + 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 + + +class OpenBLASController(_SharedPoolControllerMixin, LibController): """Controller class for OpenBLAS. When OpenBLAS uses the OpenMP threading layer, each thread that calls into @@ -470,22 +686,31 @@ 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 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()`` 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 + entered while a plain call or a direct ``set_num_threads()`` is still + holding the claim, does not release it. """ user_api = "blas" @@ -529,54 +754,14 @@ def get_num_threads(self): return get_num_threads_func() return None - class _SHARED_STATE: - """Shared state across instances.""" - - controlling_thread: Thread | None = None - lock = Lock() - - @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. - - 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``. - """ - with cls._SHARED_STATE.lock: - if force or cls._SHARED_STATE.controlling_thread is current_thread(): - cls._SHARED_STATE.controlling_thread = None + def _has_shared_pool(self): + return self.threading_layer != "openmp" 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: - this_thread = current_thread() - controlling_thread = self._SHARED_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: - return self.set_num_threads(num_threads) - elif controlling_thread is this_thread: - # This thread controls OpenBLAS: - return self._real_set_num_threads(num_threads) - else: - # Some other thread controlls OpenBLAS: do nothing. - 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 @@ -638,25 +823,127 @@ def _get_architecture(self): class _ControllingThreadOwnership(thread_local): - """Count how many paired limit-unlimits have modified OpenBLAS. + """Per-pool count of limiter scopes on this 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 ``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 of + one library drops that library's counts. ``clear_all_controlling_threads()`` + drops counts for that controller class. + """ - Clear controlling thread if it hits zero. + def __init__(self): + # threading.local runs __init__ once per thread. + 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) + 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. + + 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) + self._generations[filepath] = generation + + def _reset(self, controller_class: type): + """Forget this thread's scopes for one controller class.""" + with _SHARED_POOL_STATES_LOCK: + 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) + local[filepath] = generations[filepath] - In the future this could potentialy also apply to non-OpenBLAS libraries - that have a shared process pool. - """ + def _count_for(self, filepath: str) -> int: + """Active scopes for one shared pool on this thread.""" + self._sync() + return self._counts.get(filepath, 0) + + def increment(self, controller: ThreadpoolController | None = None): + """Increment every shared pool of ``controller`` this thread can control. - count = 0 + ``controller`` defaults to ``get_cached_controller()``. Call this + before ``set_num_threads`` so the claim stays scoped. + """ + if controller is None: + controller = get_cached_controller() + for lib, _state 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[filepath] = counts.get(filepath, 0) + 1 + + def decrement(self, controller: ThreadpoolController | None = None): + """Decrement pools of ``controller`` that this thread currently counts. + + ``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() + seen = set() + for lib in controller.lib_controllers: + if lib.filepath in seen or lib.filepath not in self._counts: + continue + seen.add(lib.filepath) + self._decrement(lib.filepath) - def increment(self): - """Increment ownership.""" - self.count += 1 + def _decrement(self, filepath: str): + """Count one fewer scope for a single shared pool. - def decrement(self): - """Decrement ownership.""" - self.count -= 1 - if self.count <= 0: - OpenBLASController.clear_controlling_thread() + Reaching zero releases that pool's non-sticky claim. + """ + self._sync() + counts = self._counts + current = counts.get(filepath, 0) + if current <= 0: + return + current -= 1 + if current == 0: + counts.pop(filepath, None) + _release_scoped_claim(filepath) + else: + counts[filepath] = current CONTROLLING_THREAD_OWNERSHIP = _ControllingThreadOwnership() @@ -1047,15 +1334,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._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_in_shared_pool_scopes() 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._exit_shared_pool_scopes() @classmethod def wrap(cls, controller, *, limits=None, user_api=None): @@ -1166,44 +1458,117 @@ 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: + 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 _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 + 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_scopes = [] + targets = [ + lib + for lib in self._controller.lib_controllers + if self._num_threads_for(lib) is not None + ] + for lib, state in _iter_claimable_shared_pools(targets): + CONTROLLING_THREAD_OWNERSHIP._increment(lib.filepath) + self._shared_pool_scopes.append((lib.filepath, state)) + + def _exit_shared_pool_scopes(self): + 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. + + 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: + with state.lock: + held = state.controlling_thread is this + if held: + kept.append((filepath, state)) + else: + CONTROLLING_THREAD_OWNERSHIP._decrement(filepath) + self._shared_pool_scopes = kept + + def _apply_limits_in_shared_pool_scopes(self): + self._enter_shared_pool_scopes() + try: + self._set_threadpool_limits() + except Exception: + self._exit_shared_pool_scopes() + raise + self._drop_unheld_shared_pool_scopes() + 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) -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_scopes: list[tuple[str, _SharedPoolState]] = [] 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_in_shared_pool_scopes() return self + 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 + + 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( USER_APIS=", ".join(f'"{api}"' for api in _ALL_USER_APIS),