diff --git a/CHANGES.md b/CHANGES.md index 71e96c18..7f013581 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -5,6 +5,15 @@ quite a bit of overhead. https://github.com/joblib/threadpoolctl/pull/250 +- 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. + https://github.com/joblib/threadpoolctl/pull/275 + - Add a new API to get a cached `ThreadpoolController` instance (`get_cached_controller()`), as an optimization for cases where limits are set frequently. See the docstring for caveats. diff --git a/README.md b/README.md index e915e14a..fcc26221 100644 --- a/README.md +++ b/README.md @@ -269,6 +269,24 @@ libraries) that can take some time. To prevent the performance cost of doing this work in all the threads, you can reuse a `ThreadpoolController` object across the threads. +When setting limits, you want to set limits both in the main thread (in case you're using +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. +> +> 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 +> the underlying BLAS library: in main thread for the process-wide limits, and +> per-thread for thread-specific limits like OpenMP. + ```python from threadpoolctl import ThreadpoolController @@ -279,11 +297,10 @@ from threadpoolctl import ThreadpoolController controller = ThreadpoolController() with ( - # This top-level limiter doesn't actually change the limits initially; it - # is there to ensure the limits are reset _after_ the Python thread pool is - # done. This is necessary because some underlying limiting APIs operate on - # a process-wide basis. - controller.limit(), + # Limit on main thread. This is necessary because some underlying limiting + # APIs operate on a process-wide basis. Because it's a context manager, it + # will restore the limits after the ``with`` is done. + controller.limit(limits=1), # Make sure each Python worker thread also calls threadpool_limits(). If # you're using another thread pool class, you will need to do so some other # way. @@ -295,7 +312,7 @@ with ( # Later... controller = ThreadpoolController() with ( - controller.limit(), + controller.limit(limits=2), ThreadPoolExecutor(4, initializer=lambda: controller.limit(limits=2)) as pool, ): # ... run some BLAS-using code in the thread pool ... @@ -309,7 +326,7 @@ from threadpoolctl import ThreadpoolController controller = ThreadpoolController() try: - limiter = controller.limit() + limiter = controller.limit(limits=1) with ThreadPoolExecutor(4, initializer=lambda: controller.limit(limits=1)) as pool: # ... run some BLAS-using code in the thread pool ... pool.map(somefunc, someargs) @@ -361,6 +378,22 @@ with CONTROLLER.limit(limits=2): ## Usage: Additional APIs and details +### 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: + +* 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 + have a Python thread pool calling BLAS, and the number of Python threads is + close to the number of available cores. +* A limit equal to the number of available cores. This is a good setting when + you are calling BLAS from the main thread and there are no Python threads. + +While `threadpoolctl` does not limit you to these two values, you may wish to do +so yourself in your code. + ### Switching the FlexiBLAS backend `FlexiBLAS` is a BLAS wrapper for which the BLAS backend can be switched at runtime. diff --git a/conftest.py b/conftest.py index 451cff36..101040c7 100644 --- a/conftest.py +++ b/conftest.py @@ -1,3 +1,5 @@ +import pytest + collect_ignore = [ "tests/empirical_scope_observation.py", "tests/_openmp_test_helper.py", @@ -5,3 +7,13 @@ "tests/_dl_iterate_phdr_deadlock.py", "tests/_get_cached_controller.py", ] + + +@pytest.fixture(autouse=True) +def reset_openblas_controlling_thread(): + from threadpoolctl import OpenBLASController + + try: + yield + finally: + OpenBLASController.clear_controlling_thread(force=True) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index ae7f4b91..080cc244 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -7,6 +7,7 @@ import shutil import subprocess import sys +from threading import Thread, current_thread import pytest @@ -14,7 +15,9 @@ from threadpoolctl import ( _ALL_PREFIXES, _ALL_USER_APIS, + CONTROLLING_THREAD_OWNERSHIP, LibController, + OpenBLASController, ThreadpoolController, _determine_thread_limit_scope, get_cached_controller, @@ -719,6 +722,143 @@ def test_openblas_threading_layer(): assert threading_layer in expected_openblas_threading_layers +def run_in_thread(f, *args, **kwargs): + thread = Thread(target=lambda: f(*args, **kwargs)) + thread.start() + thread.join() + + +def test_openblas_non_openmp_restricted_to_controlling_thread(): + """Only one thread can set number of threads for non-OpenMP OpenBLAS.""" + openblas_controller = ThreadpoolController().select(internal_api="openblas") + + if not (openblas_controller): + pytest.skip("requires OpenBLAS.") + + for openblas_lib in openblas_controller.lib_controllers: + if openblas_lib.get_num_threads() > 1: + if openblas_lib.threading_layer == "openmp": + check_openblas_openmp_no_controlling_thread(openblas_lib) + else: + check_openblas_non_openmp_restricted_to_controlling_thread(openblas_lib) + + +def check_openblas_openmp_no_controlling_thread(openblas_lib): + orig_num_threads = openblas_lib.get_num_threads() + assert orig_num_threads > 1 + + # This thread can change num threads: + openblas_lib.set_num_threads(1) + assert openblas_lib.get_num_threads() == 1 + + # Another thread's attempts change _that_ thread: + result = [] + + def another_thread(): + openblas_lib.set_num_threads(orig_num_threads) + result.append(openblas_lib.get_num_threads()) + + run_in_thread(another_thread) + assert result[-1] == orig_num_threads + + # This thread is unchanged: + assert openblas_lib.get_num_threads() == 1 + + +def check_openblas_non_openmp_restricted_to_controlling_thread(openblas_lib): + orig_num_threads = openblas_lib.get_num_threads() + assert orig_num_threads > 1 + + # This thread can change num threads: + openblas_lib.set_num_threads(1) + assert openblas_lib.get_num_threads() == 1 + + # Another thread's attempts are no-op: + result = [] + + def another_thread(): + openblas_lib.set_num_threads(orig_num_threads) + result.append(openblas_lib.get_num_threads()) + + run_in_thread(another_thread) + assert result[-1] == 1 + assert openblas_lib.get_num_threads() == 1 + + # This thread can still make changes: + openblas_lib.set_num_threads(orig_num_threads) + assert openblas_lib.get_num_threads() == orig_num_threads + + # Another thread clearing the controlling thread has no effect: + def clear_and_set(): + openblas_lib.clear_controlling_thread() + openblas_lib.set_num_threads(1) + + run_in_thread(clear_and_set) + assert openblas_lib.get_num_threads() == orig_num_threads + + # 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) + assert openblas_lib.get_num_threads() == 1 + + # clear_controlling_thread(force=True) allows any thread to clear: + openblas_lib.clear_controlling_thread(force=True) + openblas_lib.set_num_threads(orig_num_threads) + assert openblas_lib.get_num_threads() == orig_num_threads + + +def _openblas_controlling_thread(): + openblas_controller = ThreadpoolController().select(internal_api="openblas") + + if not (openblas_controller): + pytest.skip("requires OpenBLAS.") + if {lc.threading_layer for lc in openblas_controller.lib_controllers} == {"openmp"}: + pytest.skip("require non-OpenMP OpenBLAS.") + return OpenBLASController._SHARED_STATE.controlling_thread + + +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): + assert _openblas_controlling_thread() is current_thread() + assert _openblas_controlling_thread() is current_thread() + assert _openblas_controlling_thread() is None + + +def test_openblas_controlling_thread_cleared_on_decorator_exit(): + + @threadpool_limits.wrap(limits=1) + def inner(): + assert _openblas_controlling_thread() is current_thread() + + @threadpool_limits.wrap(limits=1) + def outer(): + assert _openblas_controlling_thread() is current_thread() + inner() + assert _openblas_controlling_thread() is current_thread() + + 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) + + 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 + + # 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 d228c002..47bc8e15 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -25,7 +25,8 @@ from contextlib import ContextDecorator from dataclasses import dataclass, field from functools import lru_cache, wraps -from threading import Lock, Thread +from threading import Lock, Thread, current_thread +from threading import local as thread_local from typing import Literal, TypeVar, final # ctypes.util is not imported on Linux: on CPython 3.14 it allocates a @@ -167,6 +168,7 @@ def get_and_set() -> None: thread = Thread(target=get_and_set) thread.start() thread.join() + OpenBLASController.clear_controlling_thread(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 @@ -192,6 +194,7 @@ def get_and_set() -> None: return "unknown" finally: set_n_threads(previous) + OpenBLASController.clear_controlling_thread(force=True) class LibController(ABC): @@ -456,7 +459,30 @@ def wrapper(self): class OpenBLASController(LibController): - """Controller class for OpenBLAS""" + """Controller class for OpenBLAS. + + When OpenBLAS uses the OpenMP threading layer, each thread that calls into + OpenBLAS gets its own OpenBLAS thread pool, and correspondingly the limit + 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 + process-wide shared thread pool, so having multiple Python threads change + the limit at the same time makes no sense. 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. + """ user_api = "blas" internal_api = "openblas" @@ -499,7 +525,54 @@ 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 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() + + match self._SHARED_STATE.controlling_thread: + case None: + # 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: + 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) + case controlling_thread if controlling_thread is this_thread: + # This thread controls OpenBLAS: + return self._real_set_num_threads(num_threads) + case _: + # Some other thread controlls OpenBLAS: + return + + def _real_set_num_threads(self, num_threads): # The OpenBLAS limiting API is process-wide, and we want current thread # limit if possible. When OpenBLAS is backed by OpenMP, using the # OpenMP API allows for current thread limiting when OpenMP has that @@ -558,6 +631,31 @@ def _get_architecture(self): return None +class _ControllingThreadOwnership(thread_local): + """Count how many paired limit-unlimits have modified OpenBLAS. + + Clear controlling thread if it hits zero. + + In the future this could potentialy also apply to non-OpenBLAS libraries + that have a shared process pool. + """ + + 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() + + +CONTROLLING_THREAD_OWNERSHIP = _ControllingThreadOwnership() + + class BLISController(LibController): """Controller class for BLIS""" @@ -943,10 +1041,12 @@ def __init__(self, controller, *, limits=None, user_api=None): self._set_threadpool_limits() def __enter__(self): + CONTROLLING_THREAD_OWNERSHIP.increment() return self def __exit__(self, type, value, traceback): self.restore_original_limits() + CONTROLLING_THREAD_OWNERSHIP.decrement() @classmethod def wrap(cls, controller, *, limits=None, user_api=None): @@ -1092,6 +1192,7 @@ def __enter__(self): # decorator. self._original_info = self._controller.info() self._set_threadpool_limits() + CONTROLLING_THREAD_OWNERSHIP.increment() return self