diff --git a/CHANGES.md b/CHANGES.md index 0e643c66..9d5c7805 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -9,6 +9,16 @@ in package metadata. https://github.com/joblib/threadpoolctl/issues/251 +- Faster performance when creating a `ThreadpoolController`. + https://github.com/joblib/threadpoolctl/pull/249 + +- Shared libraries are now cached across runs, whereas previously they would be + unloaded after the `ThreadpoolController` was garbage collected. + https://github.com/joblib/threadpoolctl/pull/249 + +- A single shared library can no longer have multiple `LibController` instances. + https://github.com/joblib/threadpoolctl/pull/249 + 3.7.0 (2026-09-15) ================== diff --git a/benchmarks/bench_context_manager_overhead.py b/benchmarks/bench_context_manager_overhead.py index d3b69c15..dc27a2a9 100644 --- a/benchmarks/bench_context_manager_overhead.py +++ b/benchmarks/bench_context_manager_overhead.py @@ -27,4 +27,4 @@ pass timings.append(time.time() - t) -print(f"Overhead per call: {mean(timings) * 1e3:.3f} +/-{stdev(timings) * 1e3:.3f} ms") +print(f"Overhead per call: {mean(timings) * 1e3:.4f} +/-{stdev(timings) * 1e3:.4f} ms") diff --git a/tests/test_cache.py b/tests/test_cache.py new file mode 100644 index 00000000..39d8fd96 --- /dev/null +++ b/tests/test_cache.py @@ -0,0 +1,59 @@ +"""Tests for ``threadpoolctl``'s internal caching layer.""" + +from ctypes import CDLL + +import pytest + +from threadpoolctl import ThreadpoolController, _CDLL_CACHE + + +def test_cdlls_are_cached(): + """ + When a ``LibController`` is created, it reuses the same ``CDLL``. + """ + pytest.importorskip("numpy") + + controller = ThreadpoolController() + if not controller.lib_controllers: + pytest.skip("No libraries loaded") + + cached_cdll = controller.lib_controllers[0].dynlib + assert isinstance(cached_cdll, CDLL) + + controller2 = ThreadpoolController() + assert cached_cdll is controller2.lib_controllers[0].dynlib + + +def test_cache_methods_on_dynlib(request): + """ + ``_CDLLCache.cache_method_on_dynlib()`` caches the result, tied to the + underlying ``CDLL`` as an invalidation key. + """ + pytest.importorskip("numpy") + + controller = ThreadpoolController() + libs = controller.select(user_api="blas").lib_controllers + if not libs: + pytest.skip("No libraries loaded") + + @_CDLL_CACHE.cache_method_on_dynlib + def my_extra_method(self): + return object() + + # Can't use pytest's monkeypatch since this method doesn't already exist, + # so add it manually: + libs[0].__class__.my_extra_method = my_extra_method + + def cleanup(): + del libs[0].__class__.my_extra_method + + request.addfinalizer(cleanup) + + # Accessing underlying, uncached method returns different objects each time: + initial = libs[0].my_extra_method() + assert libs[0].my_extra_method.__wrapped__(libs[0]) is not initial + assert libs[0].my_extra_method.__wrapped__(libs[0]) is not initial + + # Only one call should ever happen when using the cached method, however: + assert libs[0].my_extra_method() is initial + assert libs[0].my_extra_method() is initial diff --git a/threadpoolctl.py b/threadpoolctl.py index 8e8268ed..f34ce41a 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -11,6 +11,8 @@ # adapted from code by Intel developer @anton-malakhov available at # https://github.com/IntelPython/smp (Copyright (c) 2017, Intel Corporation) # and also published under the BSD 3-Clause license +from __future__ import annotations + import os import re import sys @@ -18,10 +20,11 @@ import itertools import textwrap from threading import Thread -from typing import Callable, Literal, final +from typing import Callable, Literal, TypeVar, final import warnings from abc import ABC, abstractmethod -from functools import lru_cache +from dataclasses import dataclass, field +from functools import lru_cache, wraps from contextlib import ContextDecorator # ctypes.util is not imported on Linux: on CPython 3.14 it allocates a @@ -223,7 +226,7 @@ def __init__(self, *, filepath=None, prefix=None, parent=None): self.parent = parent self.prefix = prefix self.filepath = filepath - self.dynlib = ctypes.CDLL(filepath, mode=_RTLD_NOLOAD) + self.dynlib = _CDLL_CACHE.get_cdll(filepath) self._symbol_prefix, self._symbol_suffix = self._find_affixes() self.version = self.get_version() self.set_additional_attributes() @@ -289,6 +292,159 @@ def _get_symbol(self, name): ) +# Internal marker for missing cache lookups: +_MISSING = object() + + +_T = TypeVar("_T") + + +@dataclass +class _CDLLCache: + """ + Cache CDLL instances and their associated method results, as well as which + ``LibController`` subclass to use for a given shared library file path. + """ + + _controller_cache: dict[str, type[LibController] | None] = field( + default_factory=dict + ) + + _cdll_cache: dict[str, ctypes.CDLL] = field(default_factory=dict) + + _method_result_cache: dict[tuple[type[LibController], str, ctypes.CDLL], object] = ( + field(default_factory=dict) + ) + + def _check_prefix( + self, library_basename: str, filename_prefixes: list[str] + ) -> str | None: + """Return the prefix library_basename starts with + + Return None if none matches. + """ + for prefix in filename_prefixes: + if library_basename.startswith(prefix): + return prefix + return None + + def create_controller( + self, filepath: str, parent: ThreadpoolController + ) -> LibController | None: + """ + Create the associated controller, if there is one, relying on cached + info for the given ``filepath``. + """ + result = self._controller_cache.get(filepath, _MISSING) + if result is not _MISSING: + controller_class, filepath, prefix = result + if controller_class is None: + return None + return controller_class(filepath=filepath, prefix=prefix, parent=parent) + + # It's not in the cache, so proceed to actually search for it. + + # Required to resolve symlinks + original_filepath = filepath + filepath = _realpath(filepath) + # `lower` required to take account of OpenMP dll case on Windows + # (vcomp, VCOMP, Vcomp, ...) + filename = os.path.basename(filepath).lower() + + # Loop through supported libraries to find if this filename corresponds + # to a supported one. + for controller_class in _ALL_CONTROLLERS: + # check if filename matches a supported prefix + prefix = self._check_prefix(filename, controller_class.filename_prefixes) + + # filename does not match any of the prefixes of the candidate + # library. move to next library. + if prefix is None: + continue + + # Legacy workaround for BLAS libraries that conda-forge used to expose + # on Windows as libblas.dll, disambiguated via implementation-specific + # symbols. Current conda-forge stacks no longer load libblas.dll (e.g. + # MKL is exposed as mkl_rt..dll instead), so this path is + # kept for older installs but cannot be exercised in today's CI. + if prefix == "libblas": + if filename.endswith(".dll"): + libblas = ctypes.CDLL(filepath, _RTLD_NOLOAD) + if not any( + hasattr(libblas, func) + for func in controller_class.check_symbols + ): + continue + else: + # Non-Windows libblas DSOs (e.g. from openblas) lack the symbols + # needed to instantiate a controller and would duplicate entries. + continue + + # filename matches a prefix. Now we check if the library has the symbols we + # are looking for. If none of the symbols exists, it's very likely not the + # expected library (e.g. a library having a common prefix with one of the + # our supported libraries). Otherwise, create and store the library + # controller. + try: + lib_controller = controller_class( + filepath=filepath, prefix=prefix, parent=parent + ) + except OSError: + # Probably because we couldn't load the filepath as a CDLL. + continue + + if not hasattr(controller_class, "check_symbols") or any( + hasattr(lib_controller.dynlib, func) + for func in controller_class.check_symbols + ): + self._controller_cache[original_filepath] = ( + controller_class, + filepath, + prefix, + ) + return lib_controller + + # Didn't find any matching libraries. + self._controller_cache[original_filepath] = (None, "", "") + return None + + def get_cdll(self, filepath: str) -> CDLL: + """ + Get the ``CDLL`` for a path, loading if necessary, using a cached + version if it was already loaded. + """ + result = self._cdll_cache.get(filepath, _MISSING) + if result is _MISSING: + result = ctypes.CDLL(filepath, mode=_RTLD_NOLOAD) + self._cdll_cache[filepath] = result + return result + + def cache_method_on_dynlib( + self, method: Callable[[LibController], _T] + ) -> Callable[[LibController], _T]: + """ + Caching decorator for idempotent read-only methods of + ``LibController``, with the cache shared across instances that have the + same ``CDLL`` instance. + """ + cache = self._method_result_cache + name = method.__name__ + + @wraps(method) + def wrapper(self): + key = (self.__class__, name, self.dynlib) + result = cache.get(key, _MISSING) + if result is _MISSING: + result = method(self) + cache[key] = result + return result + + return wrapper + + +_CDLL_CACHE = _CDLLCache() + + class OpenBLASController(LibController): """Controller class for OpenBLAS""" @@ -310,6 +466,7 @@ class OpenBLASController(LibController): for prefix, suffix in itertools.product(_symbol_prefixes, _symbol_suffixes) ) + @_CDLL_CACHE.cache_method_on_dynlib def _find_affixes(self): for prefix, suffix in itertools.product( self._symbol_prefixes, self._symbol_suffixes @@ -355,6 +512,7 @@ def set_num_threads(self, num_threads): return set_num_threads_func(num_threads) return None + @_CDLL_CACHE.cache_method_on_dynlib def get_version(self): # None means OpenBLAS is not loaded or version < 0.3.4, since OpenBLAS # did not expose its version before that. @@ -367,6 +525,7 @@ def get_version(self): return None return None + @_CDLL_CACHE.cache_method_on_dynlib def _get_threading_layer(self): """Return the threading layer of OpenBLAS""" get_threading_layer_func = self._get_symbol("openblas_get_parallel") @@ -379,6 +538,7 @@ def _get_threading_layer(self): return "disabled" return "unknown" + @_CDLL_CACHE.cache_method_on_dynlib def _get_architecture(self): """Return the architecture detected by OpenBLAS""" get_architecture_func = self._get_symbol("openblas_get_corename") @@ -424,6 +584,7 @@ def set_num_threads(self, num_threads): ) return set_func(num_threads) + @_CDLL_CACHE.cache_method_on_dynlib def get_version(self): get_version_ = getattr(self.dynlib, "bli_info_get_version_str", None) if get_version_ is None: @@ -432,6 +593,7 @@ def get_version(self): get_version_.restype = ctypes.c_char_p return get_version_().decode("utf-8") + @_CDLL_CACHE.cache_method_on_dynlib def _get_threading_layer(self): """Return the threading layer of BLIS""" if getattr(self.dynlib, "bli_info_get_enable_openmp", lambda: False)(): @@ -440,6 +602,7 @@ def _get_threading_layer(self): return "pthreads" return "disabled" + @_CDLL_CACHE.cache_method_on_dynlib def _get_architecture(self): """Return the architecture detected by BLIS""" bli_arch_query_id = getattr(self.dynlib, "bli_arch_query_id", None) @@ -503,6 +666,7 @@ def set_num_threads(self, num_threads): ) return set_func(num_threads) + @_CDLL_CACHE.cache_method_on_dynlib def get_version(self): get_version_ = getattr(self.dynlib, "flexiblas_get_version", None) if get_version_ is None: @@ -612,6 +776,7 @@ def set_num_threads(self, num_threads): ) return set_func(num_threads) + @_CDLL_CACHE.cache_method_on_dynlib def get_version(self): if not hasattr(self.dynlib, "MKL_Get_Version_String"): return None @@ -625,6 +790,7 @@ def get_version(self): version = group.groups()[0] return version.strip() + @_CDLL_CACHE.cache_method_on_dynlib def _get_threading_layer(self): """Return the threading layer of MKL""" # The function mkl_set_threading_layer returns the current threading @@ -972,6 +1138,17 @@ def wrap(cls, limits=None, user_api=None): return super().wrap(ThreadpoolController(), limits=limits, user_api=user_api) +_INCOMPATIBLE_OPENMP_MESSAGE = """ +Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at +the same time. Both libraries are known to be incompatible and this +can cause random crashes or deadlocks on Linux when loaded in the +same Python program. +Using threadpoolctl may cause crashes or deadlocks. For more +information and possible workarounds, please see + https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md +""" + + class ThreadpoolController: """Collection of LibController objects for all loaded supported libraries @@ -1586,73 +1763,15 @@ def _find_libraries_pyodide(self): def _make_controller_from_path(self, filepath): """Store a library controller if it is supported and selected""" - # Required to resolve symlinks - filepath = _realpath(filepath) - # `lower` required to take account of OpenMP dll case on Windows - # (vcomp, VCOMP, Vcomp, ...) - filename = os.path.basename(filepath).lower() - - # Loop through supported libraries to find if this filename corresponds - # to a supported one. - for controller_class in _ALL_CONTROLLERS: - # check if filename matches a supported prefix - prefix = self._check_prefix(filename, controller_class.filename_prefixes) - - # filename does not match any of the prefixes of the candidate - # library. move to next library. - if prefix is None: - continue - - # Legacy workaround for BLAS libraries that conda-forge used to expose - # on Windows as libblas.dll, disambiguated via implementation-specific - # symbols. Current conda-forge stacks no longer load libblas.dll (e.g. - # MKL is exposed as mkl_rt..dll instead), so this path is - # kept for older installs but cannot be exercised in today's CI. - if prefix == "libblas": - if filename.endswith(".dll"): - libblas = ctypes.CDLL(filepath, _RTLD_NOLOAD) - if not any( - hasattr(libblas, func) - for func in controller_class.check_symbols - ): - continue - else: - # Non-Windows libblas DSOs (e.g. from openblas) lack the symbols - # needed to instantiate a controller and would duplicate entries. - continue - - # filename matches a prefix. Now we check if the library has the symbols we - # are looking for. If none of the symbols exists, it's very likely not the - # expected library (e.g. a library having a common prefix with one of the - # our supported libraries). Otherwise, create and store the library - # controller. - try: - lib_controller = controller_class( - filepath=filepath, prefix=prefix, parent=self - ) - except OSError: - # Probably because we couldn't load the filepath as a CDLL. - continue - - if filepath in (lib.filepath for lib in self.lib_controllers): - # We already have a controller for this library. - continue - - if not hasattr(controller_class, "check_symbols") or any( - hasattr(lib_controller.dynlib, func) - for func in controller_class.check_symbols + lib_controller = _CDLL_CACHE.create_controller(filepath, self) + if lib_controller is not None: + if lib_controller.filepath in ( + lib.filepath for lib in self.lib_controllers ): - self.lib_controllers.append(lib_controller) - - def _check_prefix(self, library_basename, filename_prefixes): - """Return the prefix library_basename starts with + # We already have a controller for this library. + return - Return None if none matches. - """ - for prefix in filename_prefixes: - if library_basename.startswith(prefix): - return prefix - return None + self.lib_controllers.append(lib_controller) def _warn_if_incompatible_openmp(self): """Raise a warning if llvm-OpenMP and intel-OpenMP are both loaded""" @@ -1664,17 +1783,8 @@ def _warn_if_incompatible_openmp(self): return prefixes = [lib_controller.prefix for lib_controller in self.lib_controllers] - msg = textwrap.dedent(""" - Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at - the same time. Both libraries are known to be incompatible and this - can cause random crashes or deadlocks on Linux when loaded in the - same Python program. - Using threadpoolctl may cause crashes or deadlocks. For more - information and possible workarounds, please see - https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md - """) if "libomp" in prefixes and "libiomp" in prefixes: - warnings.warn(msg, RuntimeWarning) + warnings.warn(_INCOMPATIBLE_OPENMP_MESSAGE, RuntimeWarning) @classmethod def _get_libc(cls):