PERF New get_cached_controller() API

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

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

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"
+)

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",
 ]

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)

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",

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()