Aluode/PerceptionLabPortable

1"""Customizations of :mod:`joblib` and :mod:`threadpoolctl` tools for scikit-learn
2usage.
3"""
4
5# Authors: The scikit-learn developers
6# SPDX-License-Identifier: BSD-3-Clause
7
8import functools
9import warnings
10from functools import update_wrapper
11
12import joblib
13from threadpoolctl import ThreadpoolController
14
15from .._config import config_context, get_config
16
17# Global threadpool controller instance that can be used to locally limit the number of
18# threads without looping through all shared libraries every time.
19# It should not be accessed directly and _get_threadpool_controller should be used
20# instead.
21_threadpool_controller = None
22
23
24def _with_config_and_warning_filters(delayed_func, config, warning_filters):
25    """Helper function that intends to attach a config to a delayed function."""
26    if hasattr(delayed_func, "with_config_and_warning_filters"):
27        return delayed_func.with_config_and_warning_filters(config, warning_filters)
28    else:
29        warnings.warn(
30            (
31                "`sklearn.utils.parallel.Parallel` needs to be used in "
32                "conjunction with `sklearn.utils.parallel.delayed` instead of "
33                "`joblib.delayed` to correctly propagate the scikit-learn "
34                "configuration to the joblib workers."
35            ),
36            UserWarning,
37        )
38        return delayed_func
39
40
41class Parallel(joblib.Parallel):
42    """Tweak of :class:`joblib.Parallel` that propagates the scikit-learn configuration.
43
44    This subclass of :class:`joblib.Parallel` ensures that the active configuration
45    (thread-local) of scikit-learn is propagated to the parallel workers for the
46    duration of the execution of the parallel tasks.
47
48    The API does not change and you can refer to :class:`joblib.Parallel`
49    documentation for more details.
50
51    .. versionadded:: 1.3
52    """
53
54    def __call__(self, iterable):
55        """Dispatch the tasks and return the results.
56
57        Parameters
58        ----------
59        iterable : iterable
60            Iterable containing tuples of (delayed_function, args, kwargs) that should
61            be consumed.
62
63        Returns
64        -------
65        results : list
66            List of results of the tasks.
67        """
68        # Capture the thread-local scikit-learn configuration at the time
69        # Parallel.__call__ is issued since the tasks can be dispatched
70        # in a different thread depending on the backend and on the value of
71        # pre_dispatch and n_jobs.
72        config = get_config()
73        warning_filters = warnings.filters
74        iterable_with_config_and_warning_filters = (
75            (
76                _with_config_and_warning_filters(delayed_func, config, warning_filters),
77                args,
78                kwargs,
79            )
80            for delayed_func, args, kwargs in iterable
81        )
82        return super().__call__(iterable_with_config_and_warning_filters)
83
84
85# remove when https://github.com/joblib/joblib/issues/1071 is fixed
86def delayed(function):
87    """Decorator used to capture the arguments of a function.
88
89    This alternative to `joblib.delayed` is meant to be used in conjunction
90    with `sklearn.utils.parallel.Parallel`. The latter captures the scikit-
91    learn configuration by calling `sklearn.get_config()` in the current
92    thread, prior to dispatching the first task. The captured configuration is
93    then propagated and enabled for the duration of the execution of the
94    delayed function in the joblib workers.
95
96    .. versionchanged:: 1.3
97       `delayed` was moved from `sklearn.utils.fixes` to `sklearn.utils.parallel`
98       in scikit-learn 1.3.
99
100    Parameters
101    ----------
102    function : callable
103        The function to be delayed.
104
105    Returns
106    -------
107    output: tuple
108        Tuple containing the delayed function, the positional arguments, and the
109        keyword arguments.
110    """
111
112    @functools.wraps(function)
113    def delayed_function(*args, **kwargs):
114        return _FuncWrapper(function), args, kwargs
115
116    return delayed_function
117
118
119class _FuncWrapper:
120    """Load the global configuration before calling the function."""
121
122    def __init__(self, function):
123        self.function = function
124        update_wrapper(self, self.function)
125
126    def with_config_and_warning_filters(self, config, warning_filters):
127        self.config = config
128        self.warning_filters = warning_filters
129        return self
130
131    def __call__(self, *args, **kwargs):
132        config = getattr(self, "config", {})
133        warning_filters = getattr(self, "warning_filters", [])
134        if not config or not warning_filters:
135            warnings.warn(
136                (
137                    "`sklearn.utils.parallel.delayed` should be used with"
138                    " `sklearn.utils.parallel.Parallel` to make it possible to"
139                    " propagate the scikit-learn configuration of the current thread to"
140                    " the joblib workers."
141                ),
142                UserWarning,
143            )
144
145        with config_context(**config), warnings.catch_warnings():
146            warnings.filters = warning_filters
147            return self.function(*args, **kwargs)
148
149
150def _get_threadpool_controller():
151    """Return the global threadpool controller instance."""
152    global _threadpool_controller
153
154    if _threadpool_controller is None:
155        _threadpool_controller = ThreadpoolController()
156
157    return _threadpool_controller
158
159
160def _threadpool_controller_decorator(limits=1, user_api="blas"):
161    """Decorator to limit the number of threads used at the function level.
162
163    It should be preferred over `threadpoolctl.ThreadpoolController.wrap` because this
164    one only loads the shared libraries when the function is called while the latter
165    loads them at import time.
166    """
167
168    def decorator(func):
169        @functools.wraps(func)
170        def wrapper(*args, **kwargs):
171            controller = _get_threadpool_controller()
172            with controller.limit(limits=limits, user_api=user_api):
173                return func(*args, **kwargs)
174
175        return wrapper
176
177    return decorator
178