1715 lines
66 KiB
Python
1715 lines
66 KiB
Python
"""threadpoolctl
|
|
|
|
This module provides utilities to introspect native libraries that relies on
|
|
thread pools (notably BLAS and OpenMP implementations) and dynamically set the
|
|
maximal number of threads they can use.
|
|
"""
|
|
|
|
# License: BSD 3-Clause
|
|
|
|
# The code to introspect dynamically loaded libraries on POSIX systems is
|
|
# 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
|
|
import os
|
|
import re
|
|
import sys
|
|
import ctypes
|
|
import itertools
|
|
import textwrap
|
|
from threading import Thread
|
|
from typing import Callable, Literal, final
|
|
import warnings
|
|
from abc import ABC, abstractmethod
|
|
from functools import lru_cache
|
|
from contextlib import ContextDecorator
|
|
|
|
__version__ = "3.7.0"
|
|
__all__ = [
|
|
"threadpool_limits",
|
|
"threadpool_info",
|
|
"ThreadpoolController",
|
|
"LibController",
|
|
"register",
|
|
]
|
|
|
|
|
|
# One can get runtime errors or even segfaults due to multiple OpenMP libraries
|
|
# loaded simultaneously which can happen easily in Python when importing and
|
|
# using compiled extensions built with different compilers and therefore
|
|
# different OpenMP runtimes in the same program. In particular libiomp (used by
|
|
# Intel ICC) and libomp used by clang/llvm tend to crash. This can happen for
|
|
# instance when calling BLAS inside a prange. Setting the following environment
|
|
# variable allows multiple OpenMP libraries to be loaded. It should not degrade
|
|
# performances since we manually take care of potential over-subscription
|
|
# performance issues, in sections of the code where nested OpenMP loops can
|
|
# happen, by dynamically reconfiguring the inner OpenMP runtime to temporarily
|
|
# disable it while under the scope of the outer OpenMP parallel section.
|
|
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "True")
|
|
|
|
# Hard limit for Windows library paths resolved through GetModuleFileNameExW.
|
|
# Kept below unlimited for security reasons; see CHANGES.md and PR #189.
|
|
_WINDOWS_MAX_LIBRARY_PATH_LENGTH = 2600
|
|
|
|
# Structure to cast the info on dynamically loaded library. See
|
|
# https://linux.die.net/man/3/dl_iterate_phdr for more details.
|
|
_SYSTEM_UINT = ctypes.c_uint64 if sys.maxsize > 2**32 else ctypes.c_uint32
|
|
_SYSTEM_UINT_HALF = ctypes.c_uint32 if sys.maxsize > 2**32 else ctypes.c_uint16
|
|
|
|
|
|
class _dl_phdr_info(ctypes.Structure):
|
|
_fields_ = [
|
|
("dlpi_addr", _SYSTEM_UINT), # Base address of object
|
|
("dlpi_name", ctypes.c_char_p), # path to the library
|
|
("dlpi_phdr", ctypes.c_void_p), # pointer on dlpi_headers
|
|
("dlpi_phnum", _SYSTEM_UINT_HALF), # number of elements in dlpi_phdr
|
|
]
|
|
|
|
|
|
# The RTLD_NOLOAD flag for loading shared libraries is not defined on Windows.
|
|
try:
|
|
_RTLD_NOLOAD = os.RTLD_NOLOAD
|
|
except AttributeError:
|
|
_RTLD_NOLOAD = ctypes.DEFAULT_MODE
|
|
|
|
|
|
# What scope does the API affect?
|
|
_ThreadLimitScope = Literal[
|
|
# Using the API sets a limit only on the current thread.
|
|
"current_thread",
|
|
# Using the API sets a limit for every thread in the process; whether or
|
|
# not it's a shared process-wide pool or per-thread limit needs to be
|
|
# determined some other way.
|
|
"process",
|
|
# Something else, unexpected; perhaps another variant, perhaps information
|
|
# can't be determined under the current configuration.
|
|
"unknown",
|
|
]
|
|
|
|
|
|
def _determine_thread_limit_scope(
|
|
get_n_threads: Callable[[], int], set_n_threads: Callable[[int], None]
|
|
) -> _ThreadLimitScope:
|
|
"""
|
|
Run some experiments to determine the scope of the given get/set API.
|
|
|
|
This function might not work if you only have one core available.
|
|
|
|
This function might not work if you set a limit on a library with an
|
|
environment variable.
|
|
|
|
The function works by changing the number of threads in loaded controllers,
|
|
which can be a process-wide change. As such, it is not always thread-safe.
|
|
An attempt will be made to restore all settings to their previous state,
|
|
but the result may be subtly different, e.g. if "unset" has different
|
|
semantics than "set to the default returned value".
|
|
"""
|
|
previous = get_n_threads()
|
|
|
|
# Some plausible constraints we need to keep in mind:
|
|
#
|
|
# 1. The API might not allow setting more than the number of (available, or
|
|
# physical) cores.
|
|
# 2. Some hard limit on number of threads.
|
|
try:
|
|
# Choose a desired number of threads that is different than the current
|
|
# number, and hopefully achievable under the current configuration:
|
|
if previous < 2:
|
|
expected = 2
|
|
else:
|
|
# It's 2 or more, so shrink it slightly:
|
|
expected = previous - 1
|
|
|
|
thread_result = []
|
|
|
|
def get_and_set() -> None:
|
|
set_n_threads(expected)
|
|
thread_result.append(get_n_threads())
|
|
|
|
thread = Thread(target=get_and_set)
|
|
thread.start()
|
|
thread.join()
|
|
|
|
# First, getting in the same thread as a set should always give same
|
|
# number, if it's a number in a reasonable range. A possible exception
|
|
# fo failing this is if the number of thread is limited by available
|
|
# CPU, and only one CPU is available. In that case we can't empirically
|
|
# determine how the API works. We try to not reach that point here, but
|
|
# you can imagine a thread pool implementation that is aware of
|
|
# cgroups, in which case a Docker container limited to one core will
|
|
# pass the safety check at the start of the function. Perhaps
|
|
# cpu_count() from loky should be moved into this package...
|
|
if thread_result != [expected]:
|
|
return "unknown"
|
|
|
|
# Now, check this thread:
|
|
if get_n_threads() == expected:
|
|
# Setting modified this thread's results too:
|
|
return "process"
|
|
elif get_n_threads() == previous:
|
|
# Setting modified the other thread, but not this one:
|
|
return "current_thread"
|
|
else:
|
|
# No idea what's going on:
|
|
return "unknown"
|
|
finally:
|
|
set_n_threads(previous)
|
|
|
|
|
|
class LibController(ABC):
|
|
"""Abstract base class for the individual library controllers
|
|
|
|
A library controller must expose the following class attributes:
|
|
- user_api : str
|
|
Usually the name of the library or generic specification the library
|
|
implements, e.g. "blas" is a specification with different implementations.
|
|
- internal_api : str
|
|
Usually the name of the library or concrete implementation of some
|
|
specification, e.g. "openblas" is an implementation of the "blas"
|
|
specification.
|
|
- filename_prefixes : tuple
|
|
Possible prefixes of the shared library's filename that allow to
|
|
identify the library. e.g. "libopenblas" for libopenblas.so.
|
|
|
|
and implement the following methods: `get_num_threads`, `set_num_threads` and
|
|
`get_version`.
|
|
|
|
Threadpoolctl loops through all the loaded shared libraries and tries to match
|
|
the filename of each library with the `filename_prefixes`. If a match is found, a
|
|
controller is instantiated and a handler to the library is stored in the `dynlib`
|
|
attribute as a `ctypes.CDLL` object. It can be used to access the necessary symbols
|
|
of the shared library to implement the above methods.
|
|
|
|
The following information will be exposed in the info dictionary:
|
|
- user_api : standardized API, if any, or a copy of internal_api.
|
|
- internal_api : implementation-specific API.
|
|
- num_threads : the current thread limit.
|
|
- prefix : prefix of the shared library's filename.
|
|
- filepath : path to the loaded shared library.
|
|
- version : version of the library (if available).
|
|
|
|
In addition, each library controller may expose internal API specific entries. They
|
|
must be set as attributes in the `set_additional_attributes` method.
|
|
"""
|
|
|
|
@final
|
|
def __init__(self, *, filepath=None, prefix=None, parent=None):
|
|
"""This is not meant to be overriden by subclasses."""
|
|
self.parent = parent
|
|
self.prefix = prefix
|
|
self.filepath = filepath
|
|
self.dynlib = ctypes.CDLL(filepath, mode=_RTLD_NOLOAD)
|
|
self._symbol_prefix, self._symbol_suffix = self._find_affixes()
|
|
self.version = self.get_version()
|
|
self.set_additional_attributes()
|
|
|
|
def info(self, debugging_info: bool = False):
|
|
"""Return relevant info wrapped in a dict.
|
|
|
|
Parameters
|
|
----------
|
|
debugging_info : bool
|
|
|
|
Include extra fields which require more intrusive actions to
|
|
obtain, and which can't always reliably be determined.
|
|
"""
|
|
hidden_attrs = ("dynlib", "parent", "_symbol_prefix", "_symbol_suffix")
|
|
result = {
|
|
"user_api": self.user_api,
|
|
"internal_api": self.internal_api,
|
|
"num_threads": self.num_threads,
|
|
**{k: v for k, v in vars(self).items() if k not in hidden_attrs},
|
|
}
|
|
if debugging_info:
|
|
result["thread_limit_scope"] = _determine_thread_limit_scope(
|
|
self.get_num_threads, self.set_num_threads
|
|
)
|
|
return result
|
|
|
|
def set_additional_attributes(self):
|
|
"""Set additional attributes meant to be exposed in the info dict"""
|
|
|
|
@property
|
|
def num_threads(self):
|
|
"""Exposes the current thread limit as a dynamic property
|
|
|
|
This is not meant to be used or overriden by subclasses.
|
|
"""
|
|
return self.get_num_threads()
|
|
|
|
@abstractmethod
|
|
def get_num_threads(self):
|
|
"""Return the maximum number of threads available to use"""
|
|
|
|
@abstractmethod
|
|
def set_num_threads(self, num_threads):
|
|
"""Set the maximum number of threads to use
|
|
|
|
When possible, implementations of this method should choose a thread
|
|
limiting API that only applies to the current thread.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_version(self):
|
|
"""Return the version of the shared library"""
|
|
|
|
def _find_affixes(self):
|
|
"""Return the affixes for the symbols of the shared library"""
|
|
return "", ""
|
|
|
|
def _get_symbol(self, name):
|
|
"""Return the symbol of the shared library accounding for the affixes"""
|
|
return getattr(
|
|
self.dynlib, f"{self._symbol_prefix}{name}{self._symbol_suffix}", None
|
|
)
|
|
|
|
|
|
class OpenBLASController(LibController):
|
|
"""Controller class for OpenBLAS"""
|
|
|
|
user_api = "blas"
|
|
internal_api = "openblas"
|
|
filename_prefixes = (
|
|
"libopenblas",
|
|
"libblas", # legacy conda-forge Windows shim, see _make_controller_from_path
|
|
"libscipy_openblas",
|
|
"openblas", # Windows conda package use openblas.dll
|
|
)
|
|
|
|
_symbol_prefixes = ("", "scipy_")
|
|
_symbol_suffixes = ("", "64_", "_64")
|
|
|
|
# All variations of "openblas_get_num_threads", accounting for the affixes
|
|
check_symbols = tuple(
|
|
f"{prefix}openblas_get_num_threads{suffix}"
|
|
for prefix, suffix in itertools.product(_symbol_prefixes, _symbol_suffixes)
|
|
)
|
|
|
|
def _find_affixes(self):
|
|
for prefix, suffix in itertools.product(
|
|
self._symbol_prefixes, self._symbol_suffixes
|
|
):
|
|
if hasattr(self.dynlib, f"{prefix}openblas_get_num_threads{suffix}"):
|
|
return prefix, suffix
|
|
|
|
def set_additional_attributes(self):
|
|
self.threading_layer = self._get_threading_layer()
|
|
self.architecture = self._get_architecture()
|
|
|
|
def get_num_threads(self):
|
|
# See discussion in set_num_threads for details:
|
|
if self.threading_layer == "openmp":
|
|
symbol = "omp_get_max_threads"
|
|
else:
|
|
symbol = "openblas_get_num_threads"
|
|
get_num_threads_func = self._get_symbol(symbol)
|
|
if get_num_threads_func is not None:
|
|
return get_num_threads_func()
|
|
return None
|
|
|
|
def set_num_threads(self, num_threads):
|
|
# The OpenBLAS limiting API is process-wide, and we want current thread
|
|
# limit if possible. When OpenBLAS is backed by OpenMP, using the
|
|
# OpenMP API allows for current thread limiting when OpenMP has that
|
|
# behavior. That is the case for libgomp, libomp, and libiomp, what you
|
|
# would find on Linux or macOS.
|
|
#
|
|
# On Windows the Visual C++ OpenMP API is process-wide, unfortunately,
|
|
# though this may be fixed if the /openmp:llvm flag is used:
|
|
# https://github.com/joblib/threadpoolctl/issues/230
|
|
#
|
|
# Also worth knowing that before v0.3.34, the OpenBLAS limiting API is
|
|
# broken when using OpenMP threading:
|
|
# https://github.com/OpenMathLib/OpenBLAS/issues/5806
|
|
if self.threading_layer == "openmp":
|
|
symbol = "omp_set_num_threads"
|
|
else:
|
|
symbol = "openblas_set_num_threads"
|
|
set_num_threads_func = self._get_symbol(symbol)
|
|
if set_num_threads_func is not None:
|
|
return set_num_threads_func(num_threads)
|
|
return None
|
|
|
|
def get_version(self):
|
|
# None means OpenBLAS is not loaded or version < 0.3.4, since OpenBLAS
|
|
# did not expose its version before that.
|
|
get_version_func = self._get_symbol("openblas_get_config")
|
|
if get_version_func is not None:
|
|
get_version_func.restype = ctypes.c_char_p
|
|
config = get_version_func().split()
|
|
if config[0] == b"OpenBLAS":
|
|
return config[1].decode("utf-8")
|
|
return None
|
|
return None
|
|
|
|
def _get_threading_layer(self):
|
|
"""Return the threading layer of OpenBLAS"""
|
|
get_threading_layer_func = self._get_symbol("openblas_get_parallel")
|
|
if get_threading_layer_func is not None:
|
|
threading_layer = get_threading_layer_func()
|
|
if threading_layer == 2:
|
|
return "openmp"
|
|
elif threading_layer == 1:
|
|
return "pthreads"
|
|
return "disabled"
|
|
return "unknown"
|
|
|
|
def _get_architecture(self):
|
|
"""Return the architecture detected by OpenBLAS"""
|
|
get_architecture_func = self._get_symbol("openblas_get_corename")
|
|
if get_architecture_func is not None:
|
|
get_architecture_func.restype = ctypes.c_char_p
|
|
return get_architecture_func().decode("utf-8")
|
|
return None
|
|
|
|
|
|
class BLISController(LibController):
|
|
"""Controller class for BLIS"""
|
|
|
|
user_api = "blas"
|
|
internal_api = "blis"
|
|
filename_prefixes = (
|
|
"libblis",
|
|
"libblas",
|
|
) # libblas: legacy conda-forge Windows shim
|
|
check_symbols = (
|
|
"bli_thread_get_num_threads",
|
|
"bli_thread_set_num_threads",
|
|
"bli_info_get_version_str",
|
|
"bli_info_get_enable_openmp",
|
|
"bli_info_get_enable_pthreads",
|
|
"bli_arch_query_id",
|
|
"bli_arch_string",
|
|
)
|
|
|
|
def set_additional_attributes(self):
|
|
self.threading_layer = self._get_threading_layer()
|
|
self.architecture = self._get_architecture()
|
|
|
|
def get_num_threads(self):
|
|
get_func = getattr(self.dynlib, "bli_thread_get_num_threads", lambda: None)
|
|
num_threads = get_func()
|
|
# by default BLIS is single-threaded and get_num_threads
|
|
# returns -1. We map it to 1 for consistency with other libraries.
|
|
return 1 if num_threads == -1 else num_threads
|
|
|
|
def set_num_threads(self, num_threads):
|
|
set_func = getattr(
|
|
self.dynlib, "bli_thread_set_num_threads", lambda num_threads: None
|
|
)
|
|
return set_func(num_threads)
|
|
|
|
def get_version(self):
|
|
get_version_ = getattr(self.dynlib, "bli_info_get_version_str", None)
|
|
if get_version_ is None:
|
|
return None
|
|
|
|
get_version_.restype = ctypes.c_char_p
|
|
return get_version_().decode("utf-8")
|
|
|
|
def _get_threading_layer(self):
|
|
"""Return the threading layer of BLIS"""
|
|
if getattr(self.dynlib, "bli_info_get_enable_openmp", lambda: False)():
|
|
return "openmp"
|
|
elif getattr(self.dynlib, "bli_info_get_enable_pthreads", lambda: False)():
|
|
return "pthreads"
|
|
return "disabled"
|
|
|
|
def _get_architecture(self):
|
|
"""Return the architecture detected by BLIS"""
|
|
bli_arch_query_id = getattr(self.dynlib, "bli_arch_query_id", None)
|
|
bli_arch_string = getattr(self.dynlib, "bli_arch_string", None)
|
|
if bli_arch_query_id is None or bli_arch_string is None:
|
|
return None
|
|
|
|
# the true restype should be BLIS' arch_t (enum) but int should work
|
|
# for us:
|
|
bli_arch_query_id.restype = ctypes.c_int
|
|
bli_arch_string.restype = ctypes.c_char_p
|
|
return bli_arch_string(bli_arch_query_id()).decode("utf-8")
|
|
|
|
|
|
class FlexiBLASController(LibController):
|
|
"""Controller class for FlexiBLAS"""
|
|
|
|
user_api = "blas"
|
|
internal_api = "flexiblas"
|
|
filename_prefixes = ("libflexiblas",)
|
|
check_symbols = (
|
|
"flexiblas_get_num_threads",
|
|
"flexiblas_set_num_threads",
|
|
"flexiblas_get_version",
|
|
"flexiblas_list",
|
|
"flexiblas_list_loaded",
|
|
"flexiblas_current_backend",
|
|
)
|
|
|
|
@property
|
|
def loaded_backends(self):
|
|
return self._get_backend_list(loaded=True)
|
|
|
|
@property
|
|
def current_backend(self):
|
|
return self._get_current_backend()
|
|
|
|
def info(self, debugging_info: bool = False):
|
|
"""Return relevant info wrapped in a dict"""
|
|
# We override the info method because the loaded and current backends
|
|
# are dynamic properties
|
|
exposed_attrs = super().info(debugging_info=debugging_info)
|
|
exposed_attrs["loaded_backends"] = self.loaded_backends
|
|
exposed_attrs["current_backend"] = self.current_backend
|
|
|
|
return exposed_attrs
|
|
|
|
def set_additional_attributes(self):
|
|
self.available_backends = self._get_backend_list(loaded=False)
|
|
|
|
def get_num_threads(self):
|
|
get_func = getattr(self.dynlib, "flexiblas_get_num_threads", lambda: None)
|
|
num_threads = get_func()
|
|
# by default BLIS is single-threaded and get_num_threads
|
|
# returns -1. We map it to 1 for consistency with other libraries.
|
|
return 1 if num_threads == -1 else num_threads
|
|
|
|
def set_num_threads(self, num_threads):
|
|
set_func = getattr(
|
|
self.dynlib, "flexiblas_set_num_threads", lambda num_threads: None
|
|
)
|
|
return set_func(num_threads)
|
|
|
|
def get_version(self):
|
|
get_version_ = getattr(self.dynlib, "flexiblas_get_version", None)
|
|
if get_version_ is None:
|
|
return None
|
|
|
|
major = ctypes.c_int()
|
|
minor = ctypes.c_int()
|
|
patch = ctypes.c_int()
|
|
get_version_(ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch))
|
|
return f"{major.value}.{minor.value}.{patch.value}"
|
|
|
|
def _get_backend_list(self, loaded=False):
|
|
"""Return the list of available backends for FlexiBLAS.
|
|
|
|
If loaded is False, return the list of available backends from the FlexiBLAS
|
|
configuration. If loaded is True, return the list of actually loaded backends.
|
|
"""
|
|
func_name = f"flexiblas_list{'_loaded' if loaded else ''}"
|
|
get_backend_list_ = getattr(self.dynlib, func_name, None)
|
|
if get_backend_list_ is None:
|
|
return None
|
|
|
|
n_backends = get_backend_list_(None, 0, 0)
|
|
|
|
backends = []
|
|
for i in range(n_backends):
|
|
backend_name = ctypes.create_string_buffer(1024)
|
|
get_backend_list_(backend_name, 1024, i)
|
|
if backend_name.value.decode("utf-8") != "__FALLBACK__":
|
|
# We don't know when to expect __FALLBACK__ but it is not a real
|
|
# backend and does not show up when running flexiblas list.
|
|
backends.append(backend_name.value.decode("utf-8"))
|
|
return backends
|
|
|
|
def _get_current_backend(self):
|
|
"""Return the backend of FlexiBLAS"""
|
|
get_backend_ = getattr(self.dynlib, "flexiblas_current_backend", None)
|
|
if get_backend_ is None:
|
|
return None
|
|
|
|
backend = ctypes.create_string_buffer(1024)
|
|
get_backend_(backend, ctypes.sizeof(backend))
|
|
return backend.value.decode("utf-8")
|
|
|
|
def switch_backend(self, backend):
|
|
"""Switch the backend of FlexiBLAS
|
|
|
|
Parameters
|
|
----------
|
|
backend : str
|
|
The name or the path to the shared library of the backend to switch to. If
|
|
the backend is not already loaded, it will be loaded first.
|
|
"""
|
|
if backend not in self.loaded_backends:
|
|
if backend in self.available_backends:
|
|
load_func = getattr(self.dynlib, "flexiblas_load_backend", lambda _: -1)
|
|
else: # assume backend is a path to a shared library
|
|
load_func = getattr(
|
|
self.dynlib, "flexiblas_load_backend_library", lambda _: -1
|
|
)
|
|
res = load_func(str(backend).encode("utf-8"))
|
|
if res == -1:
|
|
raise RuntimeError(
|
|
f"Failed to load backend {backend!r}. It must either be the name of"
|
|
" a backend available in the FlexiBLAS configuration "
|
|
f"{self.available_backends} or the path to a valid shared library."
|
|
)
|
|
|
|
# Trigger a new search of loaded shared libraries since loading a new
|
|
# backend caused a dlopen.
|
|
self.parent._load_libraries()
|
|
|
|
switch_func = getattr(self.dynlib, "flexiblas_switch", lambda _: -1)
|
|
idx = self.loaded_backends.index(backend)
|
|
res = switch_func(idx)
|
|
if res == -1:
|
|
raise RuntimeError(f"Failed to switch to backend {backend!r}.")
|
|
|
|
|
|
class MKLController(LibController):
|
|
"""Controller class for MKL"""
|
|
|
|
user_api = "blas"
|
|
internal_api = "mkl"
|
|
filename_prefixes = (
|
|
"libmkl_rt",
|
|
"mkl_rt",
|
|
"libblas", # legacy conda-forge Windows shim, see _make_controller_from_path
|
|
)
|
|
check_symbols = (
|
|
"MKL_Get_Max_Threads",
|
|
"MKL_Set_Num_Threads_Local",
|
|
"MKL_Get_Version_String",
|
|
"MKL_Set_Threading_Layer",
|
|
)
|
|
|
|
def set_additional_attributes(self):
|
|
self.threading_layer = self._get_threading_layer()
|
|
|
|
def get_num_threads(self):
|
|
get_func = getattr(self.dynlib, "MKL_Get_Max_Threads", lambda: None)
|
|
return get_func()
|
|
|
|
def set_num_threads(self, num_threads):
|
|
set_func = getattr(
|
|
self.dynlib, "MKL_Set_Num_Threads_Local", lambda num_threads: None
|
|
)
|
|
return set_func(num_threads)
|
|
|
|
def get_version(self):
|
|
if not hasattr(self.dynlib, "MKL_Get_Version_String"):
|
|
return None
|
|
|
|
res = ctypes.create_string_buffer(200)
|
|
self.dynlib.MKL_Get_Version_String(res, 200)
|
|
|
|
version = res.value.decode("utf-8")
|
|
group = re.search(r"Version ([^ ]+) ", version)
|
|
if group is not None:
|
|
version = group.groups()[0]
|
|
return version.strip()
|
|
|
|
def _get_threading_layer(self):
|
|
"""Return the threading layer of MKL"""
|
|
# The function mkl_set_threading_layer returns the current threading
|
|
# layer. Calling it with an invalid threading layer allows us to safely
|
|
# get the threading layer
|
|
set_threading_layer = getattr(
|
|
self.dynlib, "MKL_Set_Threading_Layer", lambda layer: -1
|
|
)
|
|
layer_map = {
|
|
0: "intel",
|
|
1: "sequential",
|
|
2: "pgi",
|
|
3: "gnu",
|
|
4: "tbb",
|
|
-1: "not specified",
|
|
}
|
|
return layer_map[set_threading_layer(-1)]
|
|
|
|
|
|
class OpenMPController(LibController):
|
|
"""Controller class for OpenMP"""
|
|
|
|
user_api = "openmp"
|
|
internal_api = "openmp"
|
|
filename_prefixes = ("libiomp", "libgomp", "libomp", "vcomp")
|
|
check_symbols = (
|
|
"omp_get_max_threads",
|
|
"omp_get_num_threads",
|
|
)
|
|
|
|
def get_num_threads(self):
|
|
get_func = getattr(self.dynlib, "omp_get_max_threads", lambda: None)
|
|
return get_func()
|
|
|
|
def set_num_threads(self, num_threads):
|
|
set_func = getattr(self.dynlib, "omp_set_num_threads", lambda num_threads: None)
|
|
return set_func(num_threads)
|
|
|
|
def get_version(self):
|
|
# There is no way to get the version number programmatically in OpenMP.
|
|
return None
|
|
|
|
|
|
# Controllers for the libraries that we'll look for in the loaded libraries.
|
|
# Third party libraries can register their own controllers.
|
|
_ALL_CONTROLLERS = [
|
|
OpenBLASController,
|
|
BLISController,
|
|
MKLController,
|
|
OpenMPController,
|
|
FlexiBLASController,
|
|
]
|
|
|
|
# Helpers for the doc and test names
|
|
_ALL_USER_APIS = list(set(lib.user_api for lib in _ALL_CONTROLLERS))
|
|
_ALL_INTERNAL_APIS = [lib.internal_api for lib in _ALL_CONTROLLERS]
|
|
_ALL_PREFIXES = list(
|
|
set(prefix for lib in _ALL_CONTROLLERS for prefix in lib.filename_prefixes)
|
|
)
|
|
_ALL_BLAS_LIBRARIES = [
|
|
lib.internal_api for lib in _ALL_CONTROLLERS if lib.user_api == "blas"
|
|
]
|
|
_ALL_OPENMP_LIBRARIES = OpenMPController.filename_prefixes
|
|
|
|
|
|
def register(controller):
|
|
"""Register a new controller"""
|
|
_ALL_CONTROLLERS.append(controller)
|
|
_ALL_USER_APIS.append(controller.user_api)
|
|
_ALL_INTERNAL_APIS.append(controller.internal_api)
|
|
_ALL_PREFIXES.extend(controller.filename_prefixes)
|
|
|
|
|
|
def _format_docstring(*args, **kwargs):
|
|
def decorator(o):
|
|
if o.__doc__ is not None:
|
|
o.__doc__ = o.__doc__.format(*args, **kwargs)
|
|
return o
|
|
|
|
return decorator
|
|
|
|
|
|
@lru_cache(maxsize=10000)
|
|
def _realpath(filepath):
|
|
"""Small caching wrapper around os.path.realpath to limit system calls"""
|
|
return os.path.realpath(filepath)
|
|
|
|
|
|
@_format_docstring(USER_APIS=list(_ALL_USER_APIS), INTERNAL_APIS=_ALL_INTERNAL_APIS)
|
|
def threadpool_info(debugging_info: bool = False):
|
|
"""Return the maximal number of threads for each detected library.
|
|
|
|
Return a list with all the supported libraries that have been found. Each
|
|
library is represented by a dict with the following information:
|
|
|
|
- "user_api" : user API. Possible values are {USER_APIS}.
|
|
- "internal_api": internal API. Possible values are {INTERNAL_APIS}.
|
|
- "prefix" : filename prefix of the specific implementation.
|
|
- "filepath": path to the loaded library.
|
|
- "version": version of the library (if available).
|
|
- "num_threads": the current thread limit.
|
|
|
|
In addition, each library may contain internal_api specific entries.
|
|
|
|
Parameters
|
|
----------
|
|
debugging_info : bool
|
|
Include extra fields which require more intrusive actions to obtain,
|
|
and which can't always reliably be determined.
|
|
|
|
- "thread_limit_scope": When setting the number of threads, what is
|
|
affected. Possible values are "process", "current_thread", "unknown".
|
|
"""
|
|
return ThreadpoolController().info(debugging_info)
|
|
|
|
|
|
class _ThreadpoolLimiter:
|
|
"""The guts of ThreadpoolController.limit
|
|
|
|
Refer to the docstring of ThreadpoolController.limit for more details.
|
|
|
|
It will only act on the library controllers held by the provided `controller`.
|
|
Using the default constructor sets the limits right away such that it can be used as
|
|
a callable. Setting the limits can be delayed by using the `wrap` class method such
|
|
that it can be used as a decorator.
|
|
"""
|
|
|
|
def __init__(self, controller, *, limits=None, user_api=None):
|
|
self._controller = controller
|
|
self._limits, self._user_api, self._prefixes = self._check_params(
|
|
limits, user_api
|
|
)
|
|
self._original_info = self._controller.info()
|
|
self._set_threadpool_limits()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
self.restore_original_limits()
|
|
|
|
@classmethod
|
|
def wrap(cls, controller, *, limits=None, user_api=None):
|
|
"""Return an instance of this class that can be used as a decorator"""
|
|
return _ThreadpoolLimiterDecorator(
|
|
controller=controller, limits=limits, user_api=user_api
|
|
)
|
|
|
|
def restore_original_limits(self):
|
|
"""Set the limits back to their original values"""
|
|
for lib_controller, original_info in zip(
|
|
self._controller.lib_controllers, self._original_info
|
|
):
|
|
lib_controller.set_num_threads(original_info["num_threads"])
|
|
|
|
# Alias of `restore_original_limits` for backward compatibility
|
|
unregister = restore_original_limits
|
|
|
|
def get_original_num_threads(self):
|
|
"""Original num_threads from before calling threadpool_limits
|
|
|
|
Return a dict `{user_api: num_threads}`.
|
|
"""
|
|
num_threads = {}
|
|
warning_apis = []
|
|
|
|
for user_api in self._user_api:
|
|
limits = [
|
|
lib_info["num_threads"]
|
|
for lib_info in self._original_info
|
|
if lib_info["user_api"] == user_api
|
|
]
|
|
limits = set(limits)
|
|
n_limits = len(limits)
|
|
|
|
if n_limits == 1:
|
|
limit = limits.pop()
|
|
elif n_limits == 0:
|
|
limit = None
|
|
else:
|
|
limit = min(limits)
|
|
warning_apis.append(user_api)
|
|
|
|
num_threads[user_api] = limit
|
|
|
|
if warning_apis:
|
|
warnings.warn(
|
|
"Multiple values possible for following user apis: "
|
|
+ ", ".join(warning_apis)
|
|
+ ". Returning the minimum."
|
|
)
|
|
|
|
return num_threads
|
|
|
|
def _check_params(self, limits, user_api):
|
|
"""Suitable values for the _limits, _user_api and _prefixes 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}
|
|
prefixes = []
|
|
else:
|
|
if isinstance(limits, list):
|
|
# This should be a list of dicts of library info, for
|
|
# compatibility with the result from threadpool_info.
|
|
limits = {
|
|
lib_info["prefix"]: lib_info["num_threads"] for lib_info in limits
|
|
}
|
|
elif isinstance(limits, ThreadpoolController):
|
|
# To set the limits from the library controllers of a
|
|
# ThreadpoolController object.
|
|
limits = {
|
|
lib_controller.prefix: 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.
|
|
prefixes = [prefix for prefix in limits if prefix in _ALL_PREFIXES]
|
|
user_api = [api for api in limits if api in _ALL_USER_APIS]
|
|
|
|
return limits, user_api, prefixes
|
|
|
|
def _set_threadpool_limits(self):
|
|
"""Change the maximal number of threads in selected thread pools.
|
|
|
|
Return a list with all the supported libraries that have been found
|
|
matching `self._prefixes` and `self._user_api`.
|
|
"""
|
|
if self._limits is None:
|
|
return
|
|
|
|
for lib_controller in self._controller.lib_controllers:
|
|
# self._limits is a dict {key: num_threads} where key is either
|
|
# a prefix or a user_api. If a library matches both, the limit
|
|
# corresponding to the prefix is chosen.
|
|
if lib_controller.prefix in self._limits:
|
|
num_threads = self._limits[lib_controller.prefix]
|
|
elif lib_controller.user_api in self._limits:
|
|
num_threads = self._limits[lib_controller.user_api]
|
|
else:
|
|
continue
|
|
|
|
if num_threads is not None:
|
|
lib_controller.set_num_threads(num_threads)
|
|
|
|
|
|
class _ThreadpoolLimiterDecorator(_ThreadpoolLimiter, ContextDecorator):
|
|
"""Same as _ThreadpoolLimiter but to be used as a decorator"""
|
|
|
|
def __init__(self, controller, *, limits=None, user_api=None):
|
|
self._limits, self._user_api, self._prefixes = self._check_params(
|
|
limits, user_api
|
|
)
|
|
self._controller = 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()
|
|
return self
|
|
|
|
|
|
@_format_docstring(
|
|
USER_APIS=", ".join(f'"{api}"' for api in _ALL_USER_APIS),
|
|
BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES),
|
|
OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES),
|
|
)
|
|
class threadpool_limits(_ThreadpoolLimiter):
|
|
"""Change the maximal number of threads that can be used in thread pools.
|
|
|
|
This object can be used either as a callable (the construction of this object
|
|
limits the number of threads), as a context manager in a `with` block to
|
|
automatically restore the original state of the controlled libraries when exiting
|
|
the block, or as a decorator through its `wrap` method.
|
|
|
|
Set the maximal number of threads that can be used in thread pools used in
|
|
the supported libraries to `limit`. This function works for libraries that
|
|
are already loaded in the interpreter and can be changed dynamically.
|
|
|
|
This effect is global and impacts the whole Python process. There is no thread level
|
|
isolation as these libraries do not offer thread-local APIs to configure the number
|
|
of threads to use in nested parallel calls.
|
|
|
|
Parameters
|
|
----------
|
|
limits : int, dict, 'sequential_blas_under_openmp' or None (default=None)
|
|
The maximal number of threads that can be used in thread pools
|
|
|
|
- If int, sets the maximum number of threads to `limits` for each
|
|
library selected by `user_api`.
|
|
|
|
- If it is a dictionary `{{key: max_threads}}`, this function sets a
|
|
custom maximum number of threads for each `key` which can be either a
|
|
`user_api` or a `prefix` for a specific library.
|
|
|
|
- If 'sequential_blas_under_openmp', it will chose the appropriate `limits`
|
|
and `user_api` parameters for the specific use case of sequential BLAS
|
|
calls within an OpenMP parallel region. The `user_api` parameter is
|
|
ignored.
|
|
|
|
- If None, this function does not do anything.
|
|
|
|
user_api : {USER_APIS} or None (default=None)
|
|
APIs of libraries to limit. Used only if `limits` is an int.
|
|
|
|
- If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}).
|
|
|
|
- If "openmp", it will only limit OpenMP supported libraries
|
|
({OPENMP_LIBS}). Note that it can affect the number of threads used
|
|
by the BLAS libraries if they rely on OpenMP.
|
|
|
|
- If None, this function will apply to all supported libraries.
|
|
"""
|
|
|
|
def __init__(self, limits=None, user_api=None):
|
|
super().__init__(ThreadpoolController(), limits=limits, user_api=user_api)
|
|
|
|
@classmethod
|
|
def wrap(cls, limits=None, user_api=None):
|
|
return super().wrap(ThreadpoolController(), limits=limits, user_api=user_api)
|
|
|
|
|
|
class ThreadpoolController:
|
|
"""Collection of LibController objects for all loaded supported libraries
|
|
|
|
Attributes
|
|
----------
|
|
lib_controllers : list of `LibController` objects
|
|
The list of library controllers of all loaded supported libraries.
|
|
"""
|
|
|
|
# Cache for libc under POSIX and a few system libraries under Windows.
|
|
# We use a class level cache instead of an instance level cache because
|
|
# it's very unlikely that a shared library will be unloaded and reloaded
|
|
# during the lifetime of a program.
|
|
_system_libraries = dict()
|
|
|
|
def __init__(self):
|
|
self.lib_controllers = []
|
|
self._load_libraries()
|
|
self._warn_if_incompatible_openmp()
|
|
|
|
@classmethod
|
|
def _from_controllers(cls, lib_controllers):
|
|
new_controller = cls.__new__(cls)
|
|
new_controller.lib_controllers = lib_controllers
|
|
return new_controller
|
|
|
|
def info(self, debugging_info: bool = False):
|
|
"""Return lib_controllers info as a list of dicts.
|
|
|
|
Parameters
|
|
----------
|
|
debugging_info : bool
|
|
Include extra fields which require more intrusive actions to
|
|
obtain, and which can't always reliably be determined.
|
|
"""
|
|
return [
|
|
lib_controller.info(debugging_info=debugging_info)
|
|
for lib_controller in self.lib_controllers
|
|
]
|
|
|
|
def select(self, **kwargs):
|
|
"""Return a ThreadpoolController containing a subset of its current
|
|
library controllers
|
|
|
|
It will select all libraries matching at least one pair (key, value) from kwargs
|
|
where key is an entry of the library info dict (like "user_api", "internal_api",
|
|
"prefix", ...) and value is the value or a list of acceptable values for that
|
|
entry.
|
|
|
|
For instance, `ThreadpoolController().select(internal_api=["blis", "openblas"])`
|
|
will select all library controllers whose internal_api is either "blis" or
|
|
"openblas".
|
|
"""
|
|
for key, vals in kwargs.items():
|
|
kwargs[key] = [vals] if not isinstance(vals, list) else vals
|
|
|
|
lib_controllers = [
|
|
lib_controller
|
|
for lib_controller in self.lib_controllers
|
|
if any(
|
|
getattr(lib_controller, key, None) in vals
|
|
for key, vals in kwargs.items()
|
|
)
|
|
]
|
|
|
|
return ThreadpoolController._from_controllers(lib_controllers)
|
|
|
|
def _get_params_for_sequential_blas_under_openmp(self):
|
|
"""Return appropriate params to use for a sequential BLAS call in an OpenMP loop
|
|
|
|
This function takes into account the unexpected behavior of OpenBLAS with the
|
|
OpenMP threading layer.
|
|
"""
|
|
if self.select(
|
|
internal_api="openblas", threading_layer="openmp"
|
|
).lib_controllers:
|
|
return {"limits": None, "user_api": None}
|
|
return {"limits": 1, "user_api": "blas"}
|
|
|
|
@_format_docstring(
|
|
USER_APIS=", ".join('"{}"'.format(api) for api in _ALL_USER_APIS),
|
|
BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES),
|
|
OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES),
|
|
)
|
|
def limit(self, *, limits=None, user_api=None):
|
|
"""Change the maximal number of threads that can be used in thread pools.
|
|
|
|
This function returns an object that can be used either as a callable (the
|
|
construction of this object limits the number of threads) or as a context
|
|
manager, in a `with` block to automatically restore the original state of the
|
|
controlled libraries when exiting the block.
|
|
|
|
Set the maximal number of threads that can be used in thread pools used in
|
|
the supported libraries to `limits`. This function works for libraries that
|
|
are already loaded in the interpreter and can be changed dynamically.
|
|
|
|
This effect is global and impacts the whole Python process. There is no thread
|
|
level isolation as these libraries do not offer thread-local APIs to configure
|
|
the number of threads to use in nested parallel calls.
|
|
|
|
Parameters
|
|
----------
|
|
limits : int, dict, 'sequential_blas_under_openmp' or None (default=None)
|
|
The maximal number of threads that can be used in thread pools
|
|
|
|
- If int, sets the maximum number of threads to `limits` for each
|
|
library selected by `user_api`.
|
|
|
|
- If it is a dictionary `{{key: max_threads}}`, this function sets a
|
|
custom maximum number of threads for each `key` which can be either a
|
|
`user_api` or a `prefix` for a specific library.
|
|
|
|
- If 'sequential_blas_under_openmp', it will chose the appropriate `limits`
|
|
and `user_api` parameters for the specific use case of sequential BLAS
|
|
calls within an OpenMP parallel region. The `user_api` parameter is
|
|
ignored.
|
|
|
|
- If None, this function does not do anything.
|
|
|
|
user_api : {USER_APIS} or None (default=None)
|
|
APIs of libraries to limit. Used only if `limits` is an int.
|
|
|
|
- If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}).
|
|
|
|
- If "openmp", it will only limit OpenMP supported libraries
|
|
({OPENMP_LIBS}). Note that it can affect the number of threads used
|
|
by the BLAS libraries if they rely on OpenMP.
|
|
|
|
- If None, this function will apply to all supported libraries.
|
|
"""
|
|
return _ThreadpoolLimiter(self, limits=limits, user_api=user_api)
|
|
|
|
@_format_docstring(
|
|
USER_APIS=", ".join('"{}"'.format(api) for api in _ALL_USER_APIS),
|
|
BLAS_LIBS=", ".join(_ALL_BLAS_LIBRARIES),
|
|
OPENMP_LIBS=", ".join(_ALL_OPENMP_LIBRARIES),
|
|
)
|
|
def wrap(self, *, limits=None, user_api=None):
|
|
"""Change the maximal number of threads that can be used in thread pools.
|
|
|
|
This function returns an object that can be used as a decorator.
|
|
|
|
Set the maximal number of threads that can be used in thread pools used in
|
|
the supported libraries to `limits`. This function works for libraries that
|
|
are already loaded in the interpreter and can be changed dynamically.
|
|
|
|
Parameters
|
|
----------
|
|
limits : int, dict or None (default=None)
|
|
The maximal number of threads that can be used in thread pools
|
|
|
|
- If int, sets the maximum number of threads to `limits` for each
|
|
library selected by `user_api`.
|
|
|
|
- If it is a dictionary `{{key: max_threads}}`, this function sets a
|
|
custom maximum number of threads for each `key` which can be either a
|
|
`user_api` or a `prefix` for a specific library.
|
|
|
|
- If None, this function does not do anything.
|
|
|
|
user_api : {USER_APIS} or None (default=None)
|
|
APIs of libraries to limit. Used only if `limits` is an int.
|
|
|
|
- If "blas", it will only limit BLAS supported libraries ({BLAS_LIBS}).
|
|
|
|
- If "openmp", it will only limit OpenMP supported libraries
|
|
({OPENMP_LIBS}). Note that it can affect the number of threads used
|
|
by the BLAS libraries if they rely on OpenMP.
|
|
|
|
- If None, this function will apply to all supported libraries.
|
|
"""
|
|
return _ThreadpoolLimiter.wrap(self, limits=limits, user_api=user_api)
|
|
|
|
def __len__(self):
|
|
return len(self.lib_controllers)
|
|
|
|
def _load_libraries(self):
|
|
"""Loop through loaded shared libraries and store the supported ones"""
|
|
# ctypes.util is not imported on Linux: on CPython 3.14 it allocates a
|
|
# process-lifetime CFUNCTYPE callback that is not fork-safe with some
|
|
# libffi builds (#225). dllist also uses dl_iterate_phdr internally
|
|
# (#239), which we already avoid on Linux via /proc/self/maps.
|
|
dllist = None
|
|
if sys.platform not in ("linux", "emscripten"):
|
|
try:
|
|
from ctypes.util import dllist
|
|
except ImportError:
|
|
# CPython before 3.14 does not provide dll inspection.
|
|
dllist = None
|
|
|
|
if sys.platform == "linux" and os.path.exists("/proc/self/maps"):
|
|
# On glibc, dl_iterate_phdr has an internal lock, and that plus
|
|
# calling back into Python and the need to (re)acquire the GIL
|
|
# results in deadlocks. To avoid that, use a Linux-specific
|
|
# mechanism that doesn't have these issues; since it's Linux, musl
|
|
# works fine too.
|
|
self._find_libraries_with_linux()
|
|
elif dllist is not None:
|
|
# On Python 3.14+, this functionality is built-in. Once Python 3.13
|
|
# is no longer supported by threadpoolctl, most of the equivalent
|
|
# threadpoolctl implementations can be removed.
|
|
self._find_libraries_with_python(dllist)
|
|
elif sys.platform == "darwin":
|
|
self._find_libraries_with_dyld()
|
|
elif sys.platform == "win32":
|
|
self._find_libraries_on_windows()
|
|
elif "pyodide" in sys.modules:
|
|
self._find_libraries_pyodide()
|
|
else:
|
|
# Non-Linux Unix platforms.
|
|
self._find_libraries_with_dl_iterate_phdr()
|
|
|
|
def _find_libraries_with_linux(self):
|
|
"""Loop through loaded libraries and return binders on supported ones
|
|
|
|
Uses a Linux-specific mechanism:
|
|
https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html
|
|
"""
|
|
with open("/proc/self/maps") as f:
|
|
maps = f.read()
|
|
filepaths = set()
|
|
for line in maps.splitlines():
|
|
start_index = line.find("/")
|
|
if start_index == -1 or ".so" not in line:
|
|
continue
|
|
filepath = line[start_index:]
|
|
if os.path.exists(filepath):
|
|
filepaths.add(filepath)
|
|
|
|
for filepath in filepaths:
|
|
self._make_controller_from_path(filepath)
|
|
|
|
def _find_libraries_with_python(self, dllist):
|
|
"""Loop through loaded libraries and return binders on supported ones
|
|
|
|
Uses Python's built-in support for this functionality.
|
|
"""
|
|
try:
|
|
filepaths = dllist()
|
|
except OSError as exc:
|
|
# No Toolhelp fallback: if this fires in the wild, collect a
|
|
# minimal reproducer and report it upstream so CPython's dllist
|
|
# can be hardened against concurrent module list changes.
|
|
warnings.warn(
|
|
"ctypes.util.dllist failed to list loaded libraries "
|
|
f"({exc!r}). Native thread pools will not be inspected for "
|
|
"this ThreadpoolController. Please report a minimal "
|
|
"reproducer at https://github.com/joblib/threadpoolctl/issues "
|
|
"so it can be investigated and possibly be reported upstream to "
|
|
"CPython.",
|
|
RuntimeWarning,
|
|
)
|
|
return
|
|
if filepaths and filepaths[0] in ("", sys.executable):
|
|
filepaths = filepaths[1:]
|
|
for filepath in filepaths:
|
|
self._make_controller_from_path(filepath)
|
|
|
|
def _find_libraries_with_dl_iterate_phdr(self):
|
|
"""Loop through loaded libraries and return binders on supported ones
|
|
|
|
This function is expected to work on POSIX system only.
|
|
This code is adapted from code by Intel developer @anton-malakhov
|
|
available at https://github.com/IntelPython/smp
|
|
|
|
Copyright (c) 2017, Intel Corporation published under the BSD 3-Clause
|
|
license
|
|
"""
|
|
libc = self._get_libc()
|
|
if not hasattr(libc, "dl_iterate_phdr"): # pragma: no cover
|
|
warnings.warn(
|
|
"Could not find dl_iterate_phdr in the C standard library.",
|
|
RuntimeWarning,
|
|
)
|
|
return []
|
|
|
|
filepaths = []
|
|
|
|
# Callback function for `dl_iterate_phdr` which is called for every
|
|
# library loaded in the current process until it returns 1. To minimize
|
|
# the potential for deadlocks (see #228), this code should not do
|
|
# anything that might result in reentrancy into the library, the dl
|
|
# system, or anything else.
|
|
def match_library_callback(info, size, data):
|
|
# Get the path of the current library
|
|
filepath = info.contents.dlpi_name
|
|
if filepath:
|
|
filepaths.append(filepath)
|
|
return 0
|
|
|
|
c_func_signature = ctypes.CFUNCTYPE(
|
|
ctypes.c_int, # Return type
|
|
ctypes.POINTER(_dl_phdr_info),
|
|
ctypes.c_size_t,
|
|
ctypes.c_char_p,
|
|
)
|
|
c_match_library_callback = c_func_signature(match_library_callback)
|
|
|
|
data = ctypes.c_char_p(b"")
|
|
libc.dl_iterate_phdr(c_match_library_callback, data)
|
|
|
|
# Now that a list of filepaths is available, load the respective
|
|
# libraries:
|
|
for filepath in filepaths:
|
|
# Store the library controller if it is supported and selected
|
|
self._make_controller_from_path(filepath.decode("utf-8"))
|
|
|
|
def _find_libraries_with_dyld(self):
|
|
"""Loop through loaded libraries and return binders on supported ones
|
|
|
|
This function is expected to work on OSX system only
|
|
"""
|
|
libc = self._get_libc()
|
|
if not hasattr(libc, "_dyld_image_count"): # pragma: no cover
|
|
warnings.warn(
|
|
"Could not find _dyld_image_count in the C standard library.",
|
|
RuntimeWarning,
|
|
)
|
|
return []
|
|
|
|
n_dyld = libc._dyld_image_count()
|
|
libc._dyld_get_image_name.restype = ctypes.c_char_p
|
|
|
|
for i in range(n_dyld):
|
|
filepath = ctypes.string_at(libc._dyld_get_image_name(i))
|
|
filepath = filepath.decode("utf-8")
|
|
|
|
# Store the library controller if it is supported and selected
|
|
self._make_controller_from_path(filepath)
|
|
|
|
def _find_libraries_on_windows(self):
|
|
"""Loop through loaded libraries and return binders on supported ones
|
|
|
|
This function is expected to work on windows system only.
|
|
|
|
Used when ``ctypes.util.dllist`` is unavailable (Python < 3.14).
|
|
Module discovery uses a snapshot-first strategy:
|
|
``CreateToolhelp32Snapshot`` provides an atomic list of loaded
|
|
modules, which is more robust than ``EnumProcessModulesEx`` under
|
|
concurrent DLL load/unload. Paths that fit in
|
|
``MODULEENTRY32W.szExePath`` (shorter than ``MAX_PATH``) are used as-is.
|
|
Truncated or empty snapshot paths are resolved with
|
|
``GetModuleFileNameW`` and, when needed, ``GetModuleFileNameExW`` with a
|
|
larger buffer. If snapshot creation fails, enumeration falls back to
|
|
``EnumProcessModulesEx``.
|
|
"""
|
|
from ctypes.wintypes import MAX_PATH
|
|
|
|
ps_api = self._get_windll("Psapi")
|
|
kernel_32 = self._get_windll("kernel32")
|
|
self._setup_windows_module_apis(ps_api, kernel_32)
|
|
|
|
h_process = kernel_32.GetCurrentProcess()
|
|
|
|
# Allocate a buffer for long path names; see _WINDOWS_MAX_LIBRARY_PATH_LENGTH.
|
|
max_path = _WINDOWS_MAX_LIBRARY_PATH_LENGTH
|
|
path_buf = ctypes.create_unicode_buffer(max_path)
|
|
|
|
try:
|
|
modules = self._snapshot_loaded_modules(kernel_32)
|
|
except OSError:
|
|
modules = None
|
|
|
|
if modules is not None:
|
|
for h_module, snapshot_path in modules:
|
|
if snapshot_path and len(snapshot_path) < MAX_PATH - 1:
|
|
filepath = snapshot_path
|
|
else:
|
|
filepath = self._resolve_module_filepath(
|
|
ps_api,
|
|
kernel_32,
|
|
h_process,
|
|
h_module,
|
|
max_path=max_path,
|
|
path_buf=path_buf,
|
|
)
|
|
if filepath is not None:
|
|
self._make_controller_from_path(filepath)
|
|
else:
|
|
self._find_libraries_with_enum_process_modules_ex(
|
|
ps_api, kernel_32, h_process, max_path, path_buf
|
|
)
|
|
|
|
@classmethod
|
|
def _setup_windows_module_apis(cls, ps_api, kernel_32):
|
|
"""Set ctypes signatures for Windows module enumeration APIs."""
|
|
from ctypes.wintypes import BOOL, DWORD, HANDLE, HMODULE
|
|
|
|
if getattr(cls, "_windows_module_apis_configured", False):
|
|
return
|
|
|
|
kernel_32.GetCurrentProcess.restype = HANDLE
|
|
kernel_32.CreateToolhelp32Snapshot.argtypes = [DWORD, DWORD]
|
|
kernel_32.CreateToolhelp32Snapshot.restype = HANDLE
|
|
kernel_32.Module32FirstW.argtypes = [HANDLE, ctypes.c_void_p]
|
|
kernel_32.Module32FirstW.restype = BOOL
|
|
kernel_32.Module32NextW.argtypes = [HANDLE, ctypes.c_void_p]
|
|
kernel_32.Module32NextW.restype = BOOL
|
|
kernel_32.GetModuleFileNameW.argtypes = [HMODULE, ctypes.c_wchar_p, DWORD]
|
|
kernel_32.GetModuleFileNameW.restype = DWORD
|
|
kernel_32.CloseHandle.argtypes = [HANDLE]
|
|
kernel_32.CloseHandle.restype = BOOL
|
|
|
|
ps_api.EnumProcessModulesEx.argtypes = [
|
|
HANDLE,
|
|
ctypes.POINTER(HMODULE),
|
|
DWORD,
|
|
ctypes.POINTER(DWORD),
|
|
DWORD,
|
|
]
|
|
ps_api.EnumProcessModulesEx.restype = BOOL
|
|
ps_api.GetModuleFileNameExW.argtypes = [
|
|
HANDLE,
|
|
HMODULE,
|
|
ctypes.c_wchar_p,
|
|
DWORD,
|
|
]
|
|
ps_api.GetModuleFileNameExW.restype = DWORD
|
|
|
|
cls._windows_module_apis_configured = True
|
|
|
|
def _snapshot_loaded_modules(self, kernel_32):
|
|
"""Return loaded modules as (hModule, snapshot_path) pairs.
|
|
|
|
Uses CreateToolhelp32Snapshot for an atomic view of loaded modules, which
|
|
is more robust than EnumProcessModulesEx when DLLs are loaded or unloaded
|
|
concurrently. ``ERROR_BAD_LENGTH`` is retried a bounded number of times
|
|
(the documented transient race when the module list changes mid-snapshot)
|
|
and then raised as ``OSError`` so the caller can fall back.
|
|
"""
|
|
from ctypes.wintypes import DWORD, HANDLE, MAX_PATH
|
|
|
|
class MODULEENTRY32W(ctypes.Structure):
|
|
_fields_ = [
|
|
("dwSize", DWORD),
|
|
("th32ModuleID", DWORD),
|
|
("th32ProcessID", DWORD),
|
|
("GlblcntUsage", DWORD),
|
|
("ProccntUsage", DWORD),
|
|
("modBaseAddr", ctypes.POINTER(ctypes.c_byte)),
|
|
("modBaseSize", DWORD),
|
|
("hModule", HANDLE),
|
|
("szModule", ctypes.c_wchar * 256),
|
|
("szExePath", ctypes.c_wchar * MAX_PATH),
|
|
]
|
|
|
|
TH32CS_SNAPMODULE = 0x00000008
|
|
TH32CS_SNAPMODULE32 = 0x00000010
|
|
ERROR_BAD_LENGTH = 0x0018
|
|
ERROR_NO_MORE_FILES = 0x0012
|
|
INVALID_HANDLE_VALUE = HANDLE(-1).value
|
|
max_snapshot_retries = 16
|
|
|
|
for _ in range(max_snapshot_retries):
|
|
snap_handle = kernel_32.CreateToolhelp32Snapshot(
|
|
TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32, os.getpid()
|
|
)
|
|
if snap_handle != INVALID_HANDLE_VALUE:
|
|
break
|
|
err = ctypes.get_last_error()
|
|
if err == ERROR_BAD_LENGTH:
|
|
continue
|
|
msg = ctypes.FormatError(err).strip()
|
|
raise OSError(f"CreateToolhelp32Snapshot failed: {msg}")
|
|
else:
|
|
msg = ctypes.FormatError(ERROR_BAD_LENGTH).strip()
|
|
raise OSError(f"CreateToolhelp32Snapshot failed: {msg}")
|
|
|
|
modules = []
|
|
try:
|
|
lib_entry = MODULEENTRY32W()
|
|
lib_entry.dwSize = ctypes.sizeof(MODULEENTRY32W)
|
|
if not kernel_32.Module32FirstW(
|
|
snap_handle, ctypes.byref(lib_entry)
|
|
): # pragma: no cover
|
|
err = ctypes.get_last_error()
|
|
msg = ctypes.FormatError(err).strip()
|
|
raise OSError(f"Module32FirstW failed: {msg}")
|
|
|
|
while True:
|
|
modules.append((lib_entry.hModule, lib_entry.szExePath))
|
|
if not kernel_32.Module32NextW(snap_handle, ctypes.byref(lib_entry)):
|
|
err = ctypes.get_last_error()
|
|
if err != ERROR_NO_MORE_FILES: # pragma: no cover
|
|
msg = ctypes.FormatError(err).strip()
|
|
raise OSError(f"Module32NextW failed: {msg}")
|
|
break
|
|
finally:
|
|
kernel_32.CloseHandle(snap_handle)
|
|
|
|
return modules
|
|
|
|
def _resolve_module_filepath(
|
|
self,
|
|
ps_api,
|
|
kernel_32,
|
|
h_process,
|
|
h_module,
|
|
max_path,
|
|
path_buf,
|
|
):
|
|
"""Return the full path for a module, or None if it should be skipped."""
|
|
from ctypes.wintypes import MAX_PATH
|
|
|
|
n_size = kernel_32.GetModuleFileNameW(h_module, path_buf, MAX_PATH)
|
|
if n_size and n_size < MAX_PATH - 1:
|
|
return path_buf.value
|
|
|
|
n_size = ps_api.GetModuleFileNameExW(h_process, h_module, path_buf, max_path)
|
|
if n_size and n_size < max_path - 1:
|
|
return path_buf.value
|
|
if n_size: # pragma: no cover
|
|
filepath = path_buf.value
|
|
warnings.warn(
|
|
"Could not get the full path of a dynamic library (path too "
|
|
"long). This library will be ignored and threadpoolctl might "
|
|
"not be able to control or display information about all "
|
|
f"loaded libraries. Here's the truncated path: {filepath!r}",
|
|
RuntimeWarning,
|
|
)
|
|
return None
|
|
|
|
err = ctypes.get_last_error()
|
|
msg = ctypes.FormatError(err).strip()
|
|
warnings.warn(
|
|
"Could not get the path of a dynamic library "
|
|
"(GetModuleFileNameW and GetModuleFileNameExW failed). "
|
|
"This library will be ignored and threadpoolctl might not be "
|
|
"able to control or display information about all loaded "
|
|
f"libraries. {msg}",
|
|
RuntimeWarning,
|
|
)
|
|
return None
|
|
|
|
def _find_libraries_with_enum_process_modules_ex(
|
|
self, ps_api, kernel_32, h_process, max_path, path_buf
|
|
):
|
|
"""Fallback Windows module enumeration using EnumProcessModulesEx.
|
|
|
|
This code is adapted from code by Philipp Hagemeister @phihag available
|
|
at https://stackoverflow.com/questions/17474574
|
|
"""
|
|
from ctypes.wintypes import DWORD, HMODULE
|
|
|
|
LIST_LIBRARIES_ALL = 0x03
|
|
|
|
buf_count = 256
|
|
needed = DWORD()
|
|
# Grow the buffer until it becomes large enough to hold all the
|
|
# module headers
|
|
while True:
|
|
buf = (HMODULE * buf_count)()
|
|
buf_size = ctypes.sizeof(buf)
|
|
if not ps_api.EnumProcessModulesEx(
|
|
h_process,
|
|
buf,
|
|
buf_size,
|
|
ctypes.byref(needed),
|
|
LIST_LIBRARIES_ALL,
|
|
):
|
|
raise OSError("EnumProcessModulesEx failed")
|
|
if buf_size >= needed.value:
|
|
break
|
|
buf_count = needed.value // (buf_size // buf_count)
|
|
|
|
count = needed.value // (buf_size // buf_count)
|
|
for h_module in map(HMODULE, buf[:count]):
|
|
filepath = self._resolve_module_filepath(
|
|
ps_api,
|
|
kernel_32,
|
|
h_process,
|
|
h_module,
|
|
max_path=max_path,
|
|
path_buf=path_buf,
|
|
)
|
|
if filepath is not None:
|
|
self._make_controller_from_path(filepath)
|
|
|
|
def _find_libraries_pyodide(self):
|
|
"""Pyodide specific implementation for finding loaded libraries.
|
|
|
|
Adapted from suggestion in https://github.com/joblib/threadpoolctl/pull/169#issuecomment-1946696449.
|
|
|
|
One day, we may have a simpler solution. libc dl_iterate_phdr needs to
|
|
be implemented in Emscripten and exposed in Pyodide, see
|
|
https://github.com/emscripten-core/emscripten/issues/21354 for more
|
|
details.
|
|
"""
|
|
try:
|
|
from pyodide_js._module import LDSO
|
|
except ImportError:
|
|
warnings.warn(
|
|
"Unable to import LDSO from pyodide_js._module. This should never "
|
|
"happen."
|
|
)
|
|
return
|
|
|
|
if hasattr(LDSO.loadedLibsByName, "as_py_json"): # Pyodide >= 0.29
|
|
libs_iter = LDSO.loadedLibsByName.as_py_json()
|
|
else:
|
|
libs_iter = LDSO.loadedLibsByName.as_object_map() # Pyodide < 0.29
|
|
|
|
for filepath in libs_iter:
|
|
# Some libraries are duplicated by Pyodide and do not exist in the
|
|
# filesystem, so we first check for the existence of the file. For
|
|
# more details, see
|
|
# https://github.com/joblib/threadpoolctl/pull/169#issuecomment-1947946728
|
|
if os.path.exists(filepath):
|
|
self._make_controller_from_path(filepath)
|
|
|
|
def _make_controller_from_path(self, filepath):
|
|
"""Store a library controller if it is supported and selected"""
|
|
# Required to resolve symlinks
|
|
filepath = _realpath(filepath)
|
|
# `lower` required to take account of OpenMP dll case on Windows
|
|
# (vcomp, VCOMP, Vcomp, ...)
|
|
filename = os.path.basename(filepath).lower()
|
|
|
|
# Loop through supported libraries to find if this filename corresponds
|
|
# to a supported one.
|
|
for controller_class in _ALL_CONTROLLERS:
|
|
# check if filename matches a supported prefix
|
|
prefix = self._check_prefix(filename, controller_class.filename_prefixes)
|
|
|
|
# filename does not match any of the prefixes of the candidate
|
|
# library. move to next library.
|
|
if prefix is None:
|
|
continue
|
|
|
|
# Legacy workaround for BLAS libraries that conda-forge used to expose
|
|
# on Windows as libblas.dll, disambiguated via implementation-specific
|
|
# symbols. Current conda-forge stacks no longer load libblas.dll (e.g.
|
|
# MKL is exposed as mkl_rt.<version>.dll instead), so this path is
|
|
# kept for older installs but cannot be exercised in today's CI.
|
|
if prefix == "libblas":
|
|
if filename.endswith(".dll"):
|
|
libblas = ctypes.CDLL(filepath, _RTLD_NOLOAD)
|
|
if not any(
|
|
hasattr(libblas, func)
|
|
for func in controller_class.check_symbols
|
|
):
|
|
continue
|
|
else:
|
|
# Non-Windows libblas DSOs (e.g. from openblas) lack the symbols
|
|
# needed to instantiate a controller and would duplicate entries.
|
|
continue
|
|
|
|
# filename matches a prefix. Now we check if the library has the symbols we
|
|
# are looking for. If none of the symbols exists, it's very likely not the
|
|
# expected library (e.g. a library having a common prefix with one of the
|
|
# our supported libraries). Otherwise, create and store the library
|
|
# controller.
|
|
lib_controller = controller_class(
|
|
filepath=filepath, prefix=prefix, parent=self
|
|
)
|
|
|
|
if filepath in (lib.filepath for lib in self.lib_controllers):
|
|
# We already have a controller for this library.
|
|
continue
|
|
|
|
if not hasattr(controller_class, "check_symbols") or any(
|
|
hasattr(lib_controller.dynlib, func)
|
|
for func in controller_class.check_symbols
|
|
):
|
|
self.lib_controllers.append(lib_controller)
|
|
|
|
def _check_prefix(self, library_basename, filename_prefixes):
|
|
"""Return the prefix library_basename starts with
|
|
|
|
Return None if none matches.
|
|
"""
|
|
for prefix in filename_prefixes:
|
|
if library_basename.startswith(prefix):
|
|
return prefix
|
|
return None
|
|
|
|
def _warn_if_incompatible_openmp(self):
|
|
"""Raise a warning if llvm-OpenMP and intel-OpenMP are both loaded"""
|
|
if sys.platform != "linux":
|
|
# The incompatibility between libomp and libiomp is known to cause
|
|
# crashes on Linux. On other platforms, conda-forge may expose both
|
|
# libraries without the same runtime conflict (for instance when
|
|
# libiomp forwards to libomp on Windows).
|
|
return
|
|
|
|
prefixes = [lib_controller.prefix for lib_controller in self.lib_controllers]
|
|
msg = textwrap.dedent(
|
|
"""
|
|
Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
|
|
the same time. Both libraries are known to be incompatible and this
|
|
can cause random crashes or deadlocks on Linux when loaded in the
|
|
same Python program.
|
|
Using threadpoolctl may cause crashes or deadlocks. For more
|
|
information and possible workarounds, please see
|
|
https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md
|
|
"""
|
|
)
|
|
if "libomp" in prefixes and "libiomp" in prefixes:
|
|
warnings.warn(msg, RuntimeWarning)
|
|
|
|
@classmethod
|
|
def _get_libc(cls):
|
|
"""Load the lib-C for unix systems."""
|
|
libc = cls._system_libraries.get("libc")
|
|
if libc is None:
|
|
# dlopen(NULL) rather than ctypes.util.find_library("c"). Importing
|
|
# ctypes.util on CPython 3.14 Linux creates a process-lifetime
|
|
# CFUNCTYPE callback that is not fork-safe with some libffi builds
|
|
# (issue #225). If libc is statically linked or Python is linked
|
|
# against musl, the main program still exports the libc symbols we
|
|
# need. If it does not, we check for those symbols later anyway.
|
|
libc = ctypes.CDLL(None, mode=_RTLD_NOLOAD)
|
|
cls._system_libraries["libc"] = libc
|
|
return libc
|
|
|
|
@classmethod
|
|
def _get_windll(cls, dll_name):
|
|
"""Load a windows DLL"""
|
|
dll = cls._system_libraries.get(dll_name)
|
|
if dll is None:
|
|
dll = ctypes.WinDLL(f"{dll_name}.dll", use_last_error=True)
|
|
cls._system_libraries[dll_name] = dll
|
|
return dll
|
|
|
|
|
|
def _main():
|
|
"""Commandline interface to display thread-pool information and exit."""
|
|
import argparse
|
|
import importlib
|
|
import json
|
|
import sys
|
|
|
|
parser = argparse.ArgumentParser(
|
|
usage="python -m threadpoolctl -i numpy scipy.linalg xgboost",
|
|
description="Display thread-pool information and exit.",
|
|
)
|
|
parser.add_argument(
|
|
"-i",
|
|
"--import",
|
|
dest="modules",
|
|
nargs="*",
|
|
default=(),
|
|
help="Python modules to import before introspecting thread-pools.",
|
|
)
|
|
parser.add_argument(
|
|
"-c",
|
|
"--command",
|
|
help="a Python statement to execute before introspecting thread-pools.",
|
|
)
|
|
|
|
options = parser.parse_args(sys.argv[1:])
|
|
for module in options.modules:
|
|
try:
|
|
importlib.import_module(module, package=None)
|
|
except ImportError:
|
|
print("WARNING: could not import", module, file=sys.stderr)
|
|
|
|
if options.command:
|
|
exec(options.command)
|
|
|
|
print(json.dumps(threadpool_info(debugging_info=True), indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
_main()
|