FIX OpenBLAS pthreads can only be controlled from one thread at a time

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

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)

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)

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(

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