blob: b8a0b827635f07ae53d33579e674cb171eaba716 [file] [log] [blame]
Benjamin Petersone711caf2008-06-11 16:44:04 +00001#
2# Module providing the `Pool` class for managing a process pool
3#
4# multiprocessing/pool.py
5#
R. David Murray3fc969a2010-12-14 01:38:16 +00006# Copyright (c) 2006-2008, R Oudkerk
Richard Oudkerk3e268aa2012-04-30 12:13:55 +01007# Licensed to PSF under a Contributor Agreement.
Benjamin Petersone711caf2008-06-11 16:44:04 +00008#
9
Richard Oudkerk84ed9a62013-08-14 15:35:41 +010010__all__ = ['Pool', 'ThreadPool']
Benjamin Petersone711caf2008-06-11 16:44:04 +000011
12#
13# Imports
14#
15
Benjamin Petersone711caf2008-06-11 16:44:04 +000016import collections
Victor Stinner9a8d1d72018-12-20 20:33:51 +010017import itertools
Charles-François Natali37cfb0a2013-06-28 19:25:45 +020018import os
Victor Stinner9a8d1d72018-12-20 20:33:51 +010019import queue
20import threading
Benjamin Petersone711caf2008-06-11 16:44:04 +000021import time
Richard Oudkerk85757832013-05-06 11:38:25 +010022import traceback
Batuhan Taşkaya03615562020-04-10 17:46:36 +030023import types
Victor Stinner9a8d1d72018-12-20 20:33:51 +010024import warnings
Pablo Galindo7c994542019-03-16 22:34:24 +000025from queue import Empty
Benjamin Petersone711caf2008-06-11 16:44:04 +000026
Richard Oudkerk84ed9a62013-08-14 15:35:41 +010027# If threading is available then ThreadPool should be provided. Therefore
28# we avoid top-level imports which are liable to fail on some systems.
29from . import util
Victor Stinner7fa767e2014-03-20 09:16:38 +010030from . import get_context, TimeoutError
Pablo Galindo7c994542019-03-16 22:34:24 +000031from .connection import wait
Benjamin Petersone711caf2008-06-11 16:44:04 +000032
33#
34# Constants representing the state of a pool
35#
36
Victor Stinner9a8d1d72018-12-20 20:33:51 +010037INIT = "INIT"
Victor Stinner2b417fb2018-12-14 11:13:18 +010038RUN = "RUN"
39CLOSE = "CLOSE"
40TERMINATE = "TERMINATE"
Benjamin Petersone711caf2008-06-11 16:44:04 +000041
42#
43# Miscellaneous
44#
45
46job_counter = itertools.count()
47
48def mapstar(args):
49 return list(map(*args))
50
Antoine Pitroude911b22011-12-21 11:03:24 +010051def starmapstar(args):
52 return list(itertools.starmap(args[0], args[1]))
53
Benjamin Petersone711caf2008-06-11 16:44:04 +000054#
Richard Oudkerk85757832013-05-06 11:38:25 +010055# Hack to embed stringification of remote traceback in local traceback
56#
57
58class RemoteTraceback(Exception):
59 def __init__(self, tb):
60 self.tb = tb
61 def __str__(self):
62 return self.tb
63
64class ExceptionWithTraceback:
65 def __init__(self, exc, tb):
66 tb = traceback.format_exception(type(exc), exc, tb)
67 tb = ''.join(tb)
68 self.exc = exc
69 self.tb = '\n"""\n%s"""' % tb
70 def __reduce__(self):
71 return rebuild_exc, (self.exc, self.tb)
72
73def rebuild_exc(exc, tb):
74 exc.__cause__ = RemoteTraceback(tb)
75 return exc
76
77#
Benjamin Petersone711caf2008-06-11 16:44:04 +000078# Code run by worker processes
79#
80
Ask Solem2afcbf22010-11-09 20:55:52 +000081class MaybeEncodingError(Exception):
82 """Wraps possible unpickleable errors, so they can be
83 safely sent through the socket."""
84
85 def __init__(self, exc, value):
86 self.exc = repr(exc)
87 self.value = repr(value)
88 super(MaybeEncodingError, self).__init__(self.exc, self.value)
89
90 def __str__(self):
91 return "Error sending result: '%s'. Reason: '%s'" % (self.value,
92 self.exc)
93
94 def __repr__(self):
Serhiy Storchaka465e60e2014-07-25 23:36:00 +030095 return "<%s: %s>" % (self.__class__.__name__, self)
Ask Solem2afcbf22010-11-09 20:55:52 +000096
97
Richard Oudkerk80a5be12014-03-23 12:30:54 +000098def worker(inqueue, outqueue, initializer=None, initargs=(), maxtasks=None,
99 wrap_exception=False):
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500100 if (maxtasks is not None) and not (isinstance(maxtasks, int)
101 and maxtasks >= 1):
102 raise AssertionError("Maxtasks {!r} is not valid".format(maxtasks))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000103 put = outqueue.put
104 get = inqueue.get
105 if hasattr(inqueue, '_writer'):
106 inqueue._writer.close()
107 outqueue._reader.close()
108
109 if initializer is not None:
110 initializer(*initargs)
111
Jesse Noller1f0b6582010-01-27 03:36:01 +0000112 completed = 0
113 while maxtasks is None or (maxtasks and completed < maxtasks):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000114 try:
115 task = get()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200116 except (EOFError, OSError):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100117 util.debug('worker got EOFError or OSError -- exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000118 break
119
120 if task is None:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100121 util.debug('worker got sentinel -- exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000122 break
123
124 job, i, func, args, kwds = task
125 try:
126 result = (True, func(*args, **kwds))
127 except Exception as e:
Xiang Zhang794623b2017-03-29 11:58:54 +0800128 if wrap_exception and func is not _helper_reraises_exception:
Richard Oudkerk80a5be12014-03-23 12:30:54 +0000129 e = ExceptionWithTraceback(e, e.__traceback__)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000130 result = (False, e)
Ask Solem2afcbf22010-11-09 20:55:52 +0000131 try:
132 put((job, i, result))
133 except Exception as e:
134 wrapped = MaybeEncodingError(e, result[1])
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100135 util.debug("Possible encoding error while sending result: %s" % (
Ask Solem2afcbf22010-11-09 20:55:52 +0000136 wrapped))
137 put((job, i, (False, wrapped)))
Antoine Pitrou89889452017-03-24 13:52:11 +0100138
139 task = job = result = func = args = kwds = None
Jesse Noller1f0b6582010-01-27 03:36:01 +0000140 completed += 1
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100141 util.debug('worker exiting after %d tasks' % completed)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000142
Xiang Zhang794623b2017-03-29 11:58:54 +0800143def _helper_reraises_exception(ex):
144 'Pickle-able helper function for use by _guarded_task_generation.'
145 raise ex
146
Benjamin Petersone711caf2008-06-11 16:44:04 +0000147#
148# Class representing a process pool
149#
150
Pablo Galindo7c994542019-03-16 22:34:24 +0000151class _PoolCache(dict):
152 """
153 Class that implements a cache for the Pool class that will notify
154 the pool management threads every time the cache is emptied. The
155 notification is done by the use of a queue that is provided when
156 instantiating the cache.
157 """
Serhiy Storchaka2085bd02019-06-01 11:00:15 +0300158 def __init__(self, /, *args, notifier=None, **kwds):
Pablo Galindo7c994542019-03-16 22:34:24 +0000159 self.notifier = notifier
160 super().__init__(*args, **kwds)
161
162 def __delitem__(self, item):
163 super().__delitem__(item)
164
165 # Notify that the cache is empty. This is important because the
166 # pool keeps maintaining workers until the cache gets drained. This
167 # eliminates a race condition in which a task is finished after the
168 # the pool's _handle_workers method has enter another iteration of the
169 # loop. In this situation, the only event that can wake up the pool
170 # is the cache to be emptied (no more tasks available).
171 if not self:
172 self.notifier.put(None)
173
Benjamin Petersone711caf2008-06-11 16:44:04 +0000174class Pool(object):
175 '''
Georg Brandl92905032008-11-22 08:51:39 +0000176 Class which supports an async version of applying functions to arguments.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000177 '''
Richard Oudkerk80a5be12014-03-23 12:30:54 +0000178 _wrap_exception = True
179
Pablo Galindo3766f182019-02-11 17:29:00 +0000180 @staticmethod
181 def Process(ctx, *args, **kwds):
182 return ctx.Process(*args, **kwds)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000183
Jesse Noller1f0b6582010-01-27 03:36:01 +0000184 def __init__(self, processes=None, initializer=None, initargs=(),
Richard Oudkerkb1694cf2013-10-16 16:41:56 +0100185 maxtasksperchild=None, context=None):
Victor Stinner9a8d1d72018-12-20 20:33:51 +0100186 # Attributes initialized early to make sure that they exist in
187 # __del__() if __init__() raises an exception
188 self._pool = []
189 self._state = INIT
190
Richard Oudkerkb1694cf2013-10-16 16:41:56 +0100191 self._ctx = context or get_context()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000192 self._setup_queues()
Antoine Pitrouab745042018-01-18 10:38:03 +0100193 self._taskqueue = queue.SimpleQueue()
Pablo Galindo7c994542019-03-16 22:34:24 +0000194 # The _change_notifier queue exist to wake up self._handle_workers()
195 # when the cache (self._cache) is empty or when there is a change in
196 # the _state variable of the thread that runs _handle_workers.
197 self._change_notifier = self._ctx.SimpleQueue()
198 self._cache = _PoolCache(notifier=self._change_notifier)
Jesse Noller1f0b6582010-01-27 03:36:01 +0000199 self._maxtasksperchild = maxtasksperchild
200 self._initializer = initializer
201 self._initargs = initargs
Benjamin Petersone711caf2008-06-11 16:44:04 +0000202
203 if processes is None:
Charles-François Natali37cfb0a2013-06-28 19:25:45 +0200204 processes = os.cpu_count() or 1
Victor Stinner2fae27b2011-06-20 17:53:35 +0200205 if processes < 1:
206 raise ValueError("Number of processes must be at least 1")
Benjamin Petersone711caf2008-06-11 16:44:04 +0000207
Florent Xicluna5d1155c2011-10-28 14:45:05 +0200208 if initializer is not None and not callable(initializer):
Benjamin Petersonf47ed4a2009-04-11 20:45:40 +0000209 raise TypeError('initializer must be a callable')
210
Jesse Noller1f0b6582010-01-27 03:36:01 +0000211 self._processes = processes
Julien Palard5d236ca2018-11-04 23:40:32 +0100212 try:
213 self._repopulate_pool()
214 except Exception:
215 for p in self._pool:
216 if p.exitcode is None:
217 p.terminate()
218 for p in self._pool:
219 p.join()
220 raise
Jesse Noller1f0b6582010-01-27 03:36:01 +0000221
Pablo Galindo7c994542019-03-16 22:34:24 +0000222 sentinels = self._get_sentinels()
223
Jesse Noller1f0b6582010-01-27 03:36:01 +0000224 self._worker_handler = threading.Thread(
225 target=Pool._handle_workers,
Pablo Galindo3766f182019-02-11 17:29:00 +0000226 args=(self._cache, self._taskqueue, self._ctx, self.Process,
227 self._processes, self._pool, self._inqueue, self._outqueue,
228 self._initializer, self._initargs, self._maxtasksperchild,
Pablo Galindo7c994542019-03-16 22:34:24 +0000229 self._wrap_exception, sentinels, self._change_notifier)
Jesse Noller1f0b6582010-01-27 03:36:01 +0000230 )
231 self._worker_handler.daemon = True
232 self._worker_handler._state = RUN
233 self._worker_handler.start()
234
Victor Stinner9dfc7542018-12-06 08:51:47 +0100235
Benjamin Petersone711caf2008-06-11 16:44:04 +0000236 self._task_handler = threading.Thread(
237 target=Pool._handle_tasks,
Richard Oudkerke90cedb2013-10-28 23:11:58 +0000238 args=(self._taskqueue, self._quick_put, self._outqueue,
239 self._pool, self._cache)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000240 )
Benjamin Petersonfae4c622008-08-18 18:40:08 +0000241 self._task_handler.daemon = True
Benjamin Petersone711caf2008-06-11 16:44:04 +0000242 self._task_handler._state = RUN
243 self._task_handler.start()
244
245 self._result_handler = threading.Thread(
246 target=Pool._handle_results,
247 args=(self._outqueue, self._quick_get, self._cache)
248 )
Benjamin Petersonfae4c622008-08-18 18:40:08 +0000249 self._result_handler.daemon = True
Benjamin Petersone711caf2008-06-11 16:44:04 +0000250 self._result_handler._state = RUN
251 self._result_handler.start()
252
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100253 self._terminate = util.Finalize(
Benjamin Petersone711caf2008-06-11 16:44:04 +0000254 self, self._terminate_pool,
255 args=(self._taskqueue, self._inqueue, self._outqueue, self._pool,
Pablo Galindo7c994542019-03-16 22:34:24 +0000256 self._change_notifier, self._worker_handler, self._task_handler,
Jesse Noller1f0b6582010-01-27 03:36:01 +0000257 self._result_handler, self._cache),
Benjamin Petersone711caf2008-06-11 16:44:04 +0000258 exitpriority=15
259 )
Victor Stinner9a8d1d72018-12-20 20:33:51 +0100260 self._state = RUN
261
262 # Copy globals as function locals to make sure that they are available
263 # during Python shutdown when the Pool is destroyed.
264 def __del__(self, _warn=warnings.warn, RUN=RUN):
265 if self._state == RUN:
266 _warn(f"unclosed running multiprocessing pool {self!r}",
267 ResourceWarning, source=self)
Pablo Galindo7c994542019-03-16 22:34:24 +0000268 if getattr(self, '_change_notifier', None) is not None:
269 self._change_notifier.put(None)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000270
Victor Stinner2b417fb2018-12-14 11:13:18 +0100271 def __repr__(self):
272 cls = self.__class__
273 return (f'<{cls.__module__}.{cls.__qualname__} '
274 f'state={self._state} '
275 f'pool_size={len(self._pool)}>')
276
Pablo Galindo7c994542019-03-16 22:34:24 +0000277 def _get_sentinels(self):
278 task_queue_sentinels = [self._outqueue._reader]
279 self_notifier_sentinels = [self._change_notifier._reader]
280 return [*task_queue_sentinels, *self_notifier_sentinels]
281
282 @staticmethod
283 def _get_worker_sentinels(workers):
284 return [worker.sentinel for worker in
285 workers if hasattr(worker, "sentinel")]
286
Pablo Galindo3766f182019-02-11 17:29:00 +0000287 @staticmethod
288 def _join_exited_workers(pool):
Jesse Noller1f0b6582010-01-27 03:36:01 +0000289 """Cleanup after any worker processes which have exited due to reaching
290 their specified lifetime. Returns True if any workers were cleaned up.
291 """
292 cleaned = False
Pablo Galindo3766f182019-02-11 17:29:00 +0000293 for i in reversed(range(len(pool))):
294 worker = pool[i]
Jesse Noller1f0b6582010-01-27 03:36:01 +0000295 if worker.exitcode is not None:
296 # worker exited
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100297 util.debug('cleaning up worker %d' % i)
Jesse Noller1f0b6582010-01-27 03:36:01 +0000298 worker.join()
299 cleaned = True
Pablo Galindo3766f182019-02-11 17:29:00 +0000300 del pool[i]
Jesse Noller1f0b6582010-01-27 03:36:01 +0000301 return cleaned
302
303 def _repopulate_pool(self):
Pablo Galindo3766f182019-02-11 17:29:00 +0000304 return self._repopulate_pool_static(self._ctx, self.Process,
305 self._processes,
306 self._pool, self._inqueue,
307 self._outqueue, self._initializer,
308 self._initargs,
309 self._maxtasksperchild,
310 self._wrap_exception)
311
312 @staticmethod
313 def _repopulate_pool_static(ctx, Process, processes, pool, inqueue,
314 outqueue, initializer, initargs,
315 maxtasksperchild, wrap_exception):
Jesse Noller1f0b6582010-01-27 03:36:01 +0000316 """Bring the number of pool processes up to the specified number,
317 for use after reaping workers which have exited.
318 """
Pablo Galindo3766f182019-02-11 17:29:00 +0000319 for i in range(processes - len(pool)):
320 w = Process(ctx, target=worker,
321 args=(inqueue, outqueue,
322 initializer,
323 initargs, maxtasksperchild,
324 wrap_exception))
Jesse Noller1f0b6582010-01-27 03:36:01 +0000325 w.name = w.name.replace('Process', 'PoolWorker')
326 w.daemon = True
327 w.start()
Pablo Galindo3766f182019-02-11 17:29:00 +0000328 pool.append(w)
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100329 util.debug('added worker')
Jesse Noller1f0b6582010-01-27 03:36:01 +0000330
Pablo Galindo3766f182019-02-11 17:29:00 +0000331 @staticmethod
332 def _maintain_pool(ctx, Process, processes, pool, inqueue, outqueue,
333 initializer, initargs, maxtasksperchild,
334 wrap_exception):
Jesse Noller1f0b6582010-01-27 03:36:01 +0000335 """Clean up any exited workers and start replacements for them.
336 """
Pablo Galindo3766f182019-02-11 17:29:00 +0000337 if Pool._join_exited_workers(pool):
338 Pool._repopulate_pool_static(ctx, Process, processes, pool,
339 inqueue, outqueue, initializer,
340 initargs, maxtasksperchild,
341 wrap_exception)
Jesse Noller1f0b6582010-01-27 03:36:01 +0000342
Benjamin Petersone711caf2008-06-11 16:44:04 +0000343 def _setup_queues(self):
Richard Oudkerkb1694cf2013-10-16 16:41:56 +0100344 self._inqueue = self._ctx.SimpleQueue()
345 self._outqueue = self._ctx.SimpleQueue()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000346 self._quick_put = self._inqueue._writer.send
347 self._quick_get = self._outqueue._reader.recv
348
Victor Stinner08c2ba02018-12-13 02:15:30 +0100349 def _check_running(self):
350 if self._state != RUN:
351 raise ValueError("Pool not running")
352
Benjamin Petersone711caf2008-06-11 16:44:04 +0000353 def apply(self, func, args=(), kwds={}):
354 '''
Georg Brandl92905032008-11-22 08:51:39 +0000355 Equivalent of `func(*args, **kwds)`.
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500356 Pool must be running.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000357 '''
Benjamin Petersone711caf2008-06-11 16:44:04 +0000358 return self.apply_async(func, args, kwds).get()
359
360 def map(self, func, iterable, chunksize=None):
361 '''
Georg Brandl92905032008-11-22 08:51:39 +0000362 Apply `func` to each element in `iterable`, collecting the results
363 in a list that is returned.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000364 '''
Antoine Pitroude911b22011-12-21 11:03:24 +0100365 return self._map_async(func, iterable, mapstar, chunksize).get()
366
367 def starmap(self, func, iterable, chunksize=None):
368 '''
369 Like `map()` method but the elements of the `iterable` are expected to
370 be iterables as well and will be unpacked as arguments. Hence
371 `func` and (a, b) becomes func(a, b).
372 '''
Antoine Pitroude911b22011-12-21 11:03:24 +0100373 return self._map_async(func, iterable, starmapstar, chunksize).get()
374
375 def starmap_async(self, func, iterable, chunksize=None, callback=None,
376 error_callback=None):
377 '''
378 Asynchronous version of `starmap()` method.
379 '''
Antoine Pitroude911b22011-12-21 11:03:24 +0100380 return self._map_async(func, iterable, starmapstar, chunksize,
381 callback, error_callback)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000382
Xiang Zhang794623b2017-03-29 11:58:54 +0800383 def _guarded_task_generation(self, result_job, func, iterable):
384 '''Provides a generator of tasks for imap and imap_unordered with
385 appropriate handling for iterables which throw exceptions during
386 iteration.'''
387 try:
388 i = -1
389 for i, x in enumerate(iterable):
390 yield (result_job, i, func, (x,), {})
391 except Exception as e:
392 yield (result_job, i+1, _helper_reraises_exception, (e,), {})
393
Benjamin Petersone711caf2008-06-11 16:44:04 +0000394 def imap(self, func, iterable, chunksize=1):
395 '''
Georg Brandl92905032008-11-22 08:51:39 +0000396 Equivalent of `map()` -- can be MUCH slower than `Pool.map()`.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000397 '''
Victor Stinner08c2ba02018-12-13 02:15:30 +0100398 self._check_running()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000399 if chunksize == 1:
Pablo Galindo3766f182019-02-11 17:29:00 +0000400 result = IMapIterator(self)
Xiang Zhang794623b2017-03-29 11:58:54 +0800401 self._taskqueue.put(
402 (
403 self._guarded_task_generation(result._job, func, iterable),
404 result._set_length
405 ))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000406 return result
407 else:
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500408 if chunksize < 1:
409 raise ValueError(
410 "Chunksize must be 1+, not {0:n}".format(
411 chunksize))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000412 task_batches = Pool._get_tasks(func, iterable, chunksize)
Pablo Galindo3766f182019-02-11 17:29:00 +0000413 result = IMapIterator(self)
Xiang Zhang794623b2017-03-29 11:58:54 +0800414 self._taskqueue.put(
415 (
416 self._guarded_task_generation(result._job,
417 mapstar,
418 task_batches),
419 result._set_length
420 ))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000421 return (item for chunk in result for item in chunk)
422
423 def imap_unordered(self, func, iterable, chunksize=1):
424 '''
Georg Brandl92905032008-11-22 08:51:39 +0000425 Like `imap()` method but ordering of results is arbitrary.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000426 '''
Victor Stinner08c2ba02018-12-13 02:15:30 +0100427 self._check_running()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000428 if chunksize == 1:
Pablo Galindo3766f182019-02-11 17:29:00 +0000429 result = IMapUnorderedIterator(self)
Xiang Zhang794623b2017-03-29 11:58:54 +0800430 self._taskqueue.put(
431 (
432 self._guarded_task_generation(result._job, func, iterable),
433 result._set_length
434 ))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000435 return result
436 else:
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500437 if chunksize < 1:
438 raise ValueError(
439 "Chunksize must be 1+, not {0!r}".format(chunksize))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000440 task_batches = Pool._get_tasks(func, iterable, chunksize)
Pablo Galindo3766f182019-02-11 17:29:00 +0000441 result = IMapUnorderedIterator(self)
Xiang Zhang794623b2017-03-29 11:58:54 +0800442 self._taskqueue.put(
443 (
444 self._guarded_task_generation(result._job,
445 mapstar,
446 task_batches),
447 result._set_length
448 ))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000449 return (item for chunk in result for item in chunk)
450
Ask Solem2afcbf22010-11-09 20:55:52 +0000451 def apply_async(self, func, args=(), kwds={}, callback=None,
452 error_callback=None):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000453 '''
Georg Brandl92905032008-11-22 08:51:39 +0000454 Asynchronous version of `apply()` method.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000455 '''
Victor Stinner08c2ba02018-12-13 02:15:30 +0100456 self._check_running()
Pablo Galindo3766f182019-02-11 17:29:00 +0000457 result = ApplyResult(self, callback, error_callback)
Xiang Zhang794623b2017-03-29 11:58:54 +0800458 self._taskqueue.put(([(result._job, 0, func, args, kwds)], None))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000459 return result
460
Ask Solem2afcbf22010-11-09 20:55:52 +0000461 def map_async(self, func, iterable, chunksize=None, callback=None,
462 error_callback=None):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000463 '''
Georg Brandl92905032008-11-22 08:51:39 +0000464 Asynchronous version of `map()` method.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000465 '''
Hynek Schlawack254af262012-10-27 12:53:02 +0200466 return self._map_async(func, iterable, mapstar, chunksize, callback,
467 error_callback)
Antoine Pitroude911b22011-12-21 11:03:24 +0100468
469 def _map_async(self, func, iterable, mapper, chunksize=None, callback=None,
470 error_callback=None):
471 '''
472 Helper function to implement map, starmap and their async counterparts.
473 '''
Victor Stinner08c2ba02018-12-13 02:15:30 +0100474 self._check_running()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000475 if not hasattr(iterable, '__len__'):
476 iterable = list(iterable)
477
478 if chunksize is None:
479 chunksize, extra = divmod(len(iterable), len(self._pool) * 4)
480 if extra:
481 chunksize += 1
Alexandre Vassalottie52e3782009-07-17 09:18:18 +0000482 if len(iterable) == 0:
483 chunksize = 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000484
485 task_batches = Pool._get_tasks(func, iterable, chunksize)
Pablo Galindo3766f182019-02-11 17:29:00 +0000486 result = MapResult(self, chunksize, len(iterable), callback,
Ask Solem2afcbf22010-11-09 20:55:52 +0000487 error_callback=error_callback)
Xiang Zhang794623b2017-03-29 11:58:54 +0800488 self._taskqueue.put(
489 (
490 self._guarded_task_generation(result._job,
491 mapper,
492 task_batches),
493 None
494 )
495 )
Benjamin Petersone711caf2008-06-11 16:44:04 +0000496 return result
497
498 @staticmethod
Pablo Galindo7c994542019-03-16 22:34:24 +0000499 def _wait_for_updates(sentinels, change_notifier, timeout=None):
500 wait(sentinels, timeout=timeout)
501 while not change_notifier.empty():
502 change_notifier.get()
503
504 @classmethod
505 def _handle_workers(cls, cache, taskqueue, ctx, Process, processes,
506 pool, inqueue, outqueue, initializer, initargs,
507 maxtasksperchild, wrap_exception, sentinels,
508 change_notifier):
Charles-François Natalif8859e12011-10-24 18:45:29 +0200509 thread = threading.current_thread()
510
511 # Keep maintaining workers until the cache gets drained, unless the pool
512 # is terminated.
Pablo Galindo3766f182019-02-11 17:29:00 +0000513 while thread._state == RUN or (cache and thread._state != TERMINATE):
Pablo Galindo7c994542019-03-16 22:34:24 +0000514 cls._maintain_pool(ctx, Process, processes, pool, inqueue,
515 outqueue, initializer, initargs,
516 maxtasksperchild, wrap_exception)
517
518 current_sentinels = [*cls._get_worker_sentinels(pool), *sentinels]
519
520 cls._wait_for_updates(current_sentinels, change_notifier)
Antoine Pitrou81dee6b2011-04-11 00:18:59 +0200521 # send sentinel to stop workers
Pablo Galindo3766f182019-02-11 17:29:00 +0000522 taskqueue.put(None)
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100523 util.debug('worker handler exiting')
Jesse Noller1f0b6582010-01-27 03:36:01 +0000524
525 @staticmethod
Richard Oudkerke90cedb2013-10-28 23:11:58 +0000526 def _handle_tasks(taskqueue, put, outqueue, pool, cache):
Benjamin Peterson672b8032008-06-11 19:14:14 +0000527 thread = threading.current_thread()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000528
529 for taskseq, set_length in iter(taskqueue.get, None):
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200530 task = None
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200531 try:
Xiang Zhang794623b2017-03-29 11:58:54 +0800532 # iterating taskseq cannot fail
533 for task in taskseq:
Victor Stinner2b417fb2018-12-14 11:13:18 +0100534 if thread._state != RUN:
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200535 util.debug('task handler found thread._state != RUN')
536 break
Richard Oudkerke90cedb2013-10-28 23:11:58 +0000537 try:
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200538 put(task)
539 except Exception as e:
Xiang Zhang794623b2017-03-29 11:58:54 +0800540 job, idx = task[:2]
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200541 try:
Xiang Zhang794623b2017-03-29 11:58:54 +0800542 cache[job]._set(idx, (False, e))
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200543 except KeyError:
544 pass
545 else:
546 if set_length:
547 util.debug('doing set_length()')
Xiang Zhang794623b2017-03-29 11:58:54 +0800548 idx = task[1] if task else -1
549 set_length(idx + 1)
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200550 continue
551 break
Antoine Pitrou89889452017-03-24 13:52:11 +0100552 finally:
553 task = taskseq = job = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000554 else:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100555 util.debug('task handler got sentinel')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000556
Benjamin Petersone711caf2008-06-11 16:44:04 +0000557 try:
558 # tell result handler to finish when cache is empty
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100559 util.debug('task handler sending sentinel to result handler')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000560 outqueue.put(None)
561
562 # tell workers there is no more work
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100563 util.debug('task handler sending sentinel to workers')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000564 for p in pool:
565 put(None)
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200566 except OSError:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100567 util.debug('task handler got OSError when sending sentinels')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000568
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100569 util.debug('task handler exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000570
571 @staticmethod
572 def _handle_results(outqueue, get, cache):
Benjamin Peterson672b8032008-06-11 19:14:14 +0000573 thread = threading.current_thread()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000574
575 while 1:
576 try:
577 task = get()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200578 except (OSError, EOFError):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100579 util.debug('result handler got EOFError/OSError -- exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000580 return
581
Victor Stinner2dfe3512018-12-16 23:40:49 +0100582 if thread._state != RUN:
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500583 assert thread._state == TERMINATE, "Thread not in TERMINATE"
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100584 util.debug('result handler found thread._state=TERMINATE')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000585 break
586
587 if task is None:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100588 util.debug('result handler got sentinel')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000589 break
590
591 job, i, obj = task
592 try:
593 cache[job]._set(i, obj)
594 except KeyError:
595 pass
Antoine Pitrou89889452017-03-24 13:52:11 +0100596 task = job = obj = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000597
598 while cache and thread._state != TERMINATE:
599 try:
600 task = get()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200601 except (OSError, EOFError):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100602 util.debug('result handler got EOFError/OSError -- exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000603 return
604
605 if task is None:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100606 util.debug('result handler ignoring extra sentinel')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000607 continue
608 job, i, obj = task
609 try:
610 cache[job]._set(i, obj)
611 except KeyError:
612 pass
Antoine Pitrou89889452017-03-24 13:52:11 +0100613 task = job = obj = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000614
615 if hasattr(outqueue, '_reader'):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100616 util.debug('ensuring that outqueue is not full')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000617 # If we don't make room available in outqueue then
618 # attempts to add the sentinel (None) to outqueue may
619 # block. There is guaranteed to be no more than 2 sentinels.
620 try:
621 for i in range(10):
622 if not outqueue._reader.poll():
623 break
624 get()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200625 except (OSError, EOFError):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000626 pass
627
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100628 util.debug('result handler exiting: len(cache)=%s, thread._state=%s',
Benjamin Petersone711caf2008-06-11 16:44:04 +0000629 len(cache), thread._state)
630
631 @staticmethod
632 def _get_tasks(func, it, size):
633 it = iter(it)
634 while 1:
635 x = tuple(itertools.islice(it, size))
636 if not x:
637 return
638 yield (func, x)
639
640 def __reduce__(self):
641 raise NotImplementedError(
642 'pool objects cannot be passed between processes or pickled'
643 )
644
645 def close(self):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100646 util.debug('closing pool')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000647 if self._state == RUN:
648 self._state = CLOSE
Jesse Noller1f0b6582010-01-27 03:36:01 +0000649 self._worker_handler._state = CLOSE
Pablo Galindo7c994542019-03-16 22:34:24 +0000650 self._change_notifier.put(None)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000651
652 def terminate(self):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100653 util.debug('terminating pool')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000654 self._state = TERMINATE
655 self._terminate()
656
657 def join(self):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100658 util.debug('joining pool')
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500659 if self._state == RUN:
660 raise ValueError("Pool is still running")
661 elif self._state not in (CLOSE, TERMINATE):
662 raise ValueError("In unknown state")
Jesse Noller1f0b6582010-01-27 03:36:01 +0000663 self._worker_handler.join()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000664 self._task_handler.join()
665 self._result_handler.join()
666 for p in self._pool:
667 p.join()
668
669 @staticmethod
670 def _help_stuff_finish(inqueue, task_handler, size):
671 # task_handler may be blocked trying to put items on inqueue
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100672 util.debug('removing tasks from inqueue until task handler finished')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000673 inqueue._rlock.acquire()
Benjamin Peterson672b8032008-06-11 19:14:14 +0000674 while task_handler.is_alive() and inqueue._reader.poll():
Benjamin Petersone711caf2008-06-11 16:44:04 +0000675 inqueue._reader.recv()
676 time.sleep(0)
677
678 @classmethod
Pablo Galindo7c994542019-03-16 22:34:24 +0000679 def _terminate_pool(cls, taskqueue, inqueue, outqueue, pool, change_notifier,
Jesse Noller1f0b6582010-01-27 03:36:01 +0000680 worker_handler, task_handler, result_handler, cache):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000681 # this is guaranteed to only be called once
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100682 util.debug('finalizing pool')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000683
Batuhan Taşkayaac10e0c2020-03-15 22:45:56 +0300684 # Notify that the worker_handler state has been changed so the
685 # _handle_workers loop can be unblocked (and exited) in order to
686 # send the finalization sentinel all the workers.
Jesse Noller1f0b6582010-01-27 03:36:01 +0000687 worker_handler._state = TERMINATE
Batuhan Taşkayaac10e0c2020-03-15 22:45:56 +0300688 change_notifier.put(None)
689
Benjamin Petersone711caf2008-06-11 16:44:04 +0000690 task_handler._state = TERMINATE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000691
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100692 util.debug('helping task handler/workers to finish')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000693 cls._help_stuff_finish(inqueue, task_handler, len(pool))
694
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500695 if (not result_handler.is_alive()) and (len(cache) != 0):
696 raise AssertionError(
697 "Cannot have cache with result_hander not alive")
Benjamin Petersone711caf2008-06-11 16:44:04 +0000698
699 result_handler._state = TERMINATE
Pablo Galindo7c994542019-03-16 22:34:24 +0000700 change_notifier.put(None)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000701 outqueue.put(None) # sentinel
702
Antoine Pitrou81dee6b2011-04-11 00:18:59 +0200703 # We must wait for the worker handler to exit before terminating
704 # workers because we don't want workers to be restarted behind our back.
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100705 util.debug('joining worker handler')
Richard Oudkerkf29ec4b2012-06-18 15:54:57 +0100706 if threading.current_thread() is not worker_handler:
707 worker_handler.join()
Antoine Pitrou81dee6b2011-04-11 00:18:59 +0200708
Jesse Noller1f0b6582010-01-27 03:36:01 +0000709 # Terminate workers which haven't already finished.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000710 if pool and hasattr(pool[0], 'terminate'):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100711 util.debug('terminating workers')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000712 for p in pool:
Jesse Noller1f0b6582010-01-27 03:36:01 +0000713 if p.exitcode is None:
714 p.terminate()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000715
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100716 util.debug('joining task handler')
Richard Oudkerkf29ec4b2012-06-18 15:54:57 +0100717 if threading.current_thread() is not task_handler:
718 task_handler.join()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000719
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100720 util.debug('joining result handler')
Richard Oudkerkf29ec4b2012-06-18 15:54:57 +0100721 if threading.current_thread() is not result_handler:
722 result_handler.join()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000723
724 if pool and hasattr(pool[0], 'terminate'):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100725 util.debug('joining pool workers')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000726 for p in pool:
Florent Xicluna998171f2010-03-08 13:32:17 +0000727 if p.is_alive():
Jesse Noller1f0b6582010-01-27 03:36:01 +0000728 # worker has not yet exited
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100729 util.debug('cleaning up worker %d' % p.pid)
Florent Xicluna998171f2010-03-08 13:32:17 +0000730 p.join()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000731
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100732 def __enter__(self):
Victor Stinner08c2ba02018-12-13 02:15:30 +0100733 self._check_running()
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100734 return self
735
736 def __exit__(self, exc_type, exc_val, exc_tb):
737 self.terminate()
738
Benjamin Petersone711caf2008-06-11 16:44:04 +0000739#
740# Class whose instances are returned by `Pool.apply_async()`
741#
742
743class ApplyResult(object):
744
Pablo Galindo3766f182019-02-11 17:29:00 +0000745 def __init__(self, pool, callback, error_callback):
746 self._pool = pool
Richard Oudkerk692130a2012-05-25 13:26:53 +0100747 self._event = threading.Event()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000748 self._job = next(job_counter)
Pablo Galindo3766f182019-02-11 17:29:00 +0000749 self._cache = pool._cache
Benjamin Petersone711caf2008-06-11 16:44:04 +0000750 self._callback = callback
Ask Solem2afcbf22010-11-09 20:55:52 +0000751 self._error_callback = error_callback
Pablo Galindo3766f182019-02-11 17:29:00 +0000752 self._cache[self._job] = self
Benjamin Petersone711caf2008-06-11 16:44:04 +0000753
754 def ready(self):
Richard Oudkerk692130a2012-05-25 13:26:53 +0100755 return self._event.is_set()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000756
757 def successful(self):
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500758 if not self.ready():
759 raise ValueError("{0!r} not ready".format(self))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000760 return self._success
761
762 def wait(self, timeout=None):
Richard Oudkerk692130a2012-05-25 13:26:53 +0100763 self._event.wait(timeout)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000764
765 def get(self, timeout=None):
766 self.wait(timeout)
Richard Oudkerk692130a2012-05-25 13:26:53 +0100767 if not self.ready():
Benjamin Petersone711caf2008-06-11 16:44:04 +0000768 raise TimeoutError
769 if self._success:
770 return self._value
771 else:
772 raise self._value
773
774 def _set(self, i, obj):
775 self._success, self._value = obj
776 if self._callback and self._success:
777 self._callback(self._value)
Ask Solem2afcbf22010-11-09 20:55:52 +0000778 if self._error_callback and not self._success:
779 self._error_callback(self._value)
Richard Oudkerk692130a2012-05-25 13:26:53 +0100780 self._event.set()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000781 del self._cache[self._job]
Pablo Galindo3766f182019-02-11 17:29:00 +0000782 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000783
Batuhan Taşkaya03615562020-04-10 17:46:36 +0300784 __class_getitem__ = classmethod(types.GenericAlias)
785
Richard Oudkerkdef51ca2013-05-06 12:10:04 +0100786AsyncResult = ApplyResult # create alias -- see #17805
787
Benjamin Petersone711caf2008-06-11 16:44:04 +0000788#
789# Class whose instances are returned by `Pool.map_async()`
790#
791
792class MapResult(ApplyResult):
793
Pablo Galindo3766f182019-02-11 17:29:00 +0000794 def __init__(self, pool, chunksize, length, callback, error_callback):
795 ApplyResult.__init__(self, pool, callback,
Ask Solem2afcbf22010-11-09 20:55:52 +0000796 error_callback=error_callback)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000797 self._success = True
798 self._value = [None] * length
799 self._chunksize = chunksize
800 if chunksize <= 0:
801 self._number_left = 0
Richard Oudkerk692130a2012-05-25 13:26:53 +0100802 self._event.set()
Pablo Galindo3766f182019-02-11 17:29:00 +0000803 del self._cache[self._job]
Benjamin Petersone711caf2008-06-11 16:44:04 +0000804 else:
805 self._number_left = length//chunksize + bool(length % chunksize)
806
807 def _set(self, i, success_result):
Charles-François Natali78f55ff2016-02-10 22:58:18 +0000808 self._number_left -= 1
Benjamin Petersone711caf2008-06-11 16:44:04 +0000809 success, result = success_result
Charles-François Natali78f55ff2016-02-10 22:58:18 +0000810 if success and self._success:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000811 self._value[i*self._chunksize:(i+1)*self._chunksize] = result
Benjamin Petersone711caf2008-06-11 16:44:04 +0000812 if self._number_left == 0:
813 if self._callback:
814 self._callback(self._value)
815 del self._cache[self._job]
Richard Oudkerk692130a2012-05-25 13:26:53 +0100816 self._event.set()
Pablo Galindo3766f182019-02-11 17:29:00 +0000817 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000818 else:
Charles-François Natali78f55ff2016-02-10 22:58:18 +0000819 if not success and self._success:
820 # only store first exception
821 self._success = False
822 self._value = result
823 if self._number_left == 0:
824 # only consider the result ready once all jobs are done
825 if self._error_callback:
826 self._error_callback(self._value)
827 del self._cache[self._job]
828 self._event.set()
Pablo Galindo3766f182019-02-11 17:29:00 +0000829 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000830
831#
832# Class whose instances are returned by `Pool.imap()`
833#
834
835class IMapIterator(object):
836
Pablo Galindo3766f182019-02-11 17:29:00 +0000837 def __init__(self, pool):
838 self._pool = pool
Benjamin Petersone711caf2008-06-11 16:44:04 +0000839 self._cond = threading.Condition(threading.Lock())
840 self._job = next(job_counter)
Pablo Galindo3766f182019-02-11 17:29:00 +0000841 self._cache = pool._cache
Benjamin Petersone711caf2008-06-11 16:44:04 +0000842 self._items = collections.deque()
843 self._index = 0
844 self._length = None
845 self._unsorted = {}
Pablo Galindo3766f182019-02-11 17:29:00 +0000846 self._cache[self._job] = self
Benjamin Petersone711caf2008-06-11 16:44:04 +0000847
848 def __iter__(self):
849 return self
850
851 def next(self, timeout=None):
Charles-François Natalia924fc72014-05-25 14:12:12 +0100852 with self._cond:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000853 try:
854 item = self._items.popleft()
855 except IndexError:
856 if self._index == self._length:
Pablo Galindo3766f182019-02-11 17:29:00 +0000857 self._pool = None
Serhiy Storchaka5affd232017-04-05 09:37:24 +0300858 raise StopIteration from None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000859 self._cond.wait(timeout)
860 try:
861 item = self._items.popleft()
862 except IndexError:
863 if self._index == self._length:
Pablo Galindo3766f182019-02-11 17:29:00 +0000864 self._pool = None
Serhiy Storchaka5affd232017-04-05 09:37:24 +0300865 raise StopIteration from None
866 raise TimeoutError from None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000867
868 success, value = item
869 if success:
870 return value
871 raise value
872
873 __next__ = next # XXX
874
875 def _set(self, i, obj):
Charles-François Natalia924fc72014-05-25 14:12:12 +0100876 with self._cond:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000877 if self._index == i:
878 self._items.append(obj)
879 self._index += 1
880 while self._index in self._unsorted:
881 obj = self._unsorted.pop(self._index)
882 self._items.append(obj)
883 self._index += 1
884 self._cond.notify()
885 else:
886 self._unsorted[i] = obj
887
888 if self._index == self._length:
889 del self._cache[self._job]
Pablo Galindo3766f182019-02-11 17:29:00 +0000890 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000891
892 def _set_length(self, length):
Charles-François Natalia924fc72014-05-25 14:12:12 +0100893 with self._cond:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000894 self._length = length
895 if self._index == self._length:
896 self._cond.notify()
897 del self._cache[self._job]
Pablo Galindo3766f182019-02-11 17:29:00 +0000898 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000899
900#
901# Class whose instances are returned by `Pool.imap_unordered()`
902#
903
904class IMapUnorderedIterator(IMapIterator):
905
906 def _set(self, i, obj):
Charles-François Natalia924fc72014-05-25 14:12:12 +0100907 with self._cond:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000908 self._items.append(obj)
909 self._index += 1
910 self._cond.notify()
911 if self._index == self._length:
912 del self._cache[self._job]
Pablo Galindo3766f182019-02-11 17:29:00 +0000913 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000914
915#
916#
917#
918
919class ThreadPool(Pool):
Richard Oudkerk80a5be12014-03-23 12:30:54 +0000920 _wrap_exception = False
Benjamin Petersone711caf2008-06-11 16:44:04 +0000921
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100922 @staticmethod
Pablo Galindo3766f182019-02-11 17:29:00 +0000923 def Process(ctx, *args, **kwds):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100924 from .dummy import Process
925 return Process(*args, **kwds)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000926
927 def __init__(self, processes=None, initializer=None, initargs=()):
928 Pool.__init__(self, processes, initializer, initargs)
929
930 def _setup_queues(self):
Antoine Pitrouab745042018-01-18 10:38:03 +0100931 self._inqueue = queue.SimpleQueue()
932 self._outqueue = queue.SimpleQueue()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000933 self._quick_put = self._inqueue.put
934 self._quick_get = self._outqueue.get
935
Pablo Galindo7c994542019-03-16 22:34:24 +0000936 def _get_sentinels(self):
937 return [self._change_notifier._reader]
938
939 @staticmethod
940 def _get_worker_sentinels(workers):
941 return []
942
Benjamin Petersone711caf2008-06-11 16:44:04 +0000943 @staticmethod
944 def _help_stuff_finish(inqueue, task_handler, size):
Antoine Pitrouab745042018-01-18 10:38:03 +0100945 # drain inqueue, and put sentinels at its head to make workers finish
946 try:
947 while True:
948 inqueue.get(block=False)
949 except queue.Empty:
950 pass
951 for i in range(size):
952 inqueue.put(None)
Pablo Galindo7c994542019-03-16 22:34:24 +0000953
954 def _wait_for_updates(self, sentinels, change_notifier, timeout):
955 time.sleep(timeout)