1"""
2Backends for embarrassingly parallel code.
3"""
4
5from __future__ import annotations
6
7import contextlib
8import gc
9import os
10import threading
11import warnings
12from abc import ABCMeta, abstractmethod
13from dataclasses import dataclass
14from typing import Any, Callable
15
16from ._multiprocessing_helpers import mp
17from ._utils import (
18 _retrieve_traceback_capturing_wrapped_call,
19 _TracebackCapturingWrapper,
20)
21
22if mp is None:
23 from os import cpu_count as _os_cpu_count
24
25 def system_cpu_count(only_physical_cores: bool = False) -> int:
26 return _os_cpu_count() or 1
27else:
28 from multiprocessing.pool import ThreadPool
29
30 from .executor import get_memmapping_executor
31
32 # Import loky only if multiprocessing is present
33 from .externals.loky import cpu_count as system_cpu_count
34 from .externals.loky import process_executor
35 from .externals.loky.process_executor import ShutdownExecutorError
36 from .pool import MemmappingPool
37
38
39class _MaxCores(threading.local):
40 """Track the number of cores available to this thread."""
41
42 _thread_limit = None
43
44 def get(self, only_physical_cores=False) -> int:
45 """Number of cores available to this thread."""
46 process_limit = system_cpu_count(only_physical_cores)
47 if self._thread_limit is None:
48 return process_limit
49 return min(process_limit, self._thread_limit)
50
51 def set_thread_limit(self, cores: int) -> None:
52 """Set the maximum number of cores available to this thread."""
53 self._thread_limit = cores
54
55
56_MAX_CORES = _MaxCores()
57
58
59def cpu_count(only_physical_cores=False, process_wide=False) -> int:
60 """Return the number of CPU cores the current thread can use.
61
62 Per-thread limits can be constrained by ``joblib.Parallel``'s threaded
63 backend: cores will be split up across the worker threads. Additionally,
64 process-wide limits can be constrained by CPU affinity, and on Linux
65 cgroups (i.e. Docker/Kubernetes/other container systems).
66
67 If ``process_wide`` is True, ignore the per-thread limits.
68
69 If both ``process_wide`` and ``only_physical_cores`` are True, return the
70 number of physical cores instead of the number of logical cores
71 (hyperthreading / SMT). If ``only_physical_cores`` is True and
72 ``process_wide`` is False (its default) the minimum of the two (per-thread
73 and physical limit) will be chosen, but future versions may have a smarter
74 algorithm.
75 """
76 if process_wide:
77 return system_cpu_count(only_physical_cores=only_physical_cores)
78 return _MAX_CORES.get(only_physical_cores)
79
80
81def set_thread_local_cpu_limit(num_cores: int) -> None:
82 """Set the thread-local limit for result of ``cpu_count()``.
83
84 Outside of internal use, this is useful for external thread pool
85 implementations that want to rely on ``joblib``'s ``cpu_count()``. They
86 will want to call this in each worker's initializer function, dividing up
87 the cores available in the parent threads.
88 """
89 _MAX_CORES.set_thread_limit(num_cores)
90
91
92def _split_up_cores(total_cores: int, n_jobs: int) -> int:
93 """Divide up cores between workers.
94
95 Given the total number of cores and a number of workers, come up with a
96 reasonable number of cores per worker.
97
98 At the moment this just does ``total_cores // n_jobs`` but a better
99 heuristic might someday be used instead.
100 """
101 return max(total_cores // n_jobs, 1)
102
103
104class ParallelBackendBase(metaclass=ABCMeta):
105 """Helper abc which defines all methods a ParallelBackend must implement"""
106
107 default_n_jobs = 1
108
109 supports_inner_max_num_threads = False
110
111 # This flag was introduced for backward compatibility reasons.
112 # New backends should always set it to True and implement the
113 # `retrieve_result_callback` method.
114 supports_retrieve_callback = False
115
116 @property
117 def supports_return_generator(self):
118 return self.supports_retrieve_callback
119
120 @property
121 def supports_timeout(self):
122 return self.supports_retrieve_callback
123
124 nesting_level = None
125
126 def __init__(
127 self, nesting_level=None, inner_max_num_threads=None, **backend_kwargs
128 ):
129 super().__init__()
130 self.nesting_level = nesting_level
131 self.inner_max_num_threads = inner_max_num_threads
132 self.backend_kwargs = backend_kwargs
133
134 MAX_NUM_THREADS_VARS = [
135 "OMP_NUM_THREADS",
136 "OPENBLAS_NUM_THREADS",
137 "MKL_NUM_THREADS",
138 "BLIS_NUM_THREADS",
139 "VECLIB_MAXIMUM_THREADS",
140 "NUMBA_NUM_THREADS",
141 "NUMEXPR_NUM_THREADS",
142 # This sets a soft max on loky.cpu_count() in workers:
143 "LOKY_MAX_CPU_COUNT",
144 ]
145
146 TBB_ENABLE_IPC_VAR = "ENABLE_IPC"
147
148 @abstractmethod
149 def effective_n_jobs(self, n_jobs):
150 """Determine the number of jobs that can actually run in parallel
151
152 n_jobs is the number of workers requested by the callers. Passing
153 n_jobs=-1 means requesting all available workers for instance matching
154 the number of CPU cores on the worker host(s).
155
156 This method should return a guesstimate of the number of workers that
157 can actually perform work concurrently. The primary use case is to make
158 it possible for the caller to know in how many chunks to slice the
159 work.
160
161 In general working on larger data chunks is more efficient (less
162 scheduling overhead and better use of CPU cache prefetching heuristics)
163 as long as all the workers have enough work to do.
164 """
165
166 def apply_async(self, func, callback=None):
167 """Deprecated: implement `submit` instead."""
168 raise NotImplementedError("Implement `submit` instead.")
169
170 def submit(self, func, callback=None):
171 """Schedule a function to be run and return a future-like object.
172
173 This method should return a future-like object that allow tracking
174 the progress of the task.
175
176 If ``supports_retrieve_callback`` is False, the return value of this
177 method is passed to ``retrieve_result`` instead of calling
178 ``retrieve_result_callback``.
179
180 Parameters
181 ----------
182 func: callable
183 The function to be run in parallel.
184
185 callback: callable
186 A callable that will be called when the task is completed. This callable
187 is a wrapper around ``retrieve_result_callback``. This should be added
188 to the future-like object returned by this method, so that the callback
189 is called when the task is completed.
190
191 For future-like backends, this can be achieved with something like
192 ``future.add_done_callback(callback)``.
193
194 Returns
195 -------
196 future: future-like
197 A future-like object to track the execution of the submitted function.
198 """
199 warnings.warn(
200 "`apply_async` is deprecated, implement and use `submit` instead.",
201 DeprecationWarning,
202 )
203 return self.apply_async(func, callback)
204
205 def retrieve_result_callback(self, out):
206 """Called within the callback function passed to `submit`.
207
208 This method can customise how the result of the function is retrieved
209 from the future-like object.
210
211 Parameters
212 ----------
213 future: future-like
214 The future-like object returned by the `submit` method.
215
216 Returns
217 -------
218 result: object
219 The result of the function executed in parallel.
220 """
221
222 def retrieve_result(self, out, timeout=None):
223 """Hook to retrieve the result when support_retrieve_callback=False.
224
225 The argument `out` is the result of the `submit` call. This method
226 should return the result of the computation or raise an exception if
227 the computation failed.
228 """
229 if self.supports_timeout:
230 return out.get(timeout=timeout)
231 else:
232 return out.get()
233
234 def configure(
235 self, n_jobs=1, parallel=None, prefer=None, require=None, **backend_kwargs
236 ):
237 """Reconfigure the backend and return the number of workers.
238
239 This makes it possible to reuse an existing backend instance for
240 successive independent calls to Parallel with different parameters.
241 """
242 self.parallel = parallel
243 return self.effective_n_jobs(n_jobs)
244
245 def start_call(self):
246 """Call-back method called at the beginning of a Parallel call"""
247
248 def stop_call(self):
249 """Call-back method called at the end of a Parallel call"""
250
251 def terminate(self):
252 """Shutdown the workers and free the shared memory."""
253
254 def compute_batch_size(self):
255 """Determine the optimal batch size"""
256 return 1
257
258 def batch_completed(self, batch_size, duration):
259 """Callback indicate how long it took to run a batch"""
260
261 def abort_everything(self, ensure_ready=True):
262 """Abort any running tasks
263
264 This is called when an exception has been raised when executing a task
265 and all the remaining tasks will be ignored and can therefore be
266 aborted to spare computation resources.
267
268 If ensure_ready is True, the backend should be left in an operating
269 state as future tasks might be re-submitted via that same backend
270 instance.
271
272 If ensure_ready is False, the implementer of this method can decide
273 to leave the backend in a closed / terminated state as no new task
274 are expected to be submitted to this backend.
275
276 Setting ensure_ready to False is an optimization that can be leveraged
277 when aborting tasks via killing processes from a local process pool
278 managed by the backend it-self: if we expect no new tasks, there is no
279 point in re-creating new workers.
280 """
281 # Does nothing by default: to be overridden in subclasses when
282 # canceling tasks is possible.
283 pass
284
285 def get_nested_backend(self):
286 """Backend instance to be used by nested Parallel calls.
287
288 By default a thread-based backend is used for the first level of
289 nesting. Beyond, switch to sequential backend to avoid spawning too
290 many threads on the host.
291 """
292 nesting_level = getattr(self, "nesting_level", 0) + 1
293 if nesting_level > 1:
294 return SequentialBackend(nesting_level=nesting_level), None
295 else:
296 return ThreadingBackend(nesting_level=nesting_level), None
297
298 def _prepare_worker_env(self, n_jobs):
299 """Return environment variables limiting threadpools in external libs.
300
301 This function return a dict containing environment variables to pass
302 when creating a pool of process. These environment variables limit the
303 number of threads to `n_threads` for OpenMP, MKL, Accelerated and
304 OpenBLAS libraries in the child processes.
305 """
306 explicit_n_threads = self.inner_max_num_threads
307 default_n_threads = _split_up_cores(cpu_count(), n_jobs)
308
309 # Set the inner environment variables to self.inner_max_num_threads if
310 # it is given. Else, default to cpu_count // n_jobs unless the variable
311 # is already present in the parent process environment.
312 env = {}
313 for var in self.MAX_NUM_THREADS_VARS:
314 if explicit_n_threads is None:
315 if var == "LOKY_MAX_CPU_COUNT":
316 # We don't want to pass it through, it was intended for
317 # this process, and we've potentially split it up due to
318 # worker threads.
319 var_value = default_n_threads
320 else:
321 var_value = os.environ.get(var, default_n_threads)
322 else:
323 var_value = explicit_n_threads
324
325 env[var] = str(var_value)
326
327 if self.TBB_ENABLE_IPC_VAR not in os.environ:
328 # To avoid over-subscription when using TBB, let the TBB schedulers
329 # use Inter Process Communication to coordinate:
330 env[self.TBB_ENABLE_IPC_VAR] = "1"
331 return env
332
333 @contextlib.contextmanager
334 def retrieval_context(self):
335 """Context manager to manage an execution context.
336
337 Calls to Parallel.retrieve will be made inside this context.
338
339 By default, this does nothing. It may be useful for subclasses to
340 handle nested parallelism. In particular, it may be required to avoid
341 deadlocks if a backend manages a fixed number of workers, when those
342 workers may be asked to do nested Parallel calls. Without
343 'retrieval_context' this could lead to deadlock, as all the workers
344 managed by the backend may be "busy" waiting for the nested parallel
345 calls to finish, but the backend has no free workers to execute those
346 tasks.
347 """
348 yield
349
350 @staticmethod
351 def in_main_thread():
352 return isinstance(threading.current_thread(), threading._MainThread)
353
354
355class SequentialBackend(ParallelBackendBase):
356 """A ParallelBackend which will execute all batches sequentially.
357
358 Does not use/create any threading objects, and hence has minimal
359 overhead. Used when n_jobs == 1.
360 """
361
362 uses_threads = True
363 supports_timeout = False
364 supports_retrieve_callback = False
365 supports_sharedmem = True
366
367 def effective_n_jobs(self, n_jobs):
368 """Determine the number of jobs which are going to run in parallel"""
369 if n_jobs == 0:
370 raise ValueError("n_jobs == 0 in Parallel has no meaning")
371 return 1
372
373 def submit(self, func, callback=None):
374 """Schedule a func to be run"""
375 raise RuntimeError("Should never be called for SequentialBackend.")
376
377 def retrieve_result_callback(self, out):
378 raise RuntimeError("Should never be called for SequentialBackend.")
379
380 def get_nested_backend(self):
381 # import is not top level to avoid cyclic import errors.
382 from .parallel import get_active_backend
383
384 # SequentialBackend should neither change the nesting level, the
385 # default backend or the number of jobs. Just return the current one.
386 return get_active_backend()
387
388
389class PoolManagerMixin(object):
390 """A helper class for managing pool of workers."""
391
392 _pool = None
393
394 def effective_n_jobs(self, n_jobs):
395 """Determine the number of jobs which are going to run in parallel"""
396 if n_jobs == 0:
397 raise ValueError("n_jobs == 0 in Parallel has no meaning")
398 elif mp is None or n_jobs is None:
399 # multiprocessing is not available or disabled, fallback
400 # to sequential mode
401 return 1
402 elif n_jobs < 0:
403 n_jobs = max(cpu_count() + 1 + n_jobs, 1)
404 return n_jobs
405
406 def terminate(self):
407 """Shutdown the process or thread pool"""
408 if self._pool is not None:
409 self._pool.close()
410 self._pool.terminate() # terminate does a join()
411 self._pool = None
412
413 def _get_pool(self):
414 """Used by `submit` to make it possible to implement lazy init"""
415 return self._pool
416
417 def submit(self, func, callback=None):
418 """Schedule a func to be run"""
419 # Here, we need a wrapper to avoid crashes on KeyboardInterruptErrors.
420 # We also call the callback on error, to make sure the pool does not
421 # wait on crashed jobs.
422 return self._get_pool().apply_async(
423 _TracebackCapturingWrapper(func),
424 (),
425 callback=callback,
426 error_callback=callback,
427 )
428
429 def retrieve_result_callback(self, result):
430 """Mimic concurrent.futures results, raising an error if needed."""
431 # In the multiprocessing Pool API, the callback are called with the
432 # result value as an argument so `result`(`out`) is the output of
433 # job.get(). It's either the result or the exception raised while
434 # collecting the result.
435 return _retrieve_traceback_capturing_wrapped_call(result)
436
437 def abort_everything(self, ensure_ready=True):
438 """Shutdown the pool and restart a new one with the same parameters"""
439 self.terminate()
440 if ensure_ready:
441 self.configure(
442 n_jobs=self.parallel.n_jobs,
443 parallel=self.parallel,
444 **self.parallel._backend_kwargs,
445 )
446
447
448class AutoBatchingMixin(object):
449 """A helper class for automagically batching jobs."""
450
451 # In seconds, should be big enough to hide multiprocessing dispatching
452 # overhead.
453 # This settings was found by running benchmarks/bench_auto_batching.py
454 # with various parameters on various platforms.
455 MIN_IDEAL_BATCH_DURATION = 0.2
456
457 # Should not be too high to avoid stragglers: long jobs running alone
458 # on a single worker while other workers have no work to process any more.
459 MAX_IDEAL_BATCH_DURATION = 2
460
461 # Batching counters default values
462 _DEFAULT_EFFECTIVE_BATCH_SIZE = 1
463 _DEFAULT_SMOOTHED_BATCH_DURATION = 0.0
464
465 def __init__(self, **kwargs):
466 super().__init__(**kwargs)
467 self._effective_batch_size = self._DEFAULT_EFFECTIVE_BATCH_SIZE
468 self._smoothed_batch_duration = self._DEFAULT_SMOOTHED_BATCH_DURATION
469
470 def compute_batch_size(self):
471 """Determine the optimal batch size"""
472 old_batch_size = self._effective_batch_size
473 batch_duration = self._smoothed_batch_duration
474 if batch_duration > 0 and batch_duration < self.MIN_IDEAL_BATCH_DURATION:
475 # The current batch size is too small: the duration of the
476 # processing of a batch of task is not large enough to hide
477 # the scheduling overhead.
478 ideal_batch_size = int(
479 old_batch_size * self.MIN_IDEAL_BATCH_DURATION / batch_duration
480 )
481 # Multiply by two to limit oscilations between min and max.
482 ideal_batch_size *= 2
483
484 # dont increase the batch size too fast to limit huge batch sizes
485 # potentially leading to starving worker
486 batch_size = min(2 * old_batch_size, ideal_batch_size)
487
488 batch_size = max(batch_size, 1)
489
490 self._effective_batch_size = batch_size
491 if self.parallel.verbose >= 10:
492 self.parallel._print(
493 f"Batch computation too fast ({batch_duration}s.) "
494 f"Setting batch_size={batch_size}."
495 )
496 elif batch_duration > self.MAX_IDEAL_BATCH_DURATION and old_batch_size >= 2:
497 # The current batch size is too big. If we schedule overly long
498 # running batches some CPUs might wait with nothing left to do
499 # while a couple of CPUs a left processing a few long running
500 # batches. Better reduce the batch size a bit to limit the
501 # likelihood of scheduling such stragglers.
502
503 # decrease the batch size quickly to limit potential starving
504 ideal_batch_size = int(
505 old_batch_size * self.MIN_IDEAL_BATCH_DURATION / batch_duration
506 )
507 # Multiply by two to limit oscilations between min and max.
508 batch_size = max(2 * ideal_batch_size, 1)
509 self._effective_batch_size = batch_size
510 if self.parallel.verbose >= 10:
511 self.parallel._print(
512 f"Batch computation too slow ({batch_duration}s.) "
513 f"Setting batch_size={batch_size}."
514 )
515 else:
516 # No batch size adjustment
517 batch_size = old_batch_size
518
519 if batch_size != old_batch_size:
520 # Reset estimation of the smoothed mean batch duration: this
521 # estimate is updated in the multiprocessing apply_async
522 # CallBack as long as the batch_size is constant. Therefore
523 # we need to reset the estimate whenever we re-tune the batch
524 # size.
525 self._smoothed_batch_duration = self._DEFAULT_SMOOTHED_BATCH_DURATION
526
527 return batch_size
528
529 def batch_completed(self, batch_size, duration):
530 """Callback indicate how long it took to run a batch"""
531 if batch_size == self._effective_batch_size:
532 # Update the smoothed streaming estimate of the duration of a batch
533 # from dispatch to completion
534 old_duration = self._smoothed_batch_duration
535 if old_duration == self._DEFAULT_SMOOTHED_BATCH_DURATION:
536 # First record of duration for this batch size after the last
537 # reset.
538 new_duration = duration
539 else:
540 # Update the exponentially weighted average of the duration of
541 # batch for the current effective size.
542 new_duration = 0.8 * old_duration + 0.2 * duration
543 self._smoothed_batch_duration = new_duration
544
545 def reset_batch_stats(self):
546 """Reset batch statistics to default values.
547
548 This avoids interferences with future jobs.
549 """
550 self._effective_batch_size = self._DEFAULT_EFFECTIVE_BATCH_SIZE
551 self._smoothed_batch_duration = self._DEFAULT_SMOOTHED_BATCH_DURATION
552
553
554class ThreadingBackend(PoolManagerMixin, ParallelBackendBase):
555 """A ParallelBackend which will use a thread pool to execute batches in.
556
557 This is a low-overhead backend but it suffers from the Python Global
558 Interpreter Lock if the called function relies a lot on Python objects.
559 Mostly useful when the execution bottleneck is a compiled extension that
560 explicitly releases the GIL (for instance a Cython loop wrapped in a "with
561 nogil" block or an expensive call to a library such as NumPy).
562
563 The actual thread pool is lazily initialized: the actual thread pool
564 construction is delayed to the first call to apply_async.
565
566 ThreadingBackend is used as the default backend for nested calls.
567 """
568
569 supports_retrieve_callback = True
570 uses_threads = True
571 supports_sharedmem = True
572
573 def configure(self, n_jobs=1, parallel=None, **backend_kwargs):
574 """Build a process or thread pool and return the number of workers"""
575 n_jobs = self.effective_n_jobs(n_jobs)
576 if n_jobs == 1:
577 # Avoid unnecessary overhead and use sequential backend instead.
578 raise FallbackToBackend(SequentialBackend(nesting_level=self.nesting_level))
579 self.parallel = parallel
580 self._n_jobs = n_jobs
581 return n_jobs
582
583 def _get_pool(self):
584 """Lazily initialize the thread pool
585
586 The actual pool of worker threads is only initialized at the first
587 call to apply_async.
588 """
589 # Import here to prevent circular import:
590 from joblib.parallel import effective_n_jobs
591
592 if self._pool is None:
593 available_cores = effective_n_jobs(-1)
594 cores_per_thread = _split_up_cores(available_cores, self._n_jobs)
595 self._pool = ThreadPool(
596 self._n_jobs,
597 initializer=lambda: set_thread_local_cpu_limit(cores_per_thread),
598 )
599 return self._pool
600
601
602@dataclass
603class _SetEnvInitializer:
604 """
605 Pickleable initializer for multiprocessing workers.
606 """
607
608 env: dict[str, str]
609 initializer: None | Callable[..., Any]
610
611 def __call__(self, *args, **kwargs) -> Any:
612 for key, value in self.env.items():
613 os.environ[key] = value
614 if self.initializer is not None:
615 return self.initializer(*args, **kwargs)
616
617
618class MultiprocessingBackend(PoolManagerMixin, AutoBatchingMixin, ParallelBackendBase):
619 """A ParallelBackend which will use a multiprocessing.Pool.
620
621 Will introduce some communication and memory overhead when exchanging
622 input and output data with the with the worker Python processes.
623 However, does not suffer from the Python Global Interpreter Lock.
624 """
625
626 supports_retrieve_callback = True
627 supports_return_generator = False
628
629 def effective_n_jobs(self, n_jobs):
630 """Determine the number of jobs which are going to run in parallel.
631
632 This also checks if we are attempting to create a nested parallel
633 loop.
634 """
635 if n_jobs == 0:
636 raise ValueError("n_jobs == 0 in Parallel has no meaning")
637
638 if mp is None:
639 return 1
640
641 if mp.current_process().daemon:
642 # Daemonic processes cannot have children
643 if n_jobs != 1:
644 if inside_dask_worker():
645 msg = (
646 "Inside a Dask worker with daemon=True, "
647 "setting n_jobs=1.\nPossible work-arounds:\n"
648 "- dask.config.set("
649 "{'distributed.worker.daemon': False})"
650 "- set the environment variable "
651 "DASK_DISTRIBUTED__WORKER__DAEMON=False\n"
652 "before creating your Dask cluster."
653 )
654 else:
655 msg = (
656 "Multiprocessing-backed parallel loops "
657 "cannot be nested, setting n_jobs=1"
658 )
659 warnings.warn(msg, stacklevel=3)
660 return 1
661
662 if process_executor._CURRENT_DEPTH > 0:
663 # Mixing loky and multiprocessing in nested loop is not supported
664 if n_jobs != 1:
665 warnings.warn(
666 "Multiprocessing-backed parallel loops cannot be nested,"
667 " below loky, setting n_jobs=1",
668 stacklevel=3,
669 )
670 return 1
671
672 elif not (self.in_main_thread() or self.nesting_level == 0):
673 # Prevent posix fork inside in non-main posix threads
674 if n_jobs != 1:
675 warnings.warn(
676 "Multiprocessing-backed parallel loops cannot be nested"
677 " below threads, setting n_jobs=1",
678 stacklevel=3,
679 )
680 return 1
681
682 return super(MultiprocessingBackend, self).effective_n_jobs(n_jobs)
683
684 def configure(
685 self,
686 n_jobs=1,
687 parallel=None,
688 prefer=None,
689 require=None,
690 **memmapping_pool_kwargs,
691 ):
692 """Build a process or thread pool and return the number of workers"""
693 n_jobs = self.effective_n_jobs(n_jobs)
694 if n_jobs == 1:
695 raise FallbackToBackend(SequentialBackend(nesting_level=self.nesting_level))
696
697 memmapping_pool_kwargs = {
698 **self.backend_kwargs,
699 **memmapping_pool_kwargs,
700 }
701
702 # Make sure to free as much memory as possible before forking
703 gc.collect()
704 initializer = _SetEnvInitializer(
705 self._prepare_worker_env(n_jobs),
706 memmapping_pool_kwargs.pop("initializer", None),
707 )
708 self._pool = MemmappingPool(
709 n_jobs,
710 initializer=initializer,
711 **memmapping_pool_kwargs,
712 )
713 self.parallel = parallel
714 return n_jobs
715
716 def terminate(self):
717 """Shutdown the process or thread pool"""
718 super(MultiprocessingBackend, self).terminate()
719 self.reset_batch_stats()
720
721
722class LokyBackend(AutoBatchingMixin, ParallelBackendBase):
723 """Managing pool of workers with loky instead of multiprocessing."""
724
725 supports_retrieve_callback = True
726 supports_inner_max_num_threads = True
727
728 def configure(
729 self,
730 n_jobs=1,
731 parallel=None,
732 prefer=None,
733 require=None,
734 idle_worker_timeout=None,
735 **memmapping_executor_kwargs,
736 ):
737 """Build a process executor and return the number of workers"""
738 n_jobs = self.effective_n_jobs(n_jobs)
739 if n_jobs == 1:
740 raise FallbackToBackend(SequentialBackend(nesting_level=self.nesting_level))
741
742 memmapping_executor_kwargs = {
743 **self.backend_kwargs,
744 **memmapping_executor_kwargs,
745 }
746
747 # Prohibit the use of 'timeout' in the LokyBackend, as 'idle_worker_timeout'
748 # better describes the backend's behavior.
749 if "timeout" in memmapping_executor_kwargs:
750 raise ValueError(
751 "The 'timeout' parameter is not supported by the LokyBackend. "
752 "Please use the `idle_worker_timeout` parameter instead."
753 )
754 if idle_worker_timeout is None:
755 idle_worker_timeout = self.backend_kwargs.get("idle_worker_timeout", 300)
756
757 self._workers = get_memmapping_executor(
758 n_jobs,
759 timeout=idle_worker_timeout,
760 env=self._prepare_worker_env(n_jobs=n_jobs),
761 context_id=parallel._id,
762 **memmapping_executor_kwargs,
763 )
764 self.parallel = parallel
765 return n_jobs
766
767 def effective_n_jobs(self, n_jobs):
768 """Determine the number of jobs which are going to run in parallel"""
769 if n_jobs == 0:
770 raise ValueError("n_jobs == 0 in Parallel has no meaning")
771 elif mp is None or n_jobs is None:
772 # multiprocessing is not available or disabled, fallback
773 # to sequential mode
774 return 1
775 elif mp.current_process().daemon:
776 # Daemonic processes cannot have children
777 if n_jobs != 1:
778 if inside_dask_worker():
779 msg = (
780 "Inside a Dask worker with daemon=True, "
781 "setting n_jobs=1.\nPossible work-arounds:\n"
782 "- dask.config.set("
783 "{'distributed.worker.daemon': False})\n"
784 "- set the environment variable "
785 "DASK_DISTRIBUTED__WORKER__DAEMON=False\n"
786 "before creating your Dask cluster."
787 )
788 else:
789 msg = (
790 "Loky-backed parallel loops cannot be called in a"
791 " multiprocessing, setting n_jobs=1"
792 )
793 warnings.warn(msg, stacklevel=3)
794
795 return 1
796 elif not (self.in_main_thread() or self.nesting_level == 0):
797 # Prevent posix fork inside in non-main posix threads
798 if n_jobs != 1:
799 warnings.warn(
800 "Loky-backed parallel loops cannot be nested below "
801 "threads, setting n_jobs=1",
802 stacklevel=3,
803 )
804 return 1
805 elif n_jobs < 0:
806 n_jobs = max(cpu_count() + 1 + n_jobs, 1)
807 return n_jobs
808
809 def submit(self, func, callback=None):
810 """Schedule a func to be run"""
811 future = self._workers.submit(func)
812 if callback is not None:
813 future.add_done_callback(callback)
814 return future
815
816 def retrieve_result_callback(self, future):
817 """Retrieve the result, here out is the future given by submit"""
818 try:
819 return future.result()
820 except ShutdownExecutorError:
821 raise RuntimeError(
822 "The executor underlying Parallel has been shutdown. "
823 "This is likely due to the garbage collection of a previous "
824 "generator from a call to Parallel with return_as='generator'."
825 " Make sure the generator is not garbage collected when "
826 "submitting a new job or that it is first properly exhausted."
827 )
828
829 def terminate(self):
830 if self._workers is not None:
831 # Don't terminate the workers as we want to reuse them in later
832 # calls, but cleanup the temporary resources that the Parallel call
833 # created. This 'hack' requires a private, low-level operation.
834 self._workers._temp_folder_manager._clean_temporary_resources(
835 context_id=self.parallel._id, force=False
836 )
837 self._workers = None
838
839 self.reset_batch_stats()
840
841 def abort_everything(self, ensure_ready=True):
842 """Shutdown the workers and restart a new one with the same parameters"""
843 self._workers.terminate(kill_workers=True)
844 self._workers = None
845
846 if ensure_ready:
847 self.configure(n_jobs=self.parallel.n_jobs, parallel=self.parallel)
848
849
850class FallbackToBackend(Exception):
851 """Raised when configuration should fallback to another backend"""
852
853 def __init__(self, backend):
854 self.backend = backend
855
856
857def inside_dask_worker():
858 """Check whether the current function is executed inside a Dask worker."""
859 # This function can not be in joblib._dask because there would be a
860 # circular import:
861 # _dask imports _parallel_backend that imports _dask ...
862 try:
863 from distributed import get_worker
864 except ImportError:
865 return False
866
867 try:
868 get_worker()
869 return True
870 except ValueError:
871 return False