Use thread-local limiting APIs when possible, and fix deadlock

CHANGES.md

@@ -10,6 +10,26 @@
 - Only warn about simultaneous `libomp` and `libiomp` usage on Linux, where the
   incompatibility is known to cause crashes.
 
+- Fixed a deadlock triggered by getting or setting MKL's number of threads from
+  parallel threads when using MKL with libiomp (Intel threading) on Linux.
+  https://github.com/joblib/threadpoolctl/pull/228
+
+- Going forward, setting the number of threads will only have a thread-local
+  impact if feasible (for example, at minimum the underlying library must
+  support this option, and many don't.)
+  https://github.com/joblib/threadpoolctl/pull/228
+
+- For MKL, setting the number of threads is now thread-local, i.e. limiting the
+  number of threads won't impact MKL's thread pool size when using MKL in other
+  Python threads.
+  https://github.com/joblib/threadpoolctl/pull/228
+
+- For OpenBLAS compiled with OpenMP on Linux and macOS, setting the number of
+  threads is now thread-local, i.e. won't impact OpenBLAS thread pool size in
+  other Python threads. On Windows behavior is likely process-wide, but this may
+  depend on how OpenBLAS was compiled with OpenMP.
+  https://github.com/joblib/threadpoolctl/pull/228
+
 3.6.0 (2025-03-13)
 ==================
 

conftest.py

@@ -1 +1 @@
-collect_ignore = ["tests/_openmp_test_helper"]
+collect_ignore = ["tests/_openmp_test_helper", "tests/_limit_blas"]

tests/_limit_blas.py

@@ -0,0 +1,22 @@
+# Used by test_setting_limit_on_thread_local_blas_api_is_actually_thread_local()
+
+from concurrent.futures import ThreadPoolExecutor
+from time import sleep
+import sys
+
+import numpy as np
+import threadpoolctl
+
+ARR = np.ones((1500, 1500))
+
+
+def in_thread(_):
+    with threadpoolctl.threadpool_limits(limits=int(sys.argv[1]), user_api="blas"):
+        ARR.dot(ARR)
+    # Make sure jobs are evenly distributed and don't end up in one thread.
+    sleep(0.01)
+
+
+if __name__ == "__main__":
+    with ThreadPoolExecutor(2) as pool:
+        list(pool.map(in_thread, range(2)))

tests/test_api_introspection.py

@@ -60,9 +60,6 @@ def test_determine_thread_limit_scope_processwide(default: int) -> None:
     assert _determine_thread_limit_scope(api.get, api.set) == "process"
 
 
-@pytest.mark.skipif(
-    sys.platform != "linux", reason="Non-Linux OpenMP might be different"
-)
 @pytest.mark.parametrize(
     ["select_filter", "expected_thread_limit_scope", "extra_check"],
     [
@@ -77,7 +74,12 @@ def test_determine_thread_limit_scope_processwide(default: int) -> None:
             # pthreads here.
             lambda lib: lib.threading_layer == "pthreads",
         ),
-        ({"user_api": "openmp"}, "current_thread", lambda _lib: True),
+        (
+            {"user_api": "openmp"},
+            "current_thread",
+            # Windows OpenMP is process-wide:
+            lambda _lib: sys.platform in ("linux", "darwin"),
+        ),
     ],
 )
 def test_api_scope(
@@ -94,9 +96,11 @@ def test_api_scope(
     if not controller.lib_controllers:
         pytest.skip(f"{select_filter} controller not found")
 
-    for lib in controller.lib_controllers:
-        if not extra_check(lib):
-            pytest.skip("extra check returned false")
+    libs = [lib for lib in controller.lib_controllers if extra_check(lib)]
+    if not libs:
+        pytest.skip("No libraries matched the requirements")
+
+    for lib in libs:
         assert (
             _determine_thread_limit_scope(lib.get_num_threads, lib.set_num_threads)
             == expected_thread_limit_scope

tests/test_threadpoolctl.py

@@ -6,10 +6,12 @@
 import re
 import subprocess
 import sys
+from shutil import which
 
 from threadpoolctl import threadpool_limits, threadpool_info
-from threadpoolctl import ThreadpoolController
+from threadpoolctl import LibController, ThreadpoolController
 from threadpoolctl import _ALL_PREFIXES, _ALL_USER_APIS
+from threadpoolctl import _determine_thread_limit_scope
 
 from .utils import cython_extensions_compiled
 from .utils import check_nested_prange_blas
@@ -615,6 +617,7 @@ def test_architecture():
     expected_openblas_architectures = (
         # XXX: add more as needed by CI or developer laptops
         "armv8",
+        "barcelona",
         "cooperlake",
         "haswell",
         "neoversen1",
@@ -795,3 +798,118 @@ def test_custom_controller():
         assert mylib_controller.num_threads == 1
 
     assert ThreadpoolController().info() == original_info
+
+
+def parse_version(version: str) -> list[int]:
+    return list(map(int, version.split(".")))
+
+
+@pytest.fixture(
+    params=[
+        (
+            {"internal_api": "openblas"},
+            lambda lib: (
+                lib.threading_layer == "openmp"
+                # For Windows support, see
+                # https://github.com/joblib/threadpoolctl/issues/230
+                and sys.platform in ("linux", "darwin")
+                and parse_version(lib.version) >= parse_version("0.3.34")
+            ),
+        ),
+        (
+            {"internal_api": "mkl"},
+            lambda _lib: True,
+        ),
+    ],
+    # ids correspond to the params above:
+    ids=["openblas-openmp", "mkl"],
+)
+def thread_local_blas_lib(request) -> LibController:
+    """Create all LibControllers that use a thread-local setting."""
+    select_filter, extra_check = request.param
+    controller = ThreadpoolController().select(**select_filter)
+    if not controller.lib_controllers:
+        pytest.skip(f"{select_filter} controller not found")
+
+    libs = [
+        lib
+        for lib in controller.lib_controllers
+        if extra_check(lib) and lib.internal_api == select_filter["internal_api"]
+    ]
+    if not libs:
+        pytest.skip("No libraries matched the requirements")
+
+    assert len(libs) == 1
+    return libs[0]
+
+
+def test_setting_limit_on_thread_local_blas_api_is_reported_as_thread_local(
+    thread_local_blas_lib: LibController,
+) -> None:
+    """
+    Setting the number of threads for libraries that support thread-local
+    setting API is reported as doing so.
+
+    This doesn't check actual behavior, only reported behavior.
+    """
+    lib = thread_local_blas_lib
+    scope = _determine_thread_limit_scope(lib.get_num_threads, lib.set_num_threads)
+    assert scope == "current_thread"
+
+
+@pytest.mark.skipif(
+    sys.platform != "linux" or which("strace") is None,
+    reason="requires strace on Linux",
+)
+def test_setting_limit_on_thread_local_blas_api_is_actually_thread_local(
+    thread_local_blas_lib: LibController,
+) -> None:
+    """
+    Setting the number of threads for libraries that support thread-local
+    setting API actually does so.
+    """
+
+    # The test script uses NumPy, there might be multiple BLAS in this test
+    # process, and we want to only run if _NumPy_ uses that library.
+    # So check that before proceeding.
+    output = json.loads(
+        subprocess.check_output(
+            [
+                "python",
+                "-c",
+                "import numpy, json, threadpoolctl; print(json.dumps(threadpoolctl.threadpool_info()))",
+            ]
+        )
+    )
+    found_correct_blas = False
+    for library in output:
+        if library["internal_api"] == thread_local_blas_lib.internal_api:
+            found_correct_blas = True
+            break
+    if not found_correct_blas:
+        pytest.skip("NumPy doesn't use the BLAS we want to test")
+
+    def num_threads_created(limit: int) -> int:
+        result = 0
+        for line in subprocess.check_output(
+            [
+                "strace",
+                "-f",
+                "-e",
+                "clone3",
+                "python",
+                "-m",
+                "tests._limit_blas",
+                str(limit),
+            ],
+            stderr=subprocess.STDOUT,
+        ).splitlines():
+            if b"clone3(" in line and b"CLONE_THREAD" in line:
+                result += 1
+        return result
+
+    # _limit_blas runs BLAS operations in 2 Python threads, so by changing the
+    # BLAS limit from 1 to 4 we expect an extra 2 * (4 - 1) == 6 threads.
+    nmc_1 = num_threads_created(1)
+    nmc_4 = num_threads_created(4)
+    assert nmc_4 - nmc_1 == 6

threadpoolctl.py

@@ -240,7 +240,11 @@ def get_num_threads(self):
 
     @abstractmethod
     def set_num_threads(self, num_threads):
-        """Set the maximum number of threads to use"""
+        """Set the maximum number of threads to use
+
+        When possible, implementations of this method should choose a thread
+        limiting API that only applies to the current thread.
+        """
 
     @abstractmethod
     def get_version(self):
@@ -289,13 +293,35 @@ def set_additional_attributes(self):
         self.architecture = self._get_architecture()
 
     def get_num_threads(self):
-        get_num_threads_func = self._get_symbol("openblas_get_num_threads")
+        # See discussion in set_num_threads for details:
+        if self.threading_layer == "openmp":
+            symbol = "omp_get_max_threads"
+        else:
+            symbol = "openblas_get_num_threads"
+        get_num_threads_func = self._get_symbol(symbol)
         if get_num_threads_func is not None:
             return get_num_threads_func()
         return None
 
     def set_num_threads(self, num_threads):
-        set_num_threads_func = self._get_symbol("openblas_set_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
+        # behavior. That is the case for libgomp, libomp, and libiomp, what you
+        # would find on Linux or macOS.
+        #
+        # On Windows the Visual C++ OpenMP API is process-wide, unfortunately,
+        # though this may be fixed if the /openmp:llvm flag is used:
+        # https://github.com/joblib/threadpoolctl/issues/230
+        #
+        # Also worth knowing that before v0.3.34, the OpenBLAS limiting API is
+        # broken when using OpenMP threading:
+        # https://github.com/OpenMathLib/OpenBLAS/issues/5806
+        if self.threading_layer == "openmp":
+            symbol = "omp_set_num_threads"
+        else:
+            symbol = "openblas_set_num_threads"
+        set_num_threads_func = self._get_symbol(symbol)
         if set_num_threads_func is not None:
             return set_num_threads_func(num_threads)
         return None
@@ -539,7 +565,7 @@ class MKLController(LibController):
     )
     check_symbols = (
         "MKL_Get_Max_Threads",
-        "MKL_Set_Num_Threads",
+        "MKL_Set_Num_Threads_Local",
         "MKL_Get_Version_String",
         "MKL_Set_Threading_Layer",
     )
@@ -552,7 +578,9 @@ def get_num_threads(self):
         return get_func()
 
     def set_num_threads(self, num_threads):
-        set_func = getattr(self.dynlib, "MKL_Set_Num_Threads", lambda num_threads: None)
+        set_func = getattr(
+            self.dynlib, "MKL_Set_Num_Threads_Local", lambda num_threads: None
+        )
         return set_func(num_threads)
 
     def get_version(self):
@@ -1120,16 +1148,18 @@ def _find_libraries_with_dl_iterate_phdr(self):
             )
             return []
 
+        filepaths = []
+
         # Callback function for `dl_iterate_phdr` which is called for every
-        # library loaded in the current process until it returns 1.
+        # library loaded in the current process until it returns 1. To minimize
+        # the potential for deadlocks (see #228), this code should not do
+        # anything that might result in reentrancy into the library, the dl
+        # system, or anything else.
         def match_library_callback(info, size, data):
             # Get the path of the current library
             filepath = info.contents.dlpi_name
             if filepath:
-                filepath = filepath.decode("utf-8")
-
-                # Store the library controller if it is supported and selected
-                self._make_controller_from_path(filepath)
+                filepaths.append(filepath)
             return 0
 
         c_func_signature = ctypes.CFUNCTYPE(
@@ -1143,6 +1173,12 @@ def match_library_callback(info, size, data):
         data = ctypes.c_char_p(b"")
         libc.dl_iterate_phdr(c_match_library_callback, data)
 
+        # Now that a list of filepaths is available, load the respective
+        # libraries:
+        for filepath in filepaths:
+            # Store the library controller if it is supported and selected
+            self._make_controller_from_path(filepath.decode("utf-8"))
+
     def _find_libraries_with_dyld(self):
         """Loop through loaded libraries and return binders on supported ones