Threadpoolctl

from threadpoolctl import threadpool_info, threadpool_limits, ThreadpoolController

from threadpoolctl import LibController, register

from threadpoolctl import threadpool_info, threadpool_limits

# Inspect current thread pool state
info = threadpool_info()
print(f"Found {len(info)} libraries with thread pools")

# Limit threads temporarily with context manager
with threadpool_limits(limits=1):
    # All thread pools limited to 1 thread
    # Your computation here
    pass

# Limit specific API types  
with threadpool_limits(limits=2, user_api='blas'):
    # Only BLAS libraries limited to 2 threads
    # Your computation here
    pass

# Use as decorator
@threadpool_limits(limits=1)
def compute():
    # Function runs with thread pools limited to 1 thread
    pass

def threadpool_info() -> list[dict[str, Any]]:
    """
    Return thread pool information for all detected libraries.
    
    Scans loaded shared libraries and returns information for supported
    thread pool libraries (BLAS implementations and OpenMP runtimes).
    
    Returns:
        list[dict[str, Any]]: List of library info dictionaries with keys:
            - user_api: str ('blas', 'openmp')  
            - internal_api: str ('openblas', 'mkl', 'blis', 'flexiblas', 'openmp')
            - prefix: str (library filename prefix)
            - filepath: str (path to loaded shared library)
            - version: str | None (library version if available)
            - num_threads: int (current thread limit)
            - Additional implementation-specific attributes
    """

class threadpool_limits:
    """
    Context manager/decorator for limiting thread pool sizes.
    
    Can be used as callable, context manager, or decorator.
    
    Args:
        limits: int | dict[str, int] | str | None
            - int: Global thread limit for all libraries
            - dict: Per-API or per-library limits {api/prefix: limit}
            - 'sequential_blas_under_openmp': Special case for nested parallelism
            - None: No limits applied
        user_api: str | None ('blas', 'openmp', None)
            API type to limit when limits is int
    """
    
    def __init__(self, limits=None, user_api=None): ...
    
    def __enter__(self) -> 'threadpool_limits':
        """Enter context manager, returns self."""
        
    def __exit__(self, type, value, traceback) -> None:
        """Exit context manager, restores original limits."""
        
    def restore_original_limits(self) -> None:
        """Manually restore original thread limits."""
        
    def unregister(self) -> None:
        """Alias for restore_original_limits() (backward compatibility)."""
        
    def get_original_num_threads(self) -> dict[str, int]:
        """
        Get original thread counts before limiting.
        
        Returns:
            dict[str, int]: Original thread counts by user_api
        """
    
    @classmethod
    def wrap(cls, limits=None, user_api=None) -> 'Callable':
        """
        Create decorator version that delays limit setting.
        
        Returns:
            Callable: Decorator function for use with @threadpool_limits.wrap(...)
        """

class ThreadpoolController:
    """
    Controller for all loaded supported thread pool libraries.
    
    Automatically discovers supported libraries when instantiated and provides
    methods for filtering, selecting, and controlling their thread settings.
    
    Attributes:
        lib_controllers: list[LibController] - Individual library controllers
    """
    
    def __init__(self):
        """Initialize controller and discover loaded libraries."""
        
    def info(self) -> list[dict[str, Any]]:
        """
        Get information about all controlled libraries.
        
        Returns:
            list[dict[str, Any]]: List of library info dicts
        """
        
    def select(self, **kwargs) -> 'ThreadpoolController':
        """
        Filter libraries by attributes.
        
        Args:
            **kwargs: Attribute filters where key is attribute name and value
                     is either a single value or list of acceptable values.
                     
        Returns:
            ThreadpoolController: New controller with filtered library subset
        """
        
    def limit(self, *, limits=None, user_api=None) -> '_ThreadpoolLimiter':
        """
        Create thread limiter for this controller's libraries.
        
        Args:
            limits: int | dict[str, int] | str | None - Thread limit specification
            user_api: str | None - API type filter
            
        Returns:
            _ThreadpoolLimiter: Context manager for temporary limiting
        """
        
    def wrap(self, *, limits=None, user_api=None) -> 'Callable':
        """
        Create decorator for this controller's libraries.
        
        Args:
            limits: int | dict[str, int] | str | None - Thread limit specification  
            user_api: str | None - API type filter
            
        Returns:
            Callable: Decorator function for thread limiting
        """
        
    def __len__(self) -> int:
        """
        Number of controlled libraries.
        
        Returns:
            int: Count of lib_controllers
        """

class LibController:
    """
    Abstract base class for individual library controllers.
    
    Class Attributes:
        user_api: str - Standardized API name ('blas', 'openmp', etc.)
        internal_api: str - Implementation-specific identifier  
        filename_prefixes: tuple[str, ...] - Library filename prefixes to match
        check_symbols: tuple[str, ...] - Optional symbol names to verify compatibility
    """
    
    def __init__(self, *, filepath=None, prefix=None, parent=None):
        """
        Initialize library controller.
        
        Args:
            filepath: str | None - Path to shared library file
            prefix: str | None - Matched filename prefix
            parent: ThreadpoolController | None - Parent controller instance
        """
    
    def info(self) -> dict[str, Any]:
        """
        Return library information dictionary.
        
        Returns:
            dict: Library info with standard keys plus additional attributes
        """
    
    def get_num_threads(self) -> int | None:
        """
        Get current maximum thread count (abstract method).
        
        Returns:
            int | None: Current thread limit, None if unavailable
        """
    
    def set_num_threads(self, num_threads: int) -> Any:
        """
        Set maximum thread count (abstract method).
        
        Args:
            num_threads: int - New thread limit
            
        Returns:
            Any: Implementation-specific return value
        """
    
    def get_version(self) -> str | None:
        """
        Get library version (abstract method).
        
        Returns:
            str | None: Version string, None if unavailable
        """
    
    def set_additional_attributes(self) -> None:
        """
        Set implementation-specific attributes.
        
        Called during initialization to set custom attributes that will
        be included in the info() dictionary.
        """
    
    @property
    def num_threads(self) -> int | None:
        """Current thread limit (dynamic property)."""

def register(controller: type[LibController]) -> None:
    """
    Register a new library controller class.
    
    Adds the controller to the global registry so it will be used
    during library discovery in future ThreadpoolController instances.
    
    Args:
        controller: type[LibController] - LibController subclass to register
    """

# Basic usage
python -m threadpoolctl

# Import specific modules before introspection
python -m threadpoolctl -i numpy scipy.linalg xgboost

# Execute Python code before introspection  
python -m threadpoolctl -c "import sklearn; from sklearn.ensemble import RandomForestClassifier"

# Combined usage
python -m threadpoolctl -i numpy -c "import threading; print(f'Python threads: {threading.active_count()}')"

[
  {
    "filepath": "/usr/lib/libmkl_rt.so",
    "prefix": "libmkl_rt", 
    "user_api": "blas",
    "internal_api": "mkl",
    "version": "2019.0.4",
    "num_threads": 2,
    "threading_layer": "intel"
  }
]

# Library information dictionary structure
LibraryInfo = dict[str, Any]  # Keys: user_api, internal_api, prefix, filepath, version, num_threads, ...

# Thread limit specifications
ThreadLimits = int | dict[str, int] | str | None

# API type specifications  
UserAPI = Literal["blas", "openmp"] | None