diff --git a/CHANGES.md b/CHANGES.md index 0e643c66..6e0e0ab5 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -5,6 +5,14 @@ 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. + https://github.com/joblib/threadpoolctl/pull/275 + - Added official support for Python 3.15, and declared support for Python 3.14 in package metadata. https://github.com/joblib/threadpoolctl/issues/251 diff --git a/README.md b/README.md index 73a6c1f6..f4f363b6 100644 --- a/README.md +++ b/README.md @@ -231,6 +231,12 @@ 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. +OpenBLAS with pthreads or Windows threading layer, affects the whole process, +and is therefore constrained to only work from the first thread that uses it, +the controlling thread. Other third party libraries allow setting the limit +per-thread. So we set limits both in the main thread (in case you're using +OpenBLAS+pthreads) and in each worker thread. + ```python from threadpoolctl import ThreadpoolController @@ -241,11 +247,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. @@ -257,7 +262,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 ... @@ -271,7 +276,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) diff --git a/conftest.py b/conftest.py index 11855652..10e1dd32 100644 --- a/conftest.py +++ b/conftest.py @@ -1,6 +1,18 @@ +import pytest + collect_ignore = [ "tests/empirical_scope_observation.py", "tests/_openmp_test_helper.py", "tests/_limit_blas.py", "tests/_dl_iterate_phdr_deadlock.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 2b7d6065..88251d66 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 import pytest @@ -683,6 +684,91 @@ 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: + clear_and_set() + 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 + + # 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 2eab17c3..e442b264 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -22,7 +22,7 @@ from collections.abc import Callable from contextlib import ContextDecorator from functools import lru_cache -from threading import Thread +from threading import Lock, Thread, current_thread from typing import Literal, final # ctypes.util is not imported on Linux: on CPython 3.14 it allocates a @@ -154,6 +154,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 @@ -179,6 +180,7 @@ def get_and_set() -> None: return "unknown" finally: set_n_threads(previous) + OpenBLASController.clear_controlling_thread(force=True) class LibController(ABC): @@ -290,7 +292,21 @@ def _get_symbol(self, name): class OpenBLASController(LibController): - """Controller class for OpenBLAS""" + """Controller class for OpenBLAS. + + When 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. + """ user_api = "blas" internal_api = "openblas" @@ -332,7 +348,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 + moethod. 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