diff --git a/CHANGES.md b/CHANGES.md index a7b50dec..fe6e1dc0 100644 --- a/CHANGES.md +++ b/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) ================== diff --git a/conftest.py b/conftest.py index bf303839..fe7b8fcf 100644 --- a/conftest.py +++ b/conftest.py @@ -1 +1 @@ -collect_ignore = ["tests/_openmp_test_helper"] +collect_ignore = ["tests/_openmp_test_helper", "tests/_limit_blas"] diff --git a/tests/_limit_blas.py b/tests/_limit_blas.py new file mode 100644 index 00000000..0e38c9a1 --- /dev/null +++ b/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))) diff --git a/tests/test_api_introspection.py b/tests/test_api_introspection.py index 6c09d503..d534b409 100644 --- a/tests/test_api_introspection.py +++ b/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 diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 19012c5f..541662a6 100644 --- a/tests/test_threadpoolctl.py +++ b/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 diff --git a/threadpoolctl.py b/threadpoolctl.py index 299d112c..72fda34f 100644 --- a/threadpoolctl.py +++ b/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