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

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.

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.

conftest.py

@@ -1,7 +1,19 @@
+import pytest
+
 collect_ignore = [
     "tests/empirical_scope_observation.py",
     "tests/_openmp_test_helper.py",
     "tests/_limit_blas.py",
     "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)

tests/test_threadpoolctl.py

@@ -7,14 +7,17 @@
 import shutil
 import subprocess
 import sys
+from threading import Thread, current_thread
 
 import pytest
 
 import threadpoolctl
 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(