Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1###############################################################################
2# Reusable ProcessPoolExecutor
3#
4# author: Thomas Moreau and Olivier Grisel
5#
7from __future__ import annotations
9import time
10import warnings
11import threading
12import multiprocessing as mp
13from typing import Any
15from .process_executor import ProcessPoolExecutor, EXTRA_QUEUED_CALLS
16from .backend.context import cpu_count
17from .backend import get_context
19__all__ = ["get_reusable_executor"]
22class _ExecutorStorage(threading.local):
23 """Cache a thread-local executor"""
25 executor: _ReusablePoolExecutor | None = None
26 executor_kwargs: dict[str, Any] | None = None
29# Lock for id management:
30_executor_lock = threading.RLock()
31_next_executor_id = 0
33# Thread-local executor storage:
34_executor_storage = _ExecutorStorage()
37def _get_next_executor_id():
38 """Ensure that each successive executor instance has a unique, monotonic id.
40 The purpose of this monotonic id is to help debug and test automated
41 instance creation.
42 """
43 global _next_executor_id
44 with _executor_lock:
45 executor_id = _next_executor_id
46 _next_executor_id += 1
47 return executor_id
50def get_reusable_executor(
51 max_workers=None,
52 context=None,
53 timeout=10,
54 kill_workers=False,
55 reuse="auto",
56 job_reducers=None,
57 result_reducers=None,
58 initializer=None,
59 initargs=(),
60 env=None,
61):
62 """Return the current ReusableExecutor instance for the current thread.
64 Start a new instance if it has not been started already or if the previous
65 instance was left in a broken state.
67 If the previous instance does not have the requested number of workers, the
68 executor is dynamically resized to adjust the number of workers prior to
69 returning.
71 Reusing an instance (per thread) spares the overhead of starting new
72 worker processes and importing common python packages each time.
74 ``max_workers`` controls the maximum number of tasks that can be running in
75 parallel in worker processes. By default this is set to the number of
76 CPUs on the host.
78 Setting ``timeout`` (in seconds) makes idle workers automatically shutdown
79 so as to release system resources. New workers are respawn upon submission
80 of new tasks so that ``max_workers`` are available to accept the newly
81 submitted tasks. Setting ``timeout`` to around 100 times the time required
82 to spawn new processes and import packages in them (on the order of 100ms)
83 ensures that the overhead of spawning workers is negligible.
85 Setting ``kill_workers=True`` makes it possible to forcibly interrupt
86 previously spawned jobs to get a new instance of the reusable executor
87 with new constructor argument values.
89 The ``job_reducers`` and ``result_reducers`` are used to customize the
90 pickling of tasks and results send to the executor.
92 When provided, the ``initializer`` is run first in newly spawned
93 processes with argument ``initargs``.
95 The environment variable in the child process are a copy of the values in
96 the main process. One can provide a dict ``{ENV: VAL}`` where ``ENV`` and
97 ``VAL`` are string literals to overwrite the environment variable ``ENV``
98 in the child processes to value ``VAL``. The environment variables are set
99 in the children before any module is loaded. This only works with the
100 ``loky`` context.
101 """
102 _executor, _ = _ReusablePoolExecutor.get_reusable_executor(
103 max_workers=max_workers,
104 context=context,
105 timeout=timeout,
106 kill_workers=kill_workers,
107 reuse=reuse,
108 job_reducers=job_reducers,
109 result_reducers=result_reducers,
110 initializer=initializer,
111 initargs=initargs,
112 env=env,
113 )
114 return _executor
117class _ReusablePoolExecutor(ProcessPoolExecutor):
118 def __init__(
119 self,
120 max_workers=None,
121 context=None,
122 timeout=None,
123 executor_id=0,
124 job_reducers=None,
125 result_reducers=None,
126 initializer=None,
127 initargs=(),
128 env=None,
129 ):
130 super().__init__(
131 max_workers=max_workers,
132 context=context,
133 timeout=timeout,
134 job_reducers=job_reducers,
135 result_reducers=result_reducers,
136 initializer=initializer,
137 initargs=initargs,
138 env=env,
139 )
140 self.executor_id = executor_id
142 @classmethod
143 def get_reusable_executor(
144 cls,
145 max_workers=None,
146 context=None,
147 timeout=10,
148 kill_workers=False,
149 reuse="auto",
150 job_reducers=None,
151 result_reducers=None,
152 initializer=None,
153 initargs=(),
154 env=None,
155 ):
156 executor = _executor_storage.executor
158 if max_workers is None:
159 if reuse is True and executor is not None:
160 max_workers = executor._max_workers
161 else:
162 max_workers = cpu_count()
163 elif max_workers <= 0:
164 raise ValueError(
165 f"max_workers must be greater than 0, got {max_workers}."
166 )
168 if isinstance(context, str):
169 context = get_context(context)
170 if context is not None and context.get_start_method() == "fork":
171 raise ValueError(
172 "Cannot use reusable executor with the 'fork' context"
173 )
175 kwargs = dict(
176 context=context,
177 timeout=timeout,
178 job_reducers=job_reducers,
179 result_reducers=result_reducers,
180 initializer=initializer,
181 initargs=initargs,
182 env=env,
183 )
184 if executor is None:
185 is_reused = False
186 mp.util.debug(f"Create a executor with max_workers={max_workers}.")
187 executor_id = _get_next_executor_id()
188 _executor_storage.executor_kwargs = kwargs
189 _executor_storage.executor = executor = cls(
190 max_workers=max_workers,
191 executor_id=executor_id,
192 **kwargs,
193 )
194 else:
195 if reuse == "auto":
196 reuse = kwargs == _executor_storage.executor_kwargs
197 if (
198 executor._flags.broken
199 or executor._flags.shutdown
200 or not reuse
201 or executor.queue_size < max_workers
202 ):
203 if executor._flags.broken:
204 reason = "broken"
205 elif executor._flags.shutdown:
206 reason = "shutdown"
207 elif executor.queue_size < max_workers:
208 # Do not reuse the executor if the queue size is too
209 # small as this would lead to limited parallelism.
210 reason = "queue size is too small"
211 else:
212 reason = "arguments have changed"
213 mp.util.debug(
214 "Creating a new executor with max_workers="
215 f"{max_workers} as the previous instance cannot be "
216 f"reused ({reason})."
217 )
218 executor.shutdown(wait=True, kill_workers=kill_workers)
219 _executor_storage.executor = executor = (
220 _executor_storage.executor_kwargs
221 ) = None
222 return cls.get_reusable_executor(
223 max_workers=max_workers, **kwargs
224 )
225 else:
226 mp.util.debug(
227 "Reusing existing executor with "
228 f"max_workers={executor._max_workers}."
229 )
230 is_reused = True
231 executor._resize(max_workers)
233 return executor, is_reused
235 def _resize(self, max_workers):
236 if max_workers is None:
237 raise ValueError("Trying to resize with max_workers=None")
238 elif max_workers == self._max_workers:
239 return
241 if self._executor_manager_thread is None:
242 # If the executor_manager_thread has not been started
243 # then no processes have been spawned and we can just
244 # update _max_workers and return
245 self._max_workers = max_workers
246 return
248 self._wait_job_completion()
250 # Some process might have returned due to timeout so check how many
251 # children are still alive. Use the _process_management_lock to
252 # ensure that no process are spawned or timeout during the resize.
253 with self._processes_management_lock:
254 processes = list(self._processes.values())
255 nb_children_alive = sum(p.is_alive() for p in processes)
256 self._max_workers = max_workers
257 for _ in range(max_workers, nb_children_alive):
258 self._call_queue.put(None)
259 while len(self._processes) > max_workers and not self._flags.broken:
260 time.sleep(1e-3)
262 self._adjust_process_count()
263 processes = list(self._processes.values())
264 while not all(p.is_alive() for p in processes):
265 time.sleep(1e-3)
267 def _wait_job_completion(self):
268 """Wait for the cache to be empty before resizing the pool."""
269 # Issue a warning to the user about the bad effect of this usage.
270 if self._pending_work_items:
271 warnings.warn(
272 "Trying to resize an executor with running jobs: "
273 "waiting for jobs completion before resizing.",
274 UserWarning,
275 )
276 mp.util.debug(
277 f"Executor {self.executor_id} waiting for jobs completion "
278 "before resizing"
279 )
280 # Wait for the completion of the jobs
281 while self._pending_work_items:
282 time.sleep(1e-3)
284 def _setup_queues(self, job_reducers, result_reducers):
285 # As this executor can be resized, use a large queue size to avoid
286 # underestimating capacity and introducing overhead
287 # Also handle the case where the user set max_workers to a value larger
288 # than cpu_count(), to avoid limiting the number of parallel jobs.
290 min_queue_size = max(cpu_count(), self._max_workers)
291 self.queue_size = 2 * min_queue_size + EXTRA_QUEUED_CALLS
292 super()._setup_queues(
293 job_reducers, result_reducers, queue_size=self.queue_size
294 )