diff --git a/CHANGES.md b/CHANGES.md index 0e643c66..f22b69c0 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -5,6 +5,11 @@ quite a bit of overhead. https://github.com/joblib/threadpoolctl/pull/250 +- 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. + https://github.com/joblib/threadpoolctl/issues/257 + - Added official support for Python 3.15, and declared support for Python 3.14 in package metadata. https://github.com/joblib/threadpoolctl/issues/251 diff --git a/README.md b/README.md index 3f2b196f..936c6214 100644 --- a/README.md +++ b/README.md @@ -178,8 +178,9 @@ in specific sections of your Python program: The threadpools can also be controlled via the object oriented API, which is especially useful to avoid searching through all the loaded shared libraries -each time. **Note that it will not act on libraries loaded after the instantiation -of the `ThreadpoolController`!** +each time. That is, you can reuse a `ThreadpoolController` to reduce performance +overhead. **However, note that it will not act on libraries loaded after the +instantiation of the `ThreadpoolController`!** ```python >>> from threadpoolctl import ThreadpoolController @@ -191,6 +192,43 @@ of the `ThreadpoolController`!** ... a_squared = a @ a ``` +As a compromise between speed and catching newly loaded controllable libraries +you can use `get_cached_controller()`. This will notice any newly loaded +libraries triggered by importing new libraries, but might miss new libraries +loaded by compiled extensions directly using mechanisms like `dlopen()` on Linux +(as opposed to linked), or via `ctypes.CDLL`. + +```python +import numpy + +with get_cached_controller().limit(limits=1, user_api="blas"): + a = np.random.randn(1000, 1000) + a_squared = a @ a +``` + +To get a sense of how `get_cached_controller()` cache invalidation happens, +consider the following example: + +```python +from threadpoolctl import get_cached_controller + +# This ThreadpoolController will not know about the BLAS library loaded by +# NumPy, since it hasn't yet been imported: +controller = get_cached_controller() + +# Calling the API again doesn't create a new instance, so it's fast: +assert get_cached_controller() is controller + +import numpy as np + +# Now that NumPy has been imported, a new ThreadpoolController is returned that +# will know about NumPy's linked BLAS: +assert get_cached_controller() is not controller +``` + +You can clear the cached controller manually by running +``reset_cached_controller()``. + ### Restricting the Limits to the Scope of a Function `threadpool_limits` and `ThreadpoolController` can also be used as decorators to set diff --git a/benchmarks/bench_context_manager_overhead.py b/benchmarks/bench_context_manager_overhead.py index a5c9b8fa..3698911b 100644 --- a/benchmarks/bench_context_manager_overhead.py +++ b/benchmarks/bench_context_manager_overhead.py @@ -3,7 +3,7 @@ from pprint import pprint from statistics import mean, stdev -from threadpoolctl import threadpool_info, threadpool_limits +from threadpoolctl import get_cached_controller, threadpool_info, threadpool_limits parser = ArgumentParser(description="Measure threadpool_limits call overhead.") parser.add_argument( @@ -28,4 +28,18 @@ pass timings.append(time.time() - t) -print(f"Overhead per call: {mean(timings) * 1e3:.3f} +/-{stdev(timings) * 1e3:.3f} ms") +print( + f"Overhead per `with threadpool_limits()` call: {mean(timings) * 1e3:.3f} +/-{stdev(timings) * 1e3:.3f} ms" +) + + +timings = [] +for _ in range(args.n_calls): + t = time.time() + with get_cached_controller().limit(limits=1): + pass + timings.append(time.time() - t) + +print( + f"Overhead per `with get_cached_controller().limit()` call: {mean(timings) * 1e3:.3f} +/-{stdev(timings) * 1e3:.3f} ms" +) diff --git a/conftest.py b/conftest.py index 11855652..451cff36 100644 --- a/conftest.py +++ b/conftest.py @@ -3,4 +3,5 @@ "tests/_openmp_test_helper.py", "tests/_limit_blas.py", "tests/_dl_iterate_phdr_deadlock.py", + "tests/_get_cached_controller.py", ] diff --git a/tests/_get_cached_controller.py b/tests/_get_cached_controller.py new file mode 100644 index 00000000..954eed3f --- /dev/null +++ b/tests/_get_cached_controller.py @@ -0,0 +1,34 @@ +"""Test script for ``_get_cached_controller``.""" + +from threadpoolctl import ( + ThreadpoolController, + get_cached_controller, + reset_cached_controller, +) + +controller = get_cached_controller() +assert isinstance(controller, ThreadpoolController) +# Same controller is returned on another call: +assert controller is get_cached_controller() +# No libraries loaded: +assert len(controller.lib_controllers) == 0 + +import numpy as np # noqa: F401 (also gets us BLAS) + +# Now a new controller is returned: +controller2 = get_cached_controller() +assert controller is not controller2 +assert isinstance(controller2, ThreadpoolController) +# And it noticed the newly loaded controllable library: +assert len(controller2.lib_controllers) >= 1 + +# reset_cached_controller() clears the cache: +assert get_cached_controller() is controller2 +reset_cached_controller() +controller3 = get_cached_controller() +assert controller3 is not controller2 +assert get_cached_controller() is controller3 + + +# Special success exit code: +raise SystemExit(17) diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 2b7d6065..6b3bf1f6 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -17,6 +17,7 @@ LibController, ThreadpoolController, _determine_thread_limit_scope, + get_cached_controller, threadpool_info, threadpool_limits, ) @@ -741,10 +742,13 @@ def test_flexiblas_switch(): # at first, only "OPENBLAS_CONDA" is loaded assert fb_controller.current_backend == "OPENBLAS_CONDA" assert fb_controller.loaded_backends == ["OPENBLAS_CONDA"] + cached_controller = get_cached_controller() fb_controller.switch_backend("NETLIB") assert fb_controller.current_backend == "NETLIB" assert fb_controller.loaded_backends == ["OPENBLAS_CONDA", "NETLIB"] + # Switching invalidated the cached controller: + assert get_cached_controller() is not cached_controller if sys.platform == "linux": mkl_path = f"{os.getenv('CONDA_PREFIX')}/lib/libmkl_rt.so" @@ -1148,6 +1152,24 @@ def test_controller_parallelism_no_deadlocks(): assert process.returncode == 17 +def test_get_cached_controller(): + """ + ``get_cached_controller()`` returns a cached ``ThreadpoolController``, and + invalidates the cache if new modules are imported. + """ + pytest.importorskip("numpy") + controller = ThreadpoolController().select(user_api="blas") + if not controller.lib_controllers: + pytest.skip("No BLAS in NumPy") + + process = subprocess.run( + [sys.executable, "-m", "tests._get_cached_controller"], timeout=10, check=False + ) + + # Special code indicating success: + assert process.returncode == 17 + + @pytest.mark.skipif( not sys.platform.startswith("linux"), reason="ctypes.util is only avoided on Linux (#225) in Python 3.14", diff --git a/threadpoolctl.py b/threadpoolctl.py index 5ff701ef..b18c1210 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -11,6 +11,9 @@ # 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 ctypes import itertools import os @@ -22,7 +25,7 @@ from collections.abc import Callable from contextlib import ContextDecorator from functools import lru_cache -from threading import Thread +from threading import Lock, Thread from typing import Literal, final # ctypes.util is not imported on Linux: on CPython 3.14 it allocates a @@ -40,12 +43,22 @@ # CPython before 3.14 does not provide dll inspection. dllist = None +# Ensure we don't get spurious cache invalidation the first time +# ThreadpoolController is created, as the Windows dll listing methods import +# from there: +if sys.platform == "win32": + from ctypes import wintypes + + del wintypes + __version__ = "3.8.0.dev0" __all__ = [ "LibController", "ThreadpoolController", + "get_cached_controller", "register", + "reset_cached_controller", "threadpool_info", "threadpool_limits", ] @@ -571,6 +584,11 @@ def switch_backend(self, backend): f"{self.available_backends} or the path to a valid shared library." ) + # Invalidate the global cached controller: + global _CACHED_CONTROLLER + with _CACHE_LOCK: + _CACHED_CONTROLLER = None + # Trigger a new search of loaded shared libraries since loading a new # backend caused a dlopen. self.parent._load_libraries() @@ -1739,5 +1757,53 @@ def _main(): print(json.dumps(threadpool_info(debugging_info=True), indent=2)) +# A cached instance of ThreadpoolController, created by +# get_cached_controller(): +_CACHED_CONTROLLER = None +# Cached number of modules in sys.modules, used by get_cached_controller() +_CACHED_SYS_MODULES_LEN = 0 +# The lock controlling access to the above: +_CACHE_LOCK = Lock() + + +def get_cached_controller() -> ThreadpoolController: + """Return a cached ``ThreadpoolController``, as a speed optimization. + + The cached version is regenerated if the number of imported of modules has + changed, since this indicates the possibility that a new controllable + library has been loaded. Invalidation also happens if Flexiblas is in use + and has changed its backend. + + It is possible that somehow a controllable library gets loaded without + importing a Python module, e.g. using ``dlopen()`` on Linux, a ctypes + ``CDLL``, and the like. This is rare, however. Outside of these + situations, using ``get_cached_controller()`` will be much faster than + creating a new ``ThreadpoolController`` instance. + + The cached controller can be explicitly invalidated by calling the + `reset_cached_controller()` module level function. + """ + global _CACHED_CONTROLLER, _CACHED_SYS_MODULES_LEN + # Fast path, with no lock: + if _CACHED_SYS_MODULES_LEN == len(sys.modules) and _CACHED_CONTROLLER is not None: + return _CACHED_CONTROLLER + with _CACHE_LOCK: + # Do this once, so there's no race condition with imports in other + # threads: + num_modules = len(sys.modules) + if _CACHED_CONTROLLER is None or num_modules != _CACHED_SYS_MODULES_LEN: + _CACHED_CONTROLLER = ThreadpoolController() + _CACHED_SYS_MODULES_LEN = num_modules + return _CACHED_CONTROLLER + + +def reset_cached_controller() -> None: + """Clear the cached ``ThreadpoolController``, if any.""" + global _CACHED_CONTROLLER, _CACHED_SYS_MODULES_LEN + with _CACHE_LOCK: + _CACHED_CONTROLLER = None + _CACHED_SYS_MODULES_LEN = 0 + + if __name__ == "__main__": _main()