diff --git a/CHANGES.md b/CHANGES.md index 7317c063..6fd0f074 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -14,6 +14,7 @@ limit APIs will automatically call that API when they are done. https://github.com/joblib/threadpoolctl/pull/275 https://github.com/joblib/threadpoolctl/pull/289 + https://github.com/joblib/threadpoolctl/pull/294 - Add a new API to get a cached `ThreadpoolController` instance (`get_cached_controller()`), as an optimization for cases where limits are set @@ -50,6 +51,15 @@ - Dropped official support for Python 3.9. https://github.com/joblib/threadpoolctl/pull/255 +- Fixed recursive `@threadpool_limits.wrap()` calls so they restore limits + correctly. + https://github.com/joblib/threadpoolctl/pull/294 + +- Fixed `@threadpool_limits.wrap()` so it can use third-party libraries that + were registered or imported after the function was decorated. + https://github.com/joblib/threadpoolctl/pull/294 + + 3.7.0 (2026-09-15) ================== diff --git a/README.md b/README.md index fcc26221..30cf0735 100644 --- a/README.md +++ b/README.md @@ -1,705 +1,731 @@ -# Thread-pool Controls [![Build Status](https://github.com/joblib/threadpoolctl/actions/workflows/test.yml/badge.svg?branch=master)](https://github.com/joblib/threadpoolctl/actions?query=branch%3Amaster) [![codecov](https://codecov.io/gh/joblib/threadpoolctl/branch/master/graph/badge.svg)](https://codecov.io/gh/joblib/threadpoolctl) - -Python helpers to limit the number of threads used in the -threadpool-backed of common native libraries used for scientific -computing and data science (e.g. BLAS and OpenMP). - -Fine control of the underlying thread-pool size can be useful in -workloads that involve nested parallelism so as to mitigate -oversubscription issues. - -Note that "limiting the number of threads" in practice involves a variety of -potential semantics depending on the underlying third-party library; see the -section on [semantics](#semantics) below. - -## Installation - -- For users, install the last published version from PyPI: - - ```bash - pip install threadpoolctl - ``` - -- For contributors, install from the source repository in developer - mode: - - ```bash - pip install -r dev-requirements.txt - flit install --symlink - ``` - - then you run the tests with pytest: - - ```bash - pytest - ``` - -## Usage: Introspection and debugging - -### Command Line Interface - -Get a JSON description of thread-pools initialized when importing python -packages such as numpy or scipy for instance: - -``` -python -m threadpoolctl -i numpy scipy.linalg -[ - { - "filepath": "/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so", - "prefix": "libmkl_rt", - "user_api": "blas", - "internal_api": "mkl", - "version": "2019.0.4", - "num_threads": 2, - "threading_layer": "intel" - }, - { - "filepath": "/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so", - "prefix": "libiomp", - "user_api": "openmp", - "internal_api": "openmp", - "version": null, - "num_threads": 4 - } -] -``` - -The JSON information is written on STDOUT. If some of the packages are missing, -a warning message is displayed on STDERR. - -### Python Runtime Programmatic Introspection - -Introspect the current state of the threadpool-enabled runtime libraries -that are loaded when importing Python packages: - -```python ->>> from threadpoolctl import threadpool_info ->>> from pprint import pprint ->>> pprint(threadpool_info()) -[] - ->>> import numpy ->>> pprint(threadpool_info()) -[{'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so', - 'internal_api': 'mkl', - 'num_threads': 2, - 'prefix': 'libmkl_rt', - 'threading_layer': 'intel', - 'user_api': 'blas', - 'version': '2019.0.4'}, - {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so', - 'internal_api': 'openmp', - 'num_threads': 4, - 'prefix': 'libiomp', - 'user_api': 'openmp', - 'version': None}] - ->>> import xgboost ->>> pprint(threadpool_info()) -[{'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so', - 'internal_api': 'mkl', - 'num_threads': 2, - 'prefix': 'libmkl_rt', - 'threading_layer': 'intel', - 'user_api': 'blas', - 'version': '2019.0.4'}, - {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so', - 'internal_api': 'openmp', - 'num_threads': 4, - 'prefix': 'libiomp', - 'user_api': 'openmp', - 'version': None}, - {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libgomp.so.1.0.0', - 'internal_api': 'openmp', - 'num_threads': 4, - 'prefix': 'libgomp', - 'user_api': 'openmp', - 'version': None}] -``` - -In the above example, `numpy` was installed from the default anaconda channel and comes -with MKL and its Intel OpenMP (`libiomp5`) implementation while `xgboost` was installed -from pypi.org and links against GNU OpenMP (`libgomp`) so both OpenMP runtimes are -loaded in the same Python program. - -The state of these libraries is also accessible through the object oriented API: - -```python ->>> from threadpoolctl import ThreadpoolController, threadpool_info ->>> from pprint import pprint ->>> import numpy ->>> controller = ThreadpoolController() ->>> pprint(controller.info()) -[{'architecture': 'Haswell', - 'filepath': '/home/jeremie/miniconda/envs/dev/lib/libopenblasp-r0.3.17.so', - 'internal_api': 'openblas', - 'num_threads': 4, - 'prefix': 'libopenblas', - 'threading_layer': 'pthreads', - 'user_api': 'blas', - 'version': '0.3.17'}] - ->>> controller.info() == threadpool_info() -True -``` - -## Usage when not using Python threads: Restricting Controlled Library Thread Pool Sizes - +# Thread-pool Controls [![Build Status](https://github.com/joblib/threadpoolctl/actions/workflows/test.yml/badge.svg?branch=master)](https://github.com/joblib/threadpoolctl/actions?query=branch%3Amaster) [![codecov](https://codecov.io/gh/joblib/threadpoolctl/branch/master/graph/badge.svg)](https://codecov.io/gh/joblib/threadpoolctl) + +Python helpers to limit the number of threads used in the +threadpool-backed of common native libraries used for scientific +computing and data science (e.g. BLAS and OpenMP). + +Fine control of the underlying thread-pool size can be useful in +workloads that involve nested parallelism so as to mitigate +oversubscription issues. + +Note that "limiting the number of threads" in practice involves a variety of +potential semantics depending on the underlying third-party library; see the +section on [semantics](#semantics) below. + +## Installation + +- For users, install the last published version from PyPI: + + ```bash + pip install threadpoolctl + ``` + +- For contributors, install from the source repository in developer + mode: + + ```bash + pip install -r dev-requirements.txt + flit install --symlink + ``` + + then you run the tests with pytest: + + ```bash + pytest + ``` + +## Usage: Introspection and debugging + +### Command Line Interface + +Get a JSON description of thread-pools initialized when importing python +packages such as numpy or scipy for instance: + +``` +python -m threadpoolctl -i numpy scipy.linalg +[ + { + "filepath": "/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so", + "prefix": "libmkl_rt", + "user_api": "blas", + "internal_api": "mkl", + "version": "2019.0.4", + "num_threads": 2, + "threading_layer": "intel" + }, + { + "filepath": "/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so", + "prefix": "libiomp", + "user_api": "openmp", + "internal_api": "openmp", + "version": null, + "num_threads": 4 + } +] +``` + +The JSON information is written on STDOUT. If some of the packages are missing, +a warning message is displayed on STDERR. + +### Python Runtime Programmatic Introspection + +Introspect the current state of the threadpool-enabled runtime libraries +that are loaded when importing Python packages: + +```python +>>> from threadpoolctl import threadpool_info +>>> from pprint import pprint +>>> pprint(threadpool_info()) +[] + +>>> import numpy +>>> pprint(threadpool_info()) +[{'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so', + 'internal_api': 'mkl', + 'num_threads': 2, + 'prefix': 'libmkl_rt', + 'threading_layer': 'intel', + 'user_api': 'blas', + 'version': '2019.0.4'}, + {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so', + 'internal_api': 'openmp', + 'num_threads': 4, + 'prefix': 'libiomp', + 'user_api': 'openmp', + 'version': None}] + +>>> import xgboost +>>> pprint(threadpool_info()) +[{'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libmkl_rt.so', + 'internal_api': 'mkl', + 'num_threads': 2, + 'prefix': 'libmkl_rt', + 'threading_layer': 'intel', + 'user_api': 'blas', + 'version': '2019.0.4'}, + {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libiomp5.so', + 'internal_api': 'openmp', + 'num_threads': 4, + 'prefix': 'libiomp', + 'user_api': 'openmp', + 'version': None}, + {'filepath': '/home/ogrisel/miniconda3/envs/tmp/lib/libgomp.so.1.0.0', + 'internal_api': 'openmp', + 'num_threads': 4, + 'prefix': 'libgomp', + 'user_api': 'openmp', + 'version': None}] +``` + +In the above example, `numpy` was installed from the default anaconda channel and comes +with MKL and its Intel OpenMP (`libiomp5`) implementation while `xgboost` was installed +from pypi.org and links against GNU OpenMP (`libgomp`) so both OpenMP runtimes are +loaded in the same Python program. + +The state of these libraries is also accessible through the object oriented API: + +```python +>>> from threadpoolctl import ThreadpoolController, threadpool_info +>>> from pprint import pprint +>>> import numpy +>>> controller = ThreadpoolController() +>>> pprint(controller.info()) +[{'architecture': 'Haswell', + 'filepath': '/home/jeremie/miniconda/envs/dev/lib/libopenblasp-r0.3.17.so', + 'internal_api': 'openblas', + 'num_threads': 4, + 'prefix': 'libopenblas', + 'threading_layer': 'pthreads', + 'user_api': 'blas', + 'version': '0.3.17'}] + +>>> controller.info() == threadpool_info() +True +``` + +## Usage when not using Python threads: Restricting Controlled Library Thread Pool Sizes + There are two scenarios in which you might want to use `threadpoolctl`; each -requires you to use different APIs. - -1. You do not expect to use any Python threads, so all the work will be started - directly from the main thread in the process. This is a simple case - where we can globally set thread limits. -2. You will be parallelizing work using a Python thread pool, and your goal is - therefore to limit controlled libraries' thread pool sizes when - concurrently called from Python threads. This case is a bit more - complex to handle properly and requires a bit more verbose code. - -This section will cover the former case, and the latter is covered in the next -usage section. - -### Setting the Maximum Size of Thread-Pools - -Control the number of threads used by the underlying runtime libraries -in specific sections of your Python program: - -```python ->>> from threadpoolctl import threadpool_limits ->>> import numpy as np - ->>> with threadpool_limits(limits=1, user_api='blas'): -... # In this block, calls to blas implementation (like openblas or MKL) -... # will be limited to use only one thread. They can thus be used jointly -... # with thread-parallelism. -... a = np.random.randn(1000, 1000) -... a_squared = a @ a -``` - -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. 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 ->>> import numpy as np ->>> controller = ThreadpoolController() - ->>> with controller.limit(limits=1, user_api='blas'): -... a = np.random.randn(1000, 1000) -... 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 -the maximum number of threads used by the supported libraries at a function level. The -decorators are accessible through their `wrap` method. - -```python ->>> from threadpoolctl import ThreadpoolController, threadpool_limits ->>> import numpy as np ->>> controller = ThreadpoolController() - ->>> @controller.wrap(limits=1, user_api='blas') -... # or @threadpool_limits.wrap(limits=1, user_api='blas') -... def my_func(): -... # Inside this function, calls to blas implementation (like openblas or MKL) -... # will be limited to use only one thread. -... a = np.random.randn(1000, 1000) -... a_squared = a @ a -... -``` - -## Usage for Python threads: Restricting Controlled Library Thread Pool Sizes - -This section covers APIs to use when you will be using Python thread pools to parallelize work. -The usage is more complicated than one might expect because the underlying APIs have a variety of different semantics, as [explained later in the docs](#semantics). - -### Setting the Maximum Size of Thread-Pools, When Python Thread Pools Are Used - -Limiting thread pool size in controlled libraries requires a two-step process. -**Importantly, each Python worker thread must also call a method to limit -controlled libraries in that thread.** With Python's -`concurrent.futures.ThreadPoolExecutor`, you can do so by passing in an -initializer function that will get called on thread startup. - -Whenever `threadpool_limits` is called, it creates a new `ThreadpoolController`, -which needs to do some work (inspecting and getting access to third-party shared -libraries) that can take some time. To prevent the performance cost of doing -this work in all the threads, you can reuse a `ThreadpoolController` object -across the threads. - -When setting limits, you want to set limits both in the main thread (in case you're using -OpenBLAS+pthreads) and in each worker thread. - -> **Why you need both:** When setting a limit on OpenBLAS with pthreads or the -> Windows threading layer, setting the limit affects the whole process, because -> it uses a process-wide shared thread pool. (This in contrast to some other -> BLAS/OpenMP libraries, or OpenBLAS backed by OpenMP, that allow setting the -> limit per-thread.) Having multiple Python threads set limits in this case -> makes no sense, since they would effectively be arguing over a globally shared -> resource. ``threadpoolctl`` therefore constrains the limiting API in the -> OpenBLAS pthreads/Windows case to only work from the first thread that uses -> it, the so-called controlling thread. -> -> For many other libraries, for example OpenMP, the limit is set per thread. As -> a result, you want to set limits both ways in generic code that doesn't know -> the underlying BLAS library: in main thread for the process-wide limits, and -> per-thread for thread-specific limits like OpenMP. - -```python -from threadpoolctl import ThreadpoolController - -# This won't have any side-effects. Because it caches its list of loaded -# libraries, you need to create a new one if you've imported or loaded any -# relevant libraries in the interim. So storing this on module level may not be -# a good idea if you e.g. only do `import numpy` later on. -controller = ThreadpoolController() - -with ( - # Limit on main thread. This is necessary because some underlying limiting - # APIs operate on a process-wide basis. Because it's a context manager, it - # will restore the limits after the ``with`` is done. - controller.limit(limits=1), - # Make sure each Python worker thread also calls threadpool_limits(). If - # you're using another thread pool class, you will need to do so some other - # way. - ThreadPoolExecutor(4, initializer=lambda: controller.limit(limits=1)) as pool, -): - # ... run some BLAS-using code in the thread pool ... - pool.map(somefunc, someargs) - -# Later... -controller = ThreadpoolController() -with ( - controller.limit(limits=2), - ThreadPoolExecutor(4, initializer=lambda: controller.limit(limits=2)) as pool, -): - # ... run some BLAS-using code in the thread pool ... - pool.map(somefunc, someargs) -``` - -You can also operate without a context manager: - -```python -from threadpoolctl import ThreadpoolController - -controller = ThreadpoolController() -try: - limiter = controller.limit(limits=1) - with ThreadPoolExecutor(4, initializer=lambda: controller.limit(limits=1)) as pool: - # ... run some BLAS-using code in the thread pool ... - pool.map(somefunc, someargs) -finally: - limiter.restore_original_limits() -``` - -### Switching Back And Forth Between Main Thread and Python Threads - -Unfortunately not all controlled libraries providing limiting APIs that are -thread-specific, as detailed in the section on [semantics](#semantics) below. -Limiting some libraries' thread pool sizes can therefore impact the whole -process. This makes switching back and forth between running code that -uses these libraries in Python threads and running it in the main thread a bit -more complex: you need to set the limits each time you switch back and forth. - -Let's say your computer has 4 cores, and you're using some OpenMP API. - -```python -POOL = ThreadPoolExecutor(4) -CONTROLLER = ThreadpoolController() - -# 1. Run some work in a Python thread pool (OpenMP effectively disabled). -with CONTROLLER.limit(limits=1): - - def limit_then_do_work(*args, **kwargs): - # Set a limit on OpenMP in the current thread: - CONTROLLER.limit(limits=1) - # Do the actual work: - return do_real_work_with_openmp(*args, **kwargs) - - results = POOL.map(limit_then_do_work, args) - - -# 2. Run some work with 4-threads OpenMP parallelism: -with CONTROLLER.limit(limits=4): - results2 = do_more_work_with_openmp(results) - - -# 3. Nest some OpenMP parallelism under Python-level parallelism: -with CONTROLLER.limit(limits=2): - - def limit_then_do_work2(*args, **kwargs): - CONTROLLER.limit(limits=2) - return do_even_more_real_work_with_openmp(*args, **kwargs) - - results3 = POOL.map(limit_then_do_work2, results2) -``` - -## Usage: Additional APIs and details - -### OpenBLAS with pthreads/Windows - -As mentioned above, OpenBLAS with pthreads or Windows threading layer uses a -process-wide shared thread pool. As a result, there are only two limits that -make sense in most circumstances: - -* A limit of `1`, which effectively means no OpenBLAS thread pool: all - operations run inline in the calling thread. This is a good setting when you - have a Python thread pool calling BLAS, and the number of Python threads is - close to the number of available cores. -* A limit equal to the number of available cores. This is a good setting when - you are calling BLAS from the main thread and there are no Python threads. - -While `threadpoolctl` does not limit you to these two values, you may wish to do -so yourself in your code. - -### Switching the FlexiBLAS backend - -`FlexiBLAS` is a BLAS wrapper for which the BLAS backend can be switched at runtime. -`threadpoolctl` exposes python bindings for this feature. Here's an example but note -that this part of the API is experimental and subject to change without deprecation: - -```python ->>> from threadpoolctl import ThreadpoolController ->>> import numpy as np ->>> controller = ThreadpoolController() - ->>> controller.info() -[{'user_api': 'blas', - 'internal_api': 'flexiblas', - 'num_threads': 1, - 'prefix': 'libflexiblas', - 'filepath': '/usr/local/lib/libflexiblas.so.3.3', - 'version': '3.3.1', - 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], - 'loaded_backends': ['NETLIB'], - 'current_backend': 'NETLIB'}] - -# Retrieve the flexiblas controller ->>> flexiblas_ct = controller.select(internal_api="flexiblas").lib_controllers[0] - -# Switch the backend with one predefined at build time (listed in "available_backends") ->>> flexiblas_ct.switch_backend("OPENBLASPTHREAD") ->>> controller.info() -[{'user_api': 'blas', - 'internal_api': 'flexiblas', - 'num_threads': 4, - 'prefix': 'libflexiblas', - 'filepath': '/usr/local/lib/libflexiblas.so.3.3', - 'version': '3.3.1', - 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], - 'loaded_backends': ['NETLIB', 'OPENBLASPTHREAD'], - 'current_backend': 'OPENBLASPTHREAD'}, - {'user_api': 'blas', - 'internal_api': 'openblas', - 'num_threads': 4, - 'prefix': 'libopenblas', - 'filepath': '/usr/lib/x86_64-linux-gnu/openblas-pthread/libopenblasp-r0.3.8.so', - 'version': '0.3.8', - 'threading_layer': 'pthreads', - 'architecture': 'Haswell'}] - -# It's also possible to directly give the path to a shared library ->>> flexiblas_controller.switch_backend("/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so") ->>> controller.info() -[{'user_api': 'blas', - 'internal_api': 'flexiblas', - 'num_threads': 2, - 'prefix': 'libflexiblas', - 'filepath': '/usr/local/lib/libflexiblas.so.3.3', - 'version': '3.3.1', - 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], - 'loaded_backends': ['NETLIB', - 'OPENBLASPTHREAD', - '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so'], - 'current_backend': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so'}, - {'user_api': 'openmp', - 'internal_api': 'openmp', - 'num_threads': 4, - 'prefix': 'libomp', - 'filepath': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libomp.so', - 'version': None}, - {'user_api': 'blas', - 'internal_api': 'openblas', - 'num_threads': 4, - 'prefix': 'libopenblas', - 'filepath': '/usr/lib/x86_64-linux-gnu/openblas-pthread/libopenblasp-r0.3.8.so', - 'version': '0.3.8', - 'threading_layer': 'pthreads', - 'architecture': 'Haswell'}, - {'user_api': 'blas', - 'internal_api': 'mkl', - 'num_threads': 2, - 'prefix': 'libmkl_rt', - 'filepath': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so.2', - 'version': '2024.0-Product', - 'threading_layer': 'gnu'}] -``` - -You can observe that the previously linked OpenBLAS shared object stays loaded by -the Python program indefinitely, but FlexiBLAS itself no longer delegates BLAS calls -to OpenBLAS as indicated by the `current_backend` attribute. - -### Writing a custom library controller - -Currently, `threadpoolctl` has support for `OpenMP` and the main `BLAS` libraries. -However it can also be used to control the threadpool of other native libraries, -provided that they expose an API to get and set the limit on the number of threads. -For that, one must implement a controller for this library and register it to -`threadpoolctl`. - -A custom controller must be a subclass of the `LibController` class and implement -the attributes and methods described in the docstring of `LibController`. Then this -new controller class must be registered using the `threadpoolctl.register` function. -An complete example can be found [here]( - https://github.com/joblib/threadpoolctl/blob/master/tests/_pyMylib/__init__.py). - -### Sequential BLAS within OpenMP parallel region - -When one wants to have sequential BLAS calls within an OpenMP parallel region, it's -necessary to set `limits="sequential_blas_under_openmp"` since setting `limits=1` and -`user_api="blas"` leads to unexpected behavior for OpenBLAS with the OpenMP threading -layer: OpenBLAS and the calling code share the same OpenMP runtime, so it also sets -the number of OpenMP threads of the current thread to 1. Inside the limit, any OpenMP -parallel region started from this thread without an explicit number of threads runs -on a single thread, and `omp_get_max_threads()` returns 1. - -```python -from threadpoolctl import ThreadpoolController - -controller = ThreadpoolController() - -with controller.limit(limits="sequential_blas_under_openmp"): - # Compiled code with an OpenMP parallel loop that calls BLAS. - my_openmp_kernel(X) -``` - -### Known Limitations - -- `threadpool_limits` can fail to limit the number of inner threads when nesting - parallel loops managed by distinct OpenMP runtime implementations (for instance - libgomp from GCC and libomp from clang/llvm or libiomp from ICC). - - See the `test_openmp_nesting` function in [tests/test_threadpoolctl.py]( - https://github.com/joblib/threadpoolctl/blob/master/tests/test_threadpoolctl.py) - for an example. More information can be found at: - https://github.com/jeremiedbb/Nested_OpenMP - - Note however that this problem does not happen when `threadpool_limits` is - used to limit the number of threads used internally by BLAS calls that are - themselves nested under OpenMP parallel loops. `threadpool_limits` works as - expected, even if the inner BLAS implementation relies on a distinct OpenMP - implementation. - -- Using Intel OpenMP (ICC) and LLVM OpenMP (clang) in the same Python program - under Linux is known to cause problems. See the following guide for more details - and workarounds: - https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md - -- Setting the maximum number of threads of the OpenMP and BLAS libraries has - inconsistent scope and semantics (thread-local vs process-wide) depending on - the underlying library. For more details see - https://github.com/joblib/threadpoolctl/issues/208 - - For example, if you're using OpenMP with libgomp (gcc) or libomp (clang), the - setting is thread-local and sets how many OpenMP threads will be started in - the current thread. On the other hand, with OpenBLAS with pthreads backend or - on Windows, the setting is process-wide and impacts the size of a process-wide - thread pool shared across all threads in the process. - +requires you to use different APIs. + +1. You do not expect to use any Python threads, so all the work will be started + directly from the main thread in the process. This is a simple case + where we can globally set thread limits. +2. You will be parallelizing work using a Python thread pool, and your goal is + therefore to limit controlled libraries' thread pool sizes when + concurrently called from Python threads. This case is a bit more + complex to handle properly and requires a bit more verbose code. + +This section will cover the former case, and the latter is covered in the next +usage section. + +### Setting the Maximum Size of Thread-Pools + +Control the number of threads used by the underlying runtime libraries +in specific sections of your Python program: + +```python +>>> from threadpoolctl import threadpool_limits +>>> import numpy as np + +>>> with threadpool_limits(limits=1, user_api='blas'): +... # In this block, calls to blas implementation (like openblas or MKL) +... # will be limited to use only one thread. They can thus be used jointly +... # with thread-parallelism. +... a = np.random.randn(1000, 1000) +... a_squared = a @ a +``` + +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. 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 +>>> import numpy as np +>>> controller = ThreadpoolController() + +>>> with controller.limit(limits=1, user_api='blas'): +... a = np.random.randn(1000, 1000) +... 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 +the maximum number of threads used by the supported libraries at a function level. The +decorators are accessible through their `wrap` method. + +```python +>>> from threadpoolctl import ThreadpoolController, threadpool_limits +>>> import numpy as np +>>> controller = ThreadpoolController() + +>>> @controller.wrap(limits=1, user_api='blas') +... # or @threadpool_limits.wrap(limits=1, user_api='blas') +... def my_func(): +... # Inside this function, calls to blas implementation (like openblas or MKL) +... # will be limited to use only one thread. +... a = np.random.randn(1000, 1000) +... a_squared = a @ a +... +``` + +## Usage for Python threads: Restricting Controlled Library Thread Pool Sizes + +This section covers APIs to use when you will be using Python thread pools to parallelize work. +The usage is more complicated than one might expect because the underlying APIs have a variety of different semantics, as [explained later in the docs](#semantics). + +### Setting the Maximum Size of Thread-Pools, When Python Thread Pools Are Used + +Limiting thread pool size in controlled libraries requires a two-step process. +**Importantly, each Python worker thread must also call a method to limit +controlled libraries in that thread.** With Python's +`concurrent.futures.ThreadPoolExecutor`, you can do so by passing in an +initializer function that will get called on thread startup. + +Whenever `threadpool_limits` is called, it creates a new `ThreadpoolController`, +which needs to do some work (inspecting and getting access to third-party shared +libraries) that can take some time. To prevent the performance cost of doing +this work in all the threads, you can reuse a `ThreadpoolController` object +across the threads. + +When setting limits, you want to set limits both in the main thread (in case you're using +OpenBLAS+pthreads) and in each worker thread. + +> **Why you need both:** When setting a limit on OpenBLAS with pthreads or the +> Windows threading layer, setting the limit affects the whole process, because +> it uses a process-wide shared thread pool. (This in contrast to some other +> BLAS/OpenMP libraries, or OpenBLAS backed by OpenMP, that allow setting the +> limit per-thread.) +> +> For many other libraries, for example OpenMP, the limit is set per thread. As +> a result, you want to set limits both ways in generic code that doesn't know +> the underlying BLAS library: in main thread for the process-wide limits, and +> per-thread for thread-specific limits like OpenMP. + +```python +from threadpoolctl import ThreadpoolController + +# This won't have any side-effects. Because it caches its list of loaded +# libraries, you need to create a new one if you've imported or loaded any +# relevant libraries in the interim. So storing this on module level may not be +# a good idea if you e.g. only do `import numpy` later on. +controller = ThreadpoolController() + +with ( + # Limit on main thread. This is necessary because some underlying limiting + # APIs operate on a process-wide basis. Because it's a context manager, it + # will restore the limits after the ``with`` is done. + controller.limit(limits=1), + # Make sure each Python worker thread also calls threadpool_limits(). If + # you're using another thread pool class, you will need to do so some other + # way. + ThreadPoolExecutor(4, initializer=lambda: controller.limit(limits=1)) as pool, +): + # ... run some BLAS-using code in the thread pool ... + pool.map(somefunc, someargs) + +# Later... +controller = ThreadpoolController() +with ( + controller.limit(limits=2), + ThreadPoolExecutor(4, initializer=lambda: controller.limit(limits=2)) as pool, +): + # ... run some BLAS-using code in the thread pool ... + pool.map(somefunc, someargs) +``` + +You can also operate without a context manager: + +```python +from threadpoolctl import ThreadpoolController + +controller = ThreadpoolController() +try: + limiter = controller.limit(limits=1) + with ThreadPoolExecutor(4, initializer=lambda: controller.limit(limits=1)) as pool: + # ... run some BLAS-using code in the thread pool ... + pool.map(somefunc, someargs) +finally: + limiter.restore_original_limits() +``` + +### Switching Back And Forth Between Main Thread and Python Threads + +Unfortunately not all controlled libraries providing limiting APIs that are +thread-specific, as detailed in the section on [semantics](#semantics) below. +Limiting some libraries' thread pool sizes can therefore impact the whole +process. This makes switching back and forth between running code that +uses these libraries in Python threads and running it in the main thread a bit +more complex: you need to set the limits each time you switch back and forth. + +Let's say your computer has 4 cores, and you're using some OpenMP API. + +```python +POOL = ThreadPoolExecutor(4) +CONTROLLER = ThreadpoolController() + +# 1. Run some work in a Python thread pool (OpenMP effectively disabled). +with CONTROLLER.limit(limits=1): + + def limit_then_do_work(*args, **kwargs): + # Set a limit on OpenMP in the current thread: + CONTROLLER.limit(limits=1) + # Do the actual work: + return do_real_work_with_openmp(*args, **kwargs) + + results = POOL.map(limit_then_do_work, args) + + +# 2. Run some work with 4-threads OpenMP parallelism: +with CONTROLLER.limit(limits=4): + results2 = do_more_work_with_openmp(results) + + +# 3. Nest some OpenMP parallelism under Python-level parallelism: +with CONTROLLER.limit(limits=2): + + def limit_then_do_work2(*args, **kwargs): + CONTROLLER.limit(limits=2) + return do_even_more_real_work_with_openmp(*args, **kwargs) + + results3 = POOL.map(limit_then_do_work2, results2) +``` + +## Usage: Additional APIs and details + +### OpenBLAS with pthreads/Windows, and the controlling thread + +As mentioned above, OpenBLAS with pthreads or Windows threading layer uses a +process-wide shared thread pool. Having multiple Python threads set limits in this case +makes no sense, since they would effectively be arguing over a globally shared +resource. + +``threadpoolctl`` therefore constrains the limiting API in the OpenBLAS +pthreads/Windows case to only work from the first thread that uses it, the +so-called controlling thread. The controlling thread can change the limits as +normal. For other thread, calls will be silently ignored. + +When using `with threadpool_limits()` as a context manager or +`@threadpool_limits.wrap()` as a decorator, the controlling thread should be +automatically cleared when the code leaves the relevant block. This allows +another thread to take control. You can also manually release the controlling +thread by calling `OpenBLASController.clear_controlling_thread()` from the +controlling thread; passing `force=True` allows clearing it from any thread. + +Note that if you call a standalone `threadpool_limits()` or similar API, and +then later use the `with threadpool_limits()` or similar context +manager/decorator API, clearing the controlling thread will still happen, as +standalone limiting calls are not tracked. + +NumPy and SciPy, when installed from PyPI for example, can each have their own +separate OpenBLAS with its own thread pool. The controlling thread will apply to +both. + +#### Choosing limits for a process-wide shared thread pool + +Given a shared OpenBLAS process pool, logically there are only two limits that +make sense in most circumstances: + +* A limit of `1`, which effectively means no OpenBLAS thread pool: all + operations run inline in the calling thread. This is a good setting when you + have a Python thread pool calling BLAS, and the number of Python threads is + close to the number of available cores. +* A limit equal to the number of available cores. This is a good setting when + you are calling BLAS from the main thread and there are no Python threads. + +While `threadpoolctl` does not limit you to these two values, you may wish to do +so yourself in your code, after benchmarking to ensure that this actually +maximizes performance. Benchmarking is important because empirically, these +limits do not always maximize performance. + +### Switching the FlexiBLAS backend + +`FlexiBLAS` is a BLAS wrapper for which the BLAS backend can be switched at runtime. +`threadpoolctl` exposes python bindings for this feature. Here's an example but note +that this part of the API is experimental and subject to change without deprecation: + +```python +>>> from threadpoolctl import ThreadpoolController +>>> import numpy as np +>>> controller = ThreadpoolController() + +>>> controller.info() +[{'user_api': 'blas', + 'internal_api': 'flexiblas', + 'num_threads': 1, + 'prefix': 'libflexiblas', + 'filepath': '/usr/local/lib/libflexiblas.so.3.3', + 'version': '3.3.1', + 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], + 'loaded_backends': ['NETLIB'], + 'current_backend': 'NETLIB'}] + +# Retrieve the flexiblas controller +>>> flexiblas_ct = controller.select(internal_api="flexiblas").lib_controllers[0] + +# Switch the backend with one predefined at build time (listed in "available_backends") +>>> flexiblas_ct.switch_backend("OPENBLASPTHREAD") +>>> controller.info() +[{'user_api': 'blas', + 'internal_api': 'flexiblas', + 'num_threads': 4, + 'prefix': 'libflexiblas', + 'filepath': '/usr/local/lib/libflexiblas.so.3.3', + 'version': '3.3.1', + 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], + 'loaded_backends': ['NETLIB', 'OPENBLASPTHREAD'], + 'current_backend': 'OPENBLASPTHREAD'}, + {'user_api': 'blas', + 'internal_api': 'openblas', + 'num_threads': 4, + 'prefix': 'libopenblas', + 'filepath': '/usr/lib/x86_64-linux-gnu/openblas-pthread/libopenblasp-r0.3.8.so', + 'version': '0.3.8', + 'threading_layer': 'pthreads', + 'architecture': 'Haswell'}] + +# It's also possible to directly give the path to a shared library +>>> flexiblas_controller.switch_backend("/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so") +>>> controller.info() +[{'user_api': 'blas', + 'internal_api': 'flexiblas', + 'num_threads': 2, + 'prefix': 'libflexiblas', + 'filepath': '/usr/local/lib/libflexiblas.so.3.3', + 'version': '3.3.1', + 'available_backends': ['NETLIB', 'OPENBLASPTHREAD', 'ATLAS'], + 'loaded_backends': ['NETLIB', + 'OPENBLASPTHREAD', + '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so'], + 'current_backend': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so'}, + {'user_api': 'openmp', + 'internal_api': 'openmp', + 'num_threads': 4, + 'prefix': 'libomp', + 'filepath': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libomp.so', + 'version': None}, + {'user_api': 'blas', + 'internal_api': 'openblas', + 'num_threads': 4, + 'prefix': 'libopenblas', + 'filepath': '/usr/lib/x86_64-linux-gnu/openblas-pthread/libopenblasp-r0.3.8.so', + 'version': '0.3.8', + 'threading_layer': 'pthreads', + 'architecture': 'Haswell'}, + {'user_api': 'blas', + 'internal_api': 'mkl', + 'num_threads': 2, + 'prefix': 'libmkl_rt', + 'filepath': '/home/jeremie/miniforge/envs/flexiblas_threadpoolctl/lib/libmkl_rt.so.2', + 'version': '2024.0-Product', + 'threading_layer': 'gnu'}] +``` + +You can observe that the previously linked OpenBLAS shared object stays loaded by +the Python program indefinitely, but FlexiBLAS itself no longer delegates BLAS calls +to OpenBLAS as indicated by the `current_backend` attribute. + +### Writing a custom library controller + +Currently, `threadpoolctl` has support for `OpenMP` and the main `BLAS` libraries. +However it can also be used to control the threadpool of other native libraries, +provided that they expose an API to get and set the limit on the number of threads. +For that, one must implement a controller for this library and register it to +`threadpoolctl`. + +A custom controller must be a subclass of the `LibController` class and implement +the attributes and methods described in the docstring of `LibController`. Then this +new controller class must be registered using the `threadpoolctl.register` function. +An complete example can be found [here]( + https://github.com/joblib/threadpoolctl/blob/master/tests/_pyMylib/__init__.py). + +### Sequential BLAS within OpenMP parallel region + +When one wants to have sequential BLAS calls within an OpenMP parallel region, it's +necessary to set `limits="sequential_blas_under_openmp"` since setting `limits=1` and +`user_api="blas"` leads to unexpected behavior for OpenBLAS with the OpenMP threading +layer: OpenBLAS and the calling code share the same OpenMP runtime, so it also sets +the number of OpenMP threads of the current thread to 1. Inside the limit, any OpenMP +parallel region started from this thread without an explicit number of threads runs +on a single thread, and `omp_get_max_threads()` returns 1. + +```python +from threadpoolctl import ThreadpoolController + +controller = ThreadpoolController() + +with controller.limit(limits="sequential_blas_under_openmp"): + # Compiled code with an OpenMP parallel loop that calls BLAS. + my_openmp_kernel(X) +``` + +### Known Limitations + +- `threadpool_limits` can fail to limit the number of inner threads when nesting + parallel loops managed by distinct OpenMP runtime implementations (for instance + libgomp from GCC and libomp from clang/llvm or libiomp from ICC). + + See the `test_openmp_nesting` function in [tests/test_threadpoolctl.py]( + https://github.com/joblib/threadpoolctl/blob/master/tests/test_threadpoolctl.py) + for an example. More information can be found at: + https://github.com/jeremiedbb/Nested_OpenMP + + Note however that this problem does not happen when `threadpool_limits` is + used to limit the number of threads used internally by BLAS calls that are + themselves nested under OpenMP parallel loops. `threadpool_limits` works as + expected, even if the inner BLAS implementation relies on a distinct OpenMP + implementation. + +- Using Intel OpenMP (ICC) and LLVM OpenMP (clang) in the same Python program + under Linux is known to cause problems. See the following guide for more details + and workarounds: + https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md + +- Setting the maximum number of threads of the OpenMP and BLAS libraries has + inconsistent scope and semantics (thread-local vs process-wide) depending on + the underlying library. For more details see + https://github.com/joblib/threadpoolctl/issues/208 + + For example, if you're using OpenMP with libgomp (gcc) or libomp (clang), the + setting is thread-local and sets how many OpenMP threads will be started in + the current thread. On the other hand, with OpenBLAS with pthreads backend or + on Windows, the setting is process-wide and impacts the size of a process-wide + thread pool shared across all threads in the process. + ##
Semantics of thread limiting
- -Setting the number of threads may seem like a simple operation, but in practice -it can do quite different things depending on the underlying third-party library -being limited. - -### Kind of thread pool - -* In some libraries, there is a shared process-wide thread pool. Setting the - number of threads changes the size of this shared pool. -* In other libraries, there is a thread pool per calling thread. So each Python - thread gets its own personal thread pool from the third-party library. - -To give a concrete example: if you have 10 cores, and you're using OpenBLAS with -the `pthreads` backend on Linux, by default there is a single 10-worker pool of -threads created by OpenBLAS. If you start 10 Python threads, each calling BLAS -routines, all those Python threads will feed into that single 10-thread pool. -In total, you will have 20 threads running (10 Python, 10 OpenBLAS). - -On the other hand, if you use MKL, each Python thread gets its own individual -pool of worker threads from MKL, so if each Python threads calls a BLAS routine -in MKL, you will get 10×10 = 100 MKL threads! - -### Scope of limiting - -If you have a third-party library that creates a thread pool per calling thread, -there is another question about what limiting the number of threads means: which -threads are affected by the limiting API? - -* **All threads in the process:** In this case, setting the limits changes it - for all threads in the process. So if you limit from the main Python thread, - all other Python threads will have the same limit. -* **Current thread only:** Only the current thread's thread pool size will be - limited. - -MKL for example has two APIs, covering both cases, -[`mkl_set_num_threads()`](https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2026-0/mkl-set-num-threads.html) -and [`mkl_set_num_threads_local()`](https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2026-0/mkl-set-num-threads-local.html) respectively. - -### `threadpoolctl`'s policy for thread limiting - -When there is a choice between different APIs, `threadpoolctl` will prefer APIs -that: - -1. Have a per-thread worker pool (so each Python thread gets its own pool from - the underlying third-party library) -2. Only affect the current thread. - -Current libraries where `threadpoolctl` explicitly makes this choice are: - -* MKL, for all threading backends. -* OpenBLAS (v0.3.34 or later) when using the OpenMP backend, on Linux and macOS. - threadpoolctl then limits OpenBLAS with `omp_set_num_threads`, so the limit - also applies to the OpenMP parallel regions of the calling code. To get - sequential BLAS calls within an OpenMP parallel region, see - [Sequential BLAS within OpenMP parallel region](#sequential-blas-within-openmp-parallel-region). - -When using OpenMP, this is also the default on Linux and macOS. On Windows -OpenMP has a per-thread worker pool but the limiting API affects all threads in -the process, not just the current one. - -### Checking the behavior of your installed libraries - -When you run `python -m threadpoolctl` it will include the scope of the API -limit in the `"thread_limit_scope"` field, with the value of `"process"` -indicating the API affects the whole process and `"current_thread"` indicating a -per-thread thread pool. Information about the kind of limiting in the former -case (shared process-wide pool, or per-thread pool) is not included. - -Here is example output for OpenBLAS with pthreads: - -```shell-session -$ python -m threadpoolctl -i numpy -[ - { - "user_api": "blas", - "internal_api": "openblas", - "num_threads": 12, - "prefix": "libscipy_openblas", - "version": "0.3.33.112.0", - "threading_layer": "pthreads", - "architecture": "Haswell", - "thread_limit_scope": "process" - } -] -``` - -You can also use another script to empirically determine both the kind and scope of the API. From a checkout of the [`threadpoolctl` GitHub repo](https://github.com/joblib/threadpoolctl): - -```shell-session -$ python -m tests.empirical_scope_observation blas -== Observed behavior: blas == -10x12 Python threads each running a parallel workload with a 2 thread limit. -Maximum number of observed threads (includes Python and blas threads): 53 -Presumed API scope: process-wide shared thread pool -``` - -You can also call this script with an `openmp` argument instead of `blas`. - -## Maintainers - -To make a release: - -- Create a PR to bump the version number (`__version__`) in `threadpoolctl.py` and - update the release date in `CHANGES.md`. - -- Merge the PR and check that the `Publish threadpoolctl distribution to TestPyPI` job - of the `publish-to-pypi.yml` workflow successfully uploaded the wheel and source - distribution to Test PyPI. - -- If everything is fine create a tag for the release and push it to github: - - ```bash - git tag -a X.Y.Z - git push git@github.com:joblib/threadpoolctl.git X.Y.Z - ``` - -- Check that the `Publish threadpoolctl distribution to PyPI` job of the - `publish-to-pypi.yml` workflow successfully uploaded the wheel and source distribution - to PyPI this time. - -- Create a PR for the release on the [conda-forge feedstock](https://github.com/conda-forge/threadpoolctl-feedstock) (or wait for the bot to make it). - -- Publish the release on github. - -If for some reason the steps above can't be achieved and a munual upload of the wheel -and source distribution is needed: - -- Build the distribution archives: - - ```bash - pip install flit - flit build - ``` - -- Upload the wheels and source distribution to PyPI using flit. Since PyPI doesn't - allow password authentication anymore, the username needs to be changed to the - generic name `__token__`: - - ```bash - FLIT_USERNAME=__token__ flit publish - ``` - - and a PyPI token has to be passed in place of the password. - -### Credits - -The initial dynamic library introspection code was written by @anton-malakhov -for the smp package available at https://github.com/IntelPython/smp . - -threadpoolctl extends this for other operating systems. Contrary to smp, -threadpoolctl does not attempt to limit the size of Python multiprocessing -pools (threads or processes) or set operating system-level CPU affinity -constraints: threadpoolctl only interacts with native libraries via their -public runtime APIs. + +Setting the number of threads may seem like a simple operation, but in practice +it can do quite different things depending on the underlying third-party library +being limited. + +### Kind of thread pool + +* In some libraries, there is a shared process-wide thread pool. Setting the + number of threads changes the size of this shared pool. +* In other libraries, there is a thread pool per calling thread. So each Python + thread gets its own personal thread pool from the third-party library. + +To give a concrete example: if you have 10 cores, and you're using OpenBLAS with +the `pthreads` backend on Linux, by default there is a single 10-worker pool of +threads created by OpenBLAS. If you start 10 Python threads, each calling BLAS +routines, all those Python threads will feed into that single 10-thread pool. +In total, you will have 20 threads running (10 Python, 10 OpenBLAS). + +On the other hand, if you use MKL, each Python thread gets its own individual +pool of worker threads from MKL, so if each Python threads calls a BLAS routine +in MKL, you will get 10×10 = 100 MKL threads! + +### Scope of limiting + +If you have a third-party library that creates a thread pool per calling thread, +there is another question about what limiting the number of threads means: which +threads are affected by the limiting API? + +* **All threads in the process:** In this case, setting the limits changes it + for all threads in the process. So if you limit from the main Python thread, + all other Python threads will have the same limit. +* **Current thread only:** Only the current thread's thread pool size will be + limited. + +MKL for example has two APIs, covering both cases, +[`mkl_set_num_threads()`](https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2026-0/mkl-set-num-threads.html) +and [`mkl_set_num_threads_local()`](https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2026-0/mkl-set-num-threads-local.html) respectively. + +### `threadpoolctl`'s policy for thread limiting + +When there is a choice between different APIs, `threadpoolctl` will prefer APIs +that: + +1. Have a per-thread worker pool (so each Python thread gets its own pool from + the underlying third-party library). +2. Only affect the current thread. + +Current libraries where `threadpoolctl` explicitly makes this choice are: + +* MKL, for all threading backends. +* OpenBLAS (v0.3.34 or later) when using the OpenMP backend, on Linux and macOS. + threadpoolctl then limits OpenBLAS with `omp_set_num_threads`, so the limit + also applies to the OpenMP parallel regions of the calling code. To get + sequential BLAS calls within an OpenMP parallel region, see + [Sequential BLAS within OpenMP parallel region](#sequential-blas-within-openmp-parallel-region). + +When using OpenMP, a limit specific to the current is also the default on Linux +and macOS. On Windows OpenMP has a per-thread worker pool but the limiting API +When using OpenMP, a limit specific to the current thread is also the default on +Linux and macOS. On Windows OpenMP has a per-thread worker pool but the limiting API + +### Checking the behavior of your installed libraries + +When you run `python -m threadpoolctl` it will include the scope of the API +limit in the `"thread_limit_scope"` field, with the value of `"process"` +indicating the API affects the whole process and `"current_thread"` indicating a +per-thread thread pool. Information about the kind of limiting in the former +case (shared process-wide pool, or per-thread pool) is not included. + +Here is example output for OpenBLAS with pthreads: + +```shell-session +$ python -m threadpoolctl -i numpy +[ + { + "user_api": "blas", + "internal_api": "openblas", + "num_threads": 12, + "prefix": "libscipy_openblas", + "version": "0.3.33.112.0", + "threading_layer": "pthreads", + "architecture": "Haswell", + "thread_limit_scope": "process" + } +] +``` + +You can also use another script to empirically determine both the kind and scope of the API. From a checkout of the [`threadpoolctl` GitHub repo](https://github.com/joblib/threadpoolctl): + +```shell-session +$ python -m tests.empirical_scope_observation blas +== Observed behavior: blas == +10x12 Python threads each running a parallel workload with a 2 thread limit. +Maximum number of observed threads (includes Python and blas threads): 53 +Presumed API scope: process-wide shared thread pool +``` + +You can also call this script with an `openmp` argument instead of `blas`. + +## Maintainers + +To make a release: + +- Create a PR to bump the version number (`__version__`) in `threadpoolctl.py` and + update the release date in `CHANGES.md`. + +- Merge the PR and check that the `Publish threadpoolctl distribution to TestPyPI` job + of the `publish-to-pypi.yml` workflow successfully uploaded the wheel and source + distribution to Test PyPI. + +- If everything is fine create a tag for the release and push it to github: + + ```bash + git tag -a X.Y.Z + git push git@github.com:joblib/threadpoolctl.git X.Y.Z + ``` + +- Check that the `Publish threadpoolctl distribution to PyPI` job of the + `publish-to-pypi.yml` workflow successfully uploaded the wheel and source distribution + to PyPI this time. + +- Create a PR for the release on the [conda-forge feedstock](https://github.com/conda-forge/threadpoolctl-feedstock) (or wait for the bot to make it). + +- Publish the release on github. + +If for some reason the steps above can't be achieved and a munual upload of the wheel +and source distribution is needed: + +- Build the distribution archives: + + ```bash + pip install flit + flit build + ``` + +- Upload the wheels and source distribution to PyPI using flit. Since PyPI doesn't + allow password authentication anymore, the username needs to be changed to the + generic name `__token__`: + + ```bash + FLIT_USERNAME=__token__ flit publish + ``` + + and a PyPI token has to be passed in place of the password. + +### Credits + +The initial dynamic library introspection code was written by @anton-malakhov +for the smp package available at https://github.com/IntelPython/smp . + +threadpoolctl extends this for other operating systems. Contrary to smp, +threadpoolctl does not attempt to limit the size of Python multiprocessing +pools (threads or processes) or set operating system-level CPU affinity +constraints: threadpoolctl only interacts with native libraries via their +public runtime APIs. diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index 82c918f9..93c32292 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -717,6 +717,7 @@ def run_in_thread(f, *args, **kwargs): thread = Thread(target=lambda: f(*args, **kwargs)) thread.start() thread.join() + return thread def test_openblas_non_openmp_restricted_to_controlling_thread(): @@ -763,6 +764,7 @@ def check_openblas_non_openmp_restricted_to_controlling_thread(openblas_lib): # This thread can change num threads: openblas_lib.set_num_threads(1) assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._SHARED_STATE.controlling_thread is current_thread() # Another thread's attempts are no-op: result = [] @@ -774,10 +776,12 @@ def another_thread(): run_in_thread(another_thread) assert result[-1] == 1 assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._SHARED_STATE.controlling_thread is current_thread() # This thread can still make changes: openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == orig_num_threads + assert openblas_lib._SHARED_STATE.controlling_thread is current_thread() # Another thread clearing the controlling thread has no effect: def clear_and_set(): @@ -786,15 +790,18 @@ def clear_and_set(): run_in_thread(clear_and_set) assert openblas_lib.get_num_threads() == orig_num_threads + assert openblas_lib._SHARED_STATE.controlling_thread is current_thread() # The controlling thread can clear, and then another thread can set: openblas_lib.clear_controlling_thread() - run_in_thread(openblas_lib.set_num_threads, 1) + worker = run_in_thread(openblas_lib.set_num_threads, 1) assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._SHARED_STATE.controlling_thread is worker # Dead threads lose control: openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == orig_num_threads + assert openblas_lib._SHARED_STATE.controlling_thread is current_thread() # Release control: openblas_lib.clear_controlling_thread() @@ -813,42 +820,75 @@ def set_and_exit_after_lock_released(): t.start() was_set.acquire() assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._SHARED_STATE.controlling_thread is t # This thread can't set since it's not the controller: openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == 1 + assert openblas_lib._SHARED_STATE.controlling_thread is t # But when forced clear happens, we can set: openblas_lib.clear_controlling_thread(force=True) openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == orig_num_threads + assert openblas_lib._SHARED_STATE.controlling_thread is current_thread() # Clean up finish.release() t.join() -def _openblas_controlling_thread(): - openblas_controller = ThreadpoolController().select(internal_api="openblas") - - if not (openblas_controller): +def _non_openmp_openblas(): + openblas = ThreadpoolController().select(internal_api="openblas") + if not openblas: pytest.skip("requires OpenBLAS.") - if {lc.threading_layer for lc in openblas_controller.lib_controllers} == {"openmp"}: + libs = [] + seen = set() + for lib in openblas.lib_controllers: + if lib.threading_layer == "openmp" or lib.filepath in seen: + continue + seen.add(lib.filepath) + libs.append(lib) + if not libs: pytest.skip("require non-OpenMP OpenBLAS.") - return OpenBLASController._SHARED_STATE.controlling_thread + return libs + + +def _openblas_controlling_thread(): + libs = _non_openmp_openblas() + return next(lib._SHARED_STATE.controlling_thread for lib in libs) -def test_openblas_controlling_thread_cleared_on_limiter_context_manager_exit(): +def _restore_openblas_thread_counts(libs): + originals = [(lib, lib.get_num_threads()) for lib in libs] + + def restore(): + OpenBLASController.clear_controlling_thread(force=True) + for lib, orig in originals: + if orig is not None: + lib.set_num_threads(orig) + OpenBLASController.clear_controlling_thread(force=True) + + return restore + + +def test_openblas_controlling_thread_cleared_on_limiter_context_manager_exit(request): + libs, lib, orig_limit = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) with threadpool_limits(limits=1): assert _openblas_controlling_thread() is current_thread() with threadpool_limits(limits=1): assert _openblas_controlling_thread() is current_thread() assert _openblas_controlling_thread() is current_thread() + assert _openblas_controlling_thread() is None + assert lib.get_num_threads() == orig_limit -def test_openblas_controlling_thread_cleared_on_decorator_exit(): +def test_openblas_controlling_thread_cleared_on_decorator_exit_non_recursive(request): + libs, lib, orig_limit = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) @threadpool_limits.wrap(limits=1) def inner(): @@ -862,6 +902,145 @@ def outer(): outer() assert _openblas_controlling_thread() is None + assert lib.get_num_threads() == orig_limit + + +def _require_resizable_non_openmp_openblas(): + libs = _non_openmp_openblas() + lib = libs[0] + orig = lib.get_num_threads() + if orig is None or orig < 2: + pytest.skip("Test requires a non-OpenMP OpenBLAS with at least 2 threads") + return libs, lib, orig + + +@pytest.mark.parametrize( + "make_wrap", [lambda: ThreadpoolController().wrap, lambda: threadpool_limits.wrap] +) +def test_recursive_wrap_releases_the_shared_pool_claim(request, make_wrap): + # One wrapped function calling itself must keep the limit between the inner + # return and the outer return, then restore the original count and drop the + # controlling thread. Another thread can then set the count. + libs, lib, orig_limit = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + # To avoid annoying libiomp/omp warning happening at import time. + wrap = make_wrap() + + def set_in_thread(): + lib.set_num_threads(2) + + @wrap(limits=1, user_api="blas") + def recurse(depth): + assert lib.get_num_threads() == 1 + assert _openblas_controlling_thread() is current_thread() + + # This should be no-op since this thread controls OpenBLAS: + run_in_thread(set_in_thread) + + assert lib.get_num_threads() == 1 + if depth: + recurse(depth - 1) + assert lib.get_num_threads() == 1 + + recurse(1) + assert lib.get_num_threads() == orig_limit + + observed = {} + + def other_thread(): + lib.set_num_threads(1) + observed["threads"] = lib.get_num_threads() + observed["holder"] = lib._SHARED_STATE.controlling_thread + + assert _openblas_controlling_thread() is None + + worker = run_in_thread(other_thread) + assert observed["threads"] == 1 + assert observed["holder"] is worker + assert _openblas_controlling_thread() is worker + + +def test_recursive_wrap_releases_the_shared_pool_claim_exception(request): + libs, lib, orig_limit = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + @threadpool_limits.wrap(limits=1, user_api="blas") + def recurse(depth): + assert lib.get_num_threads() == 1 + if depth: + assert _openblas_controlling_thread() is current_thread() + recurse(depth - 1) + assert lib.get_num_threads() == 1 + else: + raise ZeroDivisionError() + + with pytest.raises(ZeroDivisionError): + recurse(1) + assert lib.get_num_threads() == orig_limit + assert _openblas_controlling_thread() is None + + +@pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="Relies on glibc being loaded" +) +def test_wrap_notices_new_library(): + """A newly available library is noticed by subsequent wrap() calls.""" + pytest.importorskip("numpy") + _require_resizable_non_openmp_openblas() + + script = """ +import numpy + +from threadpoolctl import ( + ThreadpoolController, + LibController, + reset_cached_controller, + register, + threadpool_limits, +) + +fake_sets = [] + +class FakeController(LibController): + user_api = "fake" + internal_api = "fake" + filename_prefixes = ("libc.",) + check_symbols = ("printf",) + + def get_num_threads(self): + if not fake_sets: + return 2 + else: + return fake_sets[-1] + + def set_num_threads(self, num_threads): + fake_sets.append(num_threads) + + def get_version(self): + return None + +blas = ThreadpoolController().select(user_api="blas").lib_controllers[0] +blas.set_num_threads(2) + +registered = False + +@threadpool_limits.wrap(limits=1) +def wrapped(): + assert blas.get_num_threads() == 1 + global registered + if not registered: + registered = True + register(FakeController) + +wrapped() +assert blas.get_num_threads() == 2 +assert not fake_sets +wrapped() +assert blas.get_num_threads() == 2 +assert fake_sets == [1, 2], fake_sets +""" + subprocess.check_call([sys.executable, "-c", script]) def test_openblas_controlling_thread_manual_increment_decrement_api(request): @@ -880,6 +1059,22 @@ def test_openblas_controlling_thread_manual_increment_decrement_api(request): assert _openblas_controlling_thread() is None +def test_restore_original_limits_keeps_openblas_controlling_thread(request): + libs, lib, orig = _require_resizable_non_openmp_openblas() + request.addfinalizer(_restore_openblas_thread_counts(libs)) + + limiter = threadpool_limits(limits=1) + assert lib.get_num_threads() == 1 + assert lib._SHARED_STATE.controlling_thread is current_thread() + + limiter.restore_original_limits() + assert lib.get_num_threads() == orig + assert lib._SHARED_STATE.controlling_thread is current_thread() + + OpenBLASController.clear_controlling_thread() + assert lib._SHARED_STATE.controlling_thread is None + + # skip test if not run in a azure pipelines job since it relies on a specific flexiblas # installation. @pytest.mark.skipif( diff --git a/threadpoolctl.py b/threadpoolctl.py index 756168cb..820e1e1c 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -558,7 +558,10 @@ def set_num_threads(self, num_threads): # threads fighting over limits: this_thread = current_thread() controlling_thread = self._SHARED_STATE.controlling_thread - if controlling_thread is None or not controlling_thread.is_alive(): + if controlling_thread is this_thread: + # This thread controls OpenBLAS: + return self._real_set_num_threads(num_threads) + elif controlling_thread is None or not controlling_thread.is_alive(): # No thread controls OpenBLAS, so try to take control: with self._SHARED_STATE.lock: # Another thread might have changed it, so ensure that's not @@ -571,11 +574,8 @@ def set_num_threads(self, num_threads): # Controlling thread should now be set, to either this thread or # another, so try again: return self.set_num_threads(num_threads) - elif controlling_thread is this_thread: - # This thread controls OpenBLAS: - return self._real_set_num_threads(num_threads) else: - # Some other thread controlls OpenBLAS: do nothing. + # Some other thread controls OpenBLAS: do nothing. return def _real_set_num_threads(self, num_threads): @@ -1032,6 +1032,63 @@ def threadpool_info(debugging_info: bool = False): _missing = object() +def _check_limiter_params(controller, limits, user_api): + """Suitable values for the _limits and _user_api attributes""" + + if isinstance(limits, str) and limits == "sequential_blas_under_openmp": + ( + limits, + user_api, + ) = controller._get_params_for_sequential_blas_under_openmp().values() + + if limits is None or isinstance(limits, int): + if user_api is None: + user_api = _ALL_USER_APIS + elif user_api in _ALL_USER_APIS: + user_api = [user_api] + else: + raise ValueError( + f"user_api must be either in {_ALL_USER_APIS} or None. Got " + f"{user_api} instead." + ) + + if limits is not None: + limits = {api: limits for api in user_api} + else: + if isinstance(limits, list): + # This should be a list of dicts of library info, for + # compatibility with the result from threadpool_info. + result = {} + for lib_info in limits: + # Use prefix key for backwards compatibility with older + # versions that didn't use the filepath as the key. + if (filepath := lib_info.get("filepath", _missing)) is not _missing: + key = filepath + else: + key = lib_info["prefix"] + result[key] = lib_info["num_threads"] + limits = result + elif isinstance(limits, ThreadpoolController): + # To set the limits from the library controllers of a + # ThreadpoolController object. + limits = { + lib_controller.filepath: lib_controller.num_threads + for lib_controller in limits.lib_controllers + } + + if not isinstance(limits, dict): + raise TypeError( + "limits must either be an int, a list, a dict, or " + f"'sequential_blas_under_openmp'. Got {type(limits)} instead" + ) + + # With a dictionary, can set both specific limit for given + # libraries and global limit for user_api. Fetch each separately. + user_api = [api for api in limits if api in _ALL_USER_APIS] + + return limits, user_api + + class _ThreadpoolLimiter: """The guts of ThreadpoolController.limit @@ -1045,7 +1102,9 @@ class _ThreadpoolLimiter: def __init__(self, controller, *, limits=None, user_api=None): self._controller = controller - self._limits, self._user_api = self._check_params(limits, user_api) + self._limits, self._user_api = _check_limiter_params( + controller, limits, user_api + ) self._original_info = self._controller.info() self._set_threadpool_limits() @@ -1110,62 +1169,6 @@ def get_original_num_threads(self): return num_threads - def _check_params(self, limits, user_api): - """Suitable values for the _limits and _user_api attributes""" - - if isinstance(limits, str) and limits == "sequential_blas_under_openmp": - ( - limits, - user_api, - ) = self._controller._get_params_for_sequential_blas_under_openmp().values() - - if limits is None or isinstance(limits, int): - if user_api is None: - user_api = _ALL_USER_APIS - elif user_api in _ALL_USER_APIS: - user_api = [user_api] - else: - raise ValueError( - f"user_api must be either in {_ALL_USER_APIS} or None. Got " - f"{user_api} instead." - ) - - if limits is not None: - limits = {api: limits for api in user_api} - else: - if isinstance(limits, list): - # This should be a list of dicts of library info, for - # compatibility with the result from threadpool_info. - result = {} - for lib_info in limits: - # Use prefix key for backwards compatibility with older - # versions that didn't use the filepath as the key. - if (filepath := lib_info.get("filepath", _missing)) is not _missing: - key = filepath - else: - key = lib_info["prefix"] - result[key] = lib_info["num_threads"] - limits = result - elif isinstance(limits, ThreadpoolController): - # To set the limits from the library controllers of a - # ThreadpoolController object. - limits = { - lib_controller.filepath: lib_controller.num_threads - for lib_controller in limits.lib_controllers - } - - if not isinstance(limits, dict): - raise TypeError( - "limits must either be an int, a list, a dict, or " - f"'sequential_blas_under_openmp'. Got {type(limits)} instead" - ) - - # With a dictionary, can set both specific limit for given - # libraries and global limit for user_api. Fetch each separately. - user_api = [api for api in limits if api in _ALL_USER_APIS] - - return limits, user_api - def _set_threadpool_limits(self): """Change the maximal number of threads in selected thread pools.""" if self._limits is None: @@ -1188,22 +1191,42 @@ def _set_threadpool_limits(self): lib_controller.set_num_threads(num_threads) -class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter, ContextDecorator): - """Same as _ThreadpoolLimiter but to be used as a decorator""" +class _ThreadpoolLimiterDecorator(ContextDecorator): + """Same as _ThreadpoolLimiter but to be used as a decorator. + + If ``controller`` is None, use ``get_cached_controller()``. + """ def __init__(self, controller, *, limits=None, user_api=None): - self._limits, self._user_api = self._check_params(limits, user_api) self._controller = controller + _check_limiter_params(self._get_controller(), limits, user_api) + self._limits, self._user_api = limits, user_api + self._limiters = thread_local() + + def _get_controller(self): + if self._controller is not None: + return self._controller + return get_cached_controller() def __enter__(self): - # we need to set the limits here and not in the __init__ because we want the - # limits to be set when calling the decorated function, not when creating the - # decorator. - self._original_info = self._controller.info() - self._set_threadpool_limits() - CONTROLLING_THREAD_OWNERSHIP.increment() + # we need to set the limits here and not in the __init__ because we + # want the limits to be set when calling the decorated function, not + # when creating the decorator. + if not hasattr(self._limiters, "stack"): + self._limiters.stack = [] + controller = self._get_controller() + limiter = _ThreadpoolLimiter( + controller, limits=self._limits, user_api=self._user_api + ) + limiter.__enter__() + self._limiters.stack.append(limiter) return self + def __exit__(self, type, value, traceback): + """Set the limits back to their original values""" + limiter = self._limiters.stack.pop() + return limiter.__exit__(type, value, traceback) + @_format_docstring( USER_APIS=", ".join(f'"{api}"' for api in _ALL_USER_APIS), @@ -1271,7 +1294,8 @@ def __init__(self, limits=None, user_api=None): @classmethod def wrap(cls, limits=None, user_api=None): - return super().wrap(ThreadpoolController(), limits=limits, user_api=user_api) + """Wrap a function, using ``get_cached_controller()`` on every call.""" + return super().wrap(None, limits=limits, user_api=user_api) _INCOMPATIBLE_OPENMP_MESSAGE = """