blob: b223d6aa724bb63e0fe5e8e19786d64c6359b96d [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
Victor Stinner9a8d1d72018-12-20 20:33:51 +010023import warnings
Pablo Galindo7c994542019-03-16 22:34:24 +000024from queue import Empty
Benjamin Petersone711caf2008-06-11 16:44:04 +000025
Richard Oudkerk84ed9a62013-08-14 15:35:41 +010026# If threading is available then ThreadPool should be provided. Therefore
27# we avoid top-level imports which are liable to fail on some systems.
28from . import util
Victor Stinner7fa767e2014-03-20 09:16:38 +010029from . import get_context, TimeoutError
Pablo Galindo7c994542019-03-16 22:34:24 +000030from .connection import wait
Benjamin Petersone711caf2008-06-11 16:44:04 +000031
32#
33# Constants representing the state of a pool
34#
35
Victor Stinner9a8d1d72018-12-20 20:33:51 +010036INIT = "INIT"
Victor Stinner2b417fb2018-12-14 11:13:18 +010037RUN = "RUN"
38CLOSE = "CLOSE"
39TERMINATE = "TERMINATE"
Benjamin Petersone711caf2008-06-11 16:44:04 +000040
41#
42# Miscellaneous
43#
44
45job_counter = itertools.count()
46
47def mapstar(args):
48 return list(map(*args))
49
Antoine Pitroude911b22011-12-21 11:03:24 +010050def starmapstar(args):
51 return list(itertools.starmap(args[0], args[1]))
52
Benjamin Petersone711caf2008-06-11 16:44:04 +000053#
Richard Oudkerk85757832013-05-06 11:38:25 +010054# Hack to embed stringification of remote traceback in local traceback
55#
56
57class RemoteTraceback(Exception):
58 def __init__(self, tb):
59 self.tb = tb
60 def __str__(self):
61 return self.tb
62
63class ExceptionWithTraceback:
64 def __init__(self, exc, tb):
65 tb = traceback.format_exception(type(exc), exc, tb)
66 tb = ''.join(tb)
67 self.exc = exc
68 self.tb = '\n"""\n%s"""' % tb
69 def __reduce__(self):
70 return rebuild_exc, (self.exc, self.tb)
71
72def rebuild_exc(exc, tb):
73 exc.__cause__ = RemoteTraceback(tb)
74 return exc
75
76#
Benjamin Petersone711caf2008-06-11 16:44:04 +000077# Code run by worker processes
78#
79
Ask Solem2afcbf22010-11-09 20:55:52 +000080class MaybeEncodingError(Exception):
81 """Wraps possible unpickleable errors, so they can be
82 safely sent through the socket."""
83
84 def __init__(self, exc, value):
85 self.exc = repr(exc)
86 self.value = repr(value)
87 super(MaybeEncodingError, self).__init__(self.exc, self.value)
88
89 def __str__(self):
90 return "Error sending result: '%s'. Reason: '%s'" % (self.value,
91 self.exc)
92
93 def __repr__(self):
Serhiy Storchaka465e60e2014-07-25 23:36:00 +030094 return "<%s: %s>" % (self.__class__.__name__, self)
Ask Solem2afcbf22010-11-09 20:55:52 +000095
96
Richard Oudkerk80a5be12014-03-23 12:30:54 +000097def worker(inqueue, outqueue, initializer=None, initargs=(), maxtasks=None,
98 wrap_exception=False):
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -050099 if (maxtasks is not None) and not (isinstance(maxtasks, int)
100 and maxtasks >= 1):
101 raise AssertionError("Maxtasks {!r} is not valid".format(maxtasks))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000102 put = outqueue.put
103 get = inqueue.get
104 if hasattr(inqueue, '_writer'):
105 inqueue._writer.close()
106 outqueue._reader.close()
107
108 if initializer is not None:
109 initializer(*initargs)
110
Jesse Noller1f0b6582010-01-27 03:36:01 +0000111 completed = 0
112 while maxtasks is None or (maxtasks and completed < maxtasks):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000113 try:
114 task = get()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200115 except (EOFError, OSError):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100116 util.debug('worker got EOFError or OSError -- exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000117 break
118
119 if task is None:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100120 util.debug('worker got sentinel -- exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000121 break
122
123 job, i, func, args, kwds = task
124 try:
125 result = (True, func(*args, **kwds))
126 except Exception as e:
Xiang Zhang794623b2017-03-29 11:58:54 +0800127 if wrap_exception and func is not _helper_reraises_exception:
Richard Oudkerk80a5be12014-03-23 12:30:54 +0000128 e = ExceptionWithTraceback(e, e.__traceback__)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000129 result = (False, e)
Ask Solem2afcbf22010-11-09 20:55:52 +0000130 try:
131 put((job, i, result))
132 except Exception as e:
133 wrapped = MaybeEncodingError(e, result[1])
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100134 util.debug("Possible encoding error while sending result: %s" % (
Ask Solem2afcbf22010-11-09 20:55:52 +0000135 wrapped))
136 put((job, i, (False, wrapped)))
Antoine Pitrou89889452017-03-24 13:52:11 +0100137
138 task = job = result = func = args = kwds = None
Jesse Noller1f0b6582010-01-27 03:36:01 +0000139 completed += 1
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100140 util.debug('worker exiting after %d tasks' % completed)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000141
Xiang Zhang794623b2017-03-29 11:58:54 +0800142def _helper_reraises_exception(ex):
143 'Pickle-able helper function for use by _guarded_task_generation.'
144 raise ex
145
Benjamin Petersone711caf2008-06-11 16:44:04 +0000146#
147# Class representing a process pool
148#
149
Pablo Galindo7c994542019-03-16 22:34:24 +0000150class _PoolCache(dict):
151 """
152 Class that implements a cache for the Pool class that will notify
153 the pool management threads every time the cache is emptied. The
154 notification is done by the use of a queue that is provided when
155 instantiating the cache.
156 """
Serhiy Storchaka2085bd02019-06-01 11:00:15 +0300157 def __init__(self, /, *args, notifier=None, **kwds):
Pablo Galindo7c994542019-03-16 22:34:24 +0000158 self.notifier = notifier
159 super().__init__(*args, **kwds)
160
161 def __delitem__(self, item):
162 super().__delitem__(item)
163
164 # Notify that the cache is empty. This is important because the
165 # pool keeps maintaining workers until the cache gets drained. This
166 # eliminates a race condition in which a task is finished after the
167 # the pool's _handle_workers method has enter another iteration of the
168 # loop. In this situation, the only event that can wake up the pool
169 # is the cache to be emptied (no more tasks available).
170 if not self:
171 self.notifier.put(None)
172
Benjamin Petersone711caf2008-06-11 16:44:04 +0000173class Pool(object):
174 '''
Georg Brandl92905032008-11-22 08:51:39 +0000175 Class which supports an async version of applying functions to arguments.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000176 '''
Richard Oudkerk80a5be12014-03-23 12:30:54 +0000177 _wrap_exception = True
178
Pablo Galindo3766f182019-02-11 17:29:00 +0000179 @staticmethod
180 def Process(ctx, *args, **kwds):
181 return ctx.Process(*args, **kwds)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000182
Jesse Noller1f0b6582010-01-27 03:36:01 +0000183 def __init__(self, processes=None, initializer=None, initargs=(),
Richard Oudkerkb1694cf2013-10-16 16:41:56 +0100184 maxtasksperchild=None, context=None):
Victor Stinner9a8d1d72018-12-20 20:33:51 +0100185 # Attributes initialized early to make sure that they exist in
186 # __del__() if __init__() raises an exception
187 self._pool = []
188 self._state = INIT
189
Richard Oudkerkb1694cf2013-10-16 16:41:56 +0100190 self._ctx = context or get_context()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000191 self._setup_queues()
Antoine Pitrouab745042018-01-18 10:38:03 +0100192 self._taskqueue = queue.SimpleQueue()
Pablo Galindo7c994542019-03-16 22:34:24 +0000193 # The _change_notifier queue exist to wake up self._handle_workers()
194 # when the cache (self._cache) is empty or when there is a change in
195 # the _state variable of the thread that runs _handle_workers.
196 self._change_notifier = self._ctx.SimpleQueue()
197 self._cache = _PoolCache(notifier=self._change_notifier)
Jesse Noller1f0b6582010-01-27 03:36:01 +0000198 self._maxtasksperchild = maxtasksperchild
199 self._initializer = initializer
200 self._initargs = initargs
Benjamin Petersone711caf2008-06-11 16:44:04 +0000201
202 if processes is None:
Charles-François Natali37cfb0a2013-06-28 19:25:45 +0200203 processes = os.cpu_count() or 1
Victor Stinner2fae27b2011-06-20 17:53:35 +0200204 if processes < 1:
205 raise ValueError("Number of processes must be at least 1")
Benjamin Petersone711caf2008-06-11 16:44:04 +0000206
Florent Xicluna5d1155c2011-10-28 14:45:05 +0200207 if initializer is not None and not callable(initializer):
Benjamin Petersonf47ed4a2009-04-11 20:45:40 +0000208 raise TypeError('initializer must be a callable')
209
Jesse Noller1f0b6582010-01-27 03:36:01 +0000210 self._processes = processes
Julien Palard5d236ca2018-11-04 23:40:32 +0100211 try:
212 self._repopulate_pool()
213 except Exception:
214 for p in self._pool:
215 if p.exitcode is None:
216 p.terminate()
217 for p in self._pool:
218 p.join()
219 raise
Jesse Noller1f0b6582010-01-27 03:36:01 +0000220
Pablo Galindo7c994542019-03-16 22:34:24 +0000221 sentinels = self._get_sentinels()
222
Jesse Noller1f0b6582010-01-27 03:36:01 +0000223 self._worker_handler = threading.Thread(
224 target=Pool._handle_workers,
Pablo Galindo3766f182019-02-11 17:29:00 +0000225 args=(self._cache, self._taskqueue, self._ctx, self.Process,
226 self._processes, self._pool, self._inqueue, self._outqueue,
227 self._initializer, self._initargs, self._maxtasksperchild,
Pablo Galindo7c994542019-03-16 22:34:24 +0000228 self._wrap_exception, sentinels, self._change_notifier)
Jesse Noller1f0b6582010-01-27 03:36:01 +0000229 )
230 self._worker_handler.daemon = True
231 self._worker_handler._state = RUN
232 self._worker_handler.start()
233
Victor Stinner9dfc7542018-12-06 08:51:47 +0100234
Benjamin Petersone711caf2008-06-11 16:44:04 +0000235 self._task_handler = threading.Thread(
236 target=Pool._handle_tasks,
Richard Oudkerke90cedb2013-10-28 23:11:58 +0000237 args=(self._taskqueue, self._quick_put, self._outqueue,
238 self._pool, self._cache)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000239 )
Benjamin Petersonfae4c622008-08-18 18:40:08 +0000240 self._task_handler.daemon = True
Benjamin Petersone711caf2008-06-11 16:44:04 +0000241 self._task_handler._state = RUN
242 self._task_handler.start()
243
244 self._result_handler = threading.Thread(
245 target=Pool._handle_results,
246 args=(self._outqueue, self._quick_get, self._cache)
247 )
Benjamin Petersonfae4c622008-08-18 18:40:08 +0000248 self._result_handler.daemon = True
Benjamin Petersone711caf2008-06-11 16:44:04 +0000249 self._result_handler._state = RUN
250 self._result_handler.start()
251
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100252 self._terminate = util.Finalize(
Benjamin Petersone711caf2008-06-11 16:44:04 +0000253 self, self._terminate_pool,
254 args=(self._taskqueue, self._inqueue, self._outqueue, self._pool,
Pablo Galindo7c994542019-03-16 22:34:24 +0000255 self._change_notifier, self._worker_handler, self._task_handler,
Jesse Noller1f0b6582010-01-27 03:36:01 +0000256 self._result_handler, self._cache),
Benjamin Petersone711caf2008-06-11 16:44:04 +0000257 exitpriority=15
258 )
Victor Stinner9a8d1d72018-12-20 20:33:51 +0100259 self._state = RUN
260
261 # Copy globals as function locals to make sure that they are available
262 # during Python shutdown when the Pool is destroyed.
263 def __del__(self, _warn=warnings.warn, RUN=RUN):
264 if self._state == RUN:
265 _warn(f"unclosed running multiprocessing pool {self!r}",
266 ResourceWarning, source=self)
Pablo Galindo7c994542019-03-16 22:34:24 +0000267 if getattr(self, '_change_notifier', None) is not None:
268 self._change_notifier.put(None)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000269
Victor Stinner2b417fb2018-12-14 11:13:18 +0100270 def __repr__(self):
271 cls = self.__class__
272 return (f'<{cls.__module__}.{cls.__qualname__} '
273 f'state={self._state} '
274 f'pool_size={len(self._pool)}>')
275
Pablo Galindo7c994542019-03-16 22:34:24 +0000276 def _get_sentinels(self):
277 task_queue_sentinels = [self._outqueue._reader]
278 self_notifier_sentinels = [self._change_notifier._reader]
279 return [*task_queue_sentinels, *self_notifier_sentinels]
280
281 @staticmethod
282 def _get_worker_sentinels(workers):
283 return [worker.sentinel for worker in
284 workers if hasattr(worker, "sentinel")]
285
Pablo Galindo3766f182019-02-11 17:29:00 +0000286 @staticmethod
287 def _join_exited_workers(pool):
Jesse Noller1f0b6582010-01-27 03:36:01 +0000288 """Cleanup after any worker processes which have exited due to reaching
289 their specified lifetime. Returns True if any workers were cleaned up.
290 """
291 cleaned = False
Pablo Galindo3766f182019-02-11 17:29:00 +0000292 for i in reversed(range(len(pool))):
293 worker = pool[i]
Jesse Noller1f0b6582010-01-27 03:36:01 +0000294 if worker.exitcode is not None:
295 # worker exited
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100296 util.debug('cleaning up worker %d' % i)
Jesse Noller1f0b6582010-01-27 03:36:01 +0000297 worker.join()
298 cleaned = True
Pablo Galindo3766f182019-02-11 17:29:00 +0000299 del pool[i]
Jesse Noller1f0b6582010-01-27 03:36:01 +0000300 return cleaned
301
302 def _repopulate_pool(self):
Pablo Galindo3766f182019-02-11 17:29:00 +0000303 return self._repopulate_pool_static(self._ctx, self.Process,
304 self._processes,
305 self._pool, self._inqueue,
306 self._outqueue, self._initializer,
307 self._initargs,
308 self._maxtasksperchild,
309 self._wrap_exception)
310
311 @staticmethod
312 def _repopulate_pool_static(ctx, Process, processes, pool, inqueue,
313 outqueue, initializer, initargs,
314 maxtasksperchild, wrap_exception):
Jesse Noller1f0b6582010-01-27 03:36:01 +0000315 """Bring the number of pool processes up to the specified number,
316 for use after reaping workers which have exited.
317 """
Pablo Galindo3766f182019-02-11 17:29:00 +0000318 for i in range(processes - len(pool)):
319 w = Process(ctx, target=worker,
320 args=(inqueue, outqueue,
321 initializer,
322 initargs, maxtasksperchild,
323 wrap_exception))
Jesse Noller1f0b6582010-01-27 03:36:01 +0000324 w.name = w.name.replace('Process', 'PoolWorker')
325 w.daemon = True
326 w.start()
Pablo Galindo3766f182019-02-11 17:29:00 +0000327 pool.append(w)
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100328 util.debug('added worker')
Jesse Noller1f0b6582010-01-27 03:36:01 +0000329
Pablo Galindo3766f182019-02-11 17:29:00 +0000330 @staticmethod
331 def _maintain_pool(ctx, Process, processes, pool, inqueue, outqueue,
332 initializer, initargs, maxtasksperchild,
333 wrap_exception):
Jesse Noller1f0b6582010-01-27 03:36:01 +0000334 """Clean up any exited workers and start replacements for them.
335 """
Pablo Galindo3766f182019-02-11 17:29:00 +0000336 if Pool._join_exited_workers(pool):
337 Pool._repopulate_pool_static(ctx, Process, processes, pool,
338 inqueue, outqueue, initializer,
339 initargs, maxtasksperchild,
340 wrap_exception)
Jesse Noller1f0b6582010-01-27 03:36:01 +0000341
Benjamin Petersone711caf2008-06-11 16:44:04 +0000342 def _setup_queues(self):
Richard Oudkerkb1694cf2013-10-16 16:41:56 +0100343 self._inqueue = self._ctx.SimpleQueue()
344 self._outqueue = self._ctx.SimpleQueue()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000345 self._quick_put = self._inqueue._writer.send
346 self._quick_get = self._outqueue._reader.recv
347
Victor Stinner08c2ba02018-12-13 02:15:30 +0100348 def _check_running(self):
349 if self._state != RUN:
350 raise ValueError("Pool not running")
351
Benjamin Petersone711caf2008-06-11 16:44:04 +0000352 def apply(self, func, args=(), kwds={}):
353 '''
Georg Brandl92905032008-11-22 08:51:39 +0000354 Equivalent of `func(*args, **kwds)`.
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500355 Pool must be running.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000356 '''
Benjamin Petersone711caf2008-06-11 16:44:04 +0000357 return self.apply_async(func, args, kwds).get()
358
359 def map(self, func, iterable, chunksize=None):
360 '''
Georg Brandl92905032008-11-22 08:51:39 +0000361 Apply `func` to each element in `iterable`, collecting the results
362 in a list that is returned.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000363 '''
Antoine Pitroude911b22011-12-21 11:03:24 +0100364 return self._map_async(func, iterable, mapstar, chunksize).get()
365
366 def starmap(self, func, iterable, chunksize=None):
367 '''
368 Like `map()` method but the elements of the `iterable` are expected to
369 be iterables as well and will be unpacked as arguments. Hence
370 `func` and (a, b) becomes func(a, b).
371 '''
Antoine Pitroude911b22011-12-21 11:03:24 +0100372 return self._map_async(func, iterable, starmapstar, chunksize).get()
373
374 def starmap_async(self, func, iterable, chunksize=None, callback=None,
375 error_callback=None):
376 '''
377 Asynchronous version of `starmap()` method.
378 '''
Antoine Pitroude911b22011-12-21 11:03:24 +0100379 return self._map_async(func, iterable, starmapstar, chunksize,
380 callback, error_callback)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000381
Xiang Zhang794623b2017-03-29 11:58:54 +0800382 def _guarded_task_generation(self, result_job, func, iterable):
383 '''Provides a generator of tasks for imap and imap_unordered with
384 appropriate handling for iterables which throw exceptions during
385 iteration.'''
386 try:
387 i = -1
388 for i, x in enumerate(iterable):
389 yield (result_job, i, func, (x,), {})
390 except Exception as e:
391 yield (result_job, i+1, _helper_reraises_exception, (e,), {})
392
Benjamin Petersone711caf2008-06-11 16:44:04 +0000393 def imap(self, func, iterable, chunksize=1):
394 '''
Georg Brandl92905032008-11-22 08:51:39 +0000395 Equivalent of `map()` -- can be MUCH slower than `Pool.map()`.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000396 '''
Victor Stinner08c2ba02018-12-13 02:15:30 +0100397 self._check_running()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000398 if chunksize == 1:
Pablo Galindo3766f182019-02-11 17:29:00 +0000399 result = IMapIterator(self)
Xiang Zhang794623b2017-03-29 11:58:54 +0800400 self._taskqueue.put(
401 (
402 self._guarded_task_generation(result._job, func, iterable),
403 result._set_length
404 ))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000405 return result
406 else:
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500407 if chunksize < 1:
408 raise ValueError(
409 "Chunksize must be 1+, not {0:n}".format(
410 chunksize))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000411 task_batches = Pool._get_tasks(func, iterable, chunksize)
Pablo Galindo3766f182019-02-11 17:29:00 +0000412 result = IMapIterator(self)
Xiang Zhang794623b2017-03-29 11:58:54 +0800413 self._taskqueue.put(
414 (
415 self._guarded_task_generation(result._job,
416 mapstar,
417 task_batches),
418 result._set_length
419 ))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000420 return (item for chunk in result for item in chunk)
421
422 def imap_unordered(self, func, iterable, chunksize=1):
423 '''
Georg Brandl92905032008-11-22 08:51:39 +0000424 Like `imap()` method but ordering of results is arbitrary.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000425 '''
Victor Stinner08c2ba02018-12-13 02:15:30 +0100426 self._check_running()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000427 if chunksize == 1:
Pablo Galindo3766f182019-02-11 17:29:00 +0000428 result = IMapUnorderedIterator(self)
Xiang Zhang794623b2017-03-29 11:58:54 +0800429 self._taskqueue.put(
430 (
431 self._guarded_task_generation(result._job, func, iterable),
432 result._set_length
433 ))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000434 return result
435 else:
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500436 if chunksize < 1:
437 raise ValueError(
438 "Chunksize must be 1+, not {0!r}".format(chunksize))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000439 task_batches = Pool._get_tasks(func, iterable, chunksize)
Pablo Galindo3766f182019-02-11 17:29:00 +0000440 result = IMapUnorderedIterator(self)
Xiang Zhang794623b2017-03-29 11:58:54 +0800441 self._taskqueue.put(
442 (
443 self._guarded_task_generation(result._job,
444 mapstar,
445 task_batches),
446 result._set_length
447 ))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000448 return (item for chunk in result for item in chunk)
449
Ask Solem2afcbf22010-11-09 20:55:52 +0000450 def apply_async(self, func, args=(), kwds={}, callback=None,
451 error_callback=None):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000452 '''
Georg Brandl92905032008-11-22 08:51:39 +0000453 Asynchronous version of `apply()` method.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000454 '''
Victor Stinner08c2ba02018-12-13 02:15:30 +0100455 self._check_running()
Pablo Galindo3766f182019-02-11 17:29:00 +0000456 result = ApplyResult(self, callback, error_callback)
Xiang Zhang794623b2017-03-29 11:58:54 +0800457 self._taskqueue.put(([(result._job, 0, func, args, kwds)], None))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000458 return result
459
Ask Solem2afcbf22010-11-09 20:55:52 +0000460 def map_async(self, func, iterable, chunksize=None, callback=None,
461 error_callback=None):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000462 '''
Georg Brandl92905032008-11-22 08:51:39 +0000463 Asynchronous version of `map()` method.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000464 '''
Hynek Schlawack254af262012-10-27 12:53:02 +0200465 return self._map_async(func, iterable, mapstar, chunksize, callback,
466 error_callback)
Antoine Pitroude911b22011-12-21 11:03:24 +0100467
468 def _map_async(self, func, iterable, mapper, chunksize=None, callback=None,
469 error_callback=None):
470 '''
471 Helper function to implement map, starmap and their async counterparts.
472 '''
Victor Stinner08c2ba02018-12-13 02:15:30 +0100473 self._check_running()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000474 if not hasattr(iterable, '__len__'):
475 iterable = list(iterable)
476
477 if chunksize is None:
478 chunksize, extra = divmod(len(iterable), len(self._pool) * 4)
479 if extra:
480 chunksize += 1
Alexandre Vassalottie52e3782009-07-17 09:18:18 +0000481 if len(iterable) == 0:
482 chunksize = 0
Benjamin Petersone711caf2008-06-11 16:44:04 +0000483
484 task_batches = Pool._get_tasks(func, iterable, chunksize)
Pablo Galindo3766f182019-02-11 17:29:00 +0000485 result = MapResult(self, chunksize, len(iterable), callback,
Ask Solem2afcbf22010-11-09 20:55:52 +0000486 error_callback=error_callback)
Xiang Zhang794623b2017-03-29 11:58:54 +0800487 self._taskqueue.put(
488 (
489 self._guarded_task_generation(result._job,
490 mapper,
491 task_batches),
492 None
493 )
494 )
Benjamin Petersone711caf2008-06-11 16:44:04 +0000495 return result
496
497 @staticmethod
Pablo Galindo7c994542019-03-16 22:34:24 +0000498 def _wait_for_updates(sentinels, change_notifier, timeout=None):
499 wait(sentinels, timeout=timeout)
500 while not change_notifier.empty():
501 change_notifier.get()
502
503 @classmethod
504 def _handle_workers(cls, cache, taskqueue, ctx, Process, processes,
505 pool, inqueue, outqueue, initializer, initargs,
506 maxtasksperchild, wrap_exception, sentinels,
507 change_notifier):
Charles-François Natalif8859e12011-10-24 18:45:29 +0200508 thread = threading.current_thread()
509
510 # Keep maintaining workers until the cache gets drained, unless the pool
511 # is terminated.
Pablo Galindo3766f182019-02-11 17:29:00 +0000512 while thread._state == RUN or (cache and thread._state != TERMINATE):
Pablo Galindo7c994542019-03-16 22:34:24 +0000513 cls._maintain_pool(ctx, Process, processes, pool, inqueue,
514 outqueue, initializer, initargs,
515 maxtasksperchild, wrap_exception)
516
517 current_sentinels = [*cls._get_worker_sentinels(pool), *sentinels]
518
519 cls._wait_for_updates(current_sentinels, change_notifier)
Antoine Pitrou81dee6b2011-04-11 00:18:59 +0200520 # send sentinel to stop workers
Pablo Galindo3766f182019-02-11 17:29:00 +0000521 taskqueue.put(None)
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100522 util.debug('worker handler exiting')
Jesse Noller1f0b6582010-01-27 03:36:01 +0000523
524 @staticmethod
Richard Oudkerke90cedb2013-10-28 23:11:58 +0000525 def _handle_tasks(taskqueue, put, outqueue, pool, cache):
Benjamin Peterson672b8032008-06-11 19:14:14 +0000526 thread = threading.current_thread()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000527
528 for taskseq, set_length in iter(taskqueue.get, None):
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200529 task = None
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200530 try:
Xiang Zhang794623b2017-03-29 11:58:54 +0800531 # iterating taskseq cannot fail
532 for task in taskseq:
Victor Stinner2b417fb2018-12-14 11:13:18 +0100533 if thread._state != RUN:
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200534 util.debug('task handler found thread._state != RUN')
535 break
Richard Oudkerke90cedb2013-10-28 23:11:58 +0000536 try:
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200537 put(task)
538 except Exception as e:
Xiang Zhang794623b2017-03-29 11:58:54 +0800539 job, idx = task[:2]
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200540 try:
Xiang Zhang794623b2017-03-29 11:58:54 +0800541 cache[job]._set(idx, (False, e))
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200542 except KeyError:
543 pass
544 else:
545 if set_length:
546 util.debug('doing set_length()')
Xiang Zhang794623b2017-03-29 11:58:54 +0800547 idx = task[1] if task else -1
548 set_length(idx + 1)
Serhiy Storchaka79fbeee2015-03-13 08:25:26 +0200549 continue
550 break
Antoine Pitrou89889452017-03-24 13:52:11 +0100551 finally:
552 task = taskseq = job = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000553 else:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100554 util.debug('task handler got sentinel')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000555
Benjamin Petersone711caf2008-06-11 16:44:04 +0000556 try:
557 # tell result handler to finish when cache is empty
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100558 util.debug('task handler sending sentinel to result handler')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000559 outqueue.put(None)
560
561 # tell workers there is no more work
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100562 util.debug('task handler sending sentinel to workers')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000563 for p in pool:
564 put(None)
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200565 except OSError:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100566 util.debug('task handler got OSError when sending sentinels')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000567
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100568 util.debug('task handler exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000569
570 @staticmethod
571 def _handle_results(outqueue, get, cache):
Benjamin Peterson672b8032008-06-11 19:14:14 +0000572 thread = threading.current_thread()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000573
574 while 1:
575 try:
576 task = get()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200577 except (OSError, EOFError):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100578 util.debug('result handler got EOFError/OSError -- exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000579 return
580
Victor Stinner2dfe3512018-12-16 23:40:49 +0100581 if thread._state != RUN:
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500582 assert thread._state == TERMINATE, "Thread not in TERMINATE"
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100583 util.debug('result handler found thread._state=TERMINATE')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000584 break
585
586 if task is None:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100587 util.debug('result handler got sentinel')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000588 break
589
590 job, i, obj = task
591 try:
592 cache[job]._set(i, obj)
593 except KeyError:
594 pass
Antoine Pitrou89889452017-03-24 13:52:11 +0100595 task = job = obj = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000596
597 while cache and thread._state != TERMINATE:
598 try:
599 task = get()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200600 except (OSError, EOFError):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100601 util.debug('result handler got EOFError/OSError -- exiting')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000602 return
603
604 if task is None:
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100605 util.debug('result handler ignoring extra sentinel')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000606 continue
607 job, i, obj = task
608 try:
609 cache[job]._set(i, obj)
610 except KeyError:
611 pass
Antoine Pitrou89889452017-03-24 13:52:11 +0100612 task = job = obj = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000613
614 if hasattr(outqueue, '_reader'):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100615 util.debug('ensuring that outqueue is not full')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000616 # If we don't make room available in outqueue then
617 # attempts to add the sentinel (None) to outqueue may
618 # block. There is guaranteed to be no more than 2 sentinels.
619 try:
620 for i in range(10):
621 if not outqueue._reader.poll():
622 break
623 get()
Andrew Svetlovf7a17b42012-12-25 16:47:37 +0200624 except (OSError, EOFError):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000625 pass
626
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100627 util.debug('result handler exiting: len(cache)=%s, thread._state=%s',
Benjamin Petersone711caf2008-06-11 16:44:04 +0000628 len(cache), thread._state)
629
630 @staticmethod
631 def _get_tasks(func, it, size):
632 it = iter(it)
633 while 1:
634 x = tuple(itertools.islice(it, size))
635 if not x:
636 return
637 yield (func, x)
638
639 def __reduce__(self):
640 raise NotImplementedError(
641 'pool objects cannot be passed between processes or pickled'
642 )
643
644 def close(self):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100645 util.debug('closing pool')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000646 if self._state == RUN:
647 self._state = CLOSE
Jesse Noller1f0b6582010-01-27 03:36:01 +0000648 self._worker_handler._state = CLOSE
Pablo Galindo7c994542019-03-16 22:34:24 +0000649 self._change_notifier.put(None)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000650
651 def terminate(self):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100652 util.debug('terminating pool')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000653 self._state = TERMINATE
Jesse Noller1f0b6582010-01-27 03:36:01 +0000654 self._worker_handler._state = TERMINATE
Pablo Galindo7c994542019-03-16 22:34:24 +0000655 self._change_notifier.put(None)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000656 self._terminate()
657
658 def join(self):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100659 util.debug('joining pool')
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500660 if self._state == RUN:
661 raise ValueError("Pool is still running")
662 elif self._state not in (CLOSE, TERMINATE):
663 raise ValueError("In unknown state")
Jesse Noller1f0b6582010-01-27 03:36:01 +0000664 self._worker_handler.join()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000665 self._task_handler.join()
666 self._result_handler.join()
667 for p in self._pool:
668 p.join()
669
670 @staticmethod
671 def _help_stuff_finish(inqueue, task_handler, size):
672 # task_handler may be blocked trying to put items on inqueue
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100673 util.debug('removing tasks from inqueue until task handler finished')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000674 inqueue._rlock.acquire()
Benjamin Peterson672b8032008-06-11 19:14:14 +0000675 while task_handler.is_alive() and inqueue._reader.poll():
Benjamin Petersone711caf2008-06-11 16:44:04 +0000676 inqueue._reader.recv()
677 time.sleep(0)
678
679 @classmethod
Pablo Galindo7c994542019-03-16 22:34:24 +0000680 def _terminate_pool(cls, taskqueue, inqueue, outqueue, pool, change_notifier,
Jesse Noller1f0b6582010-01-27 03:36:01 +0000681 worker_handler, task_handler, result_handler, cache):
Benjamin Petersone711caf2008-06-11 16:44:04 +0000682 # this is guaranteed to only be called once
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100683 util.debug('finalizing pool')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000684
Jesse Noller1f0b6582010-01-27 03:36:01 +0000685 worker_handler._state = TERMINATE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000686 task_handler._state = TERMINATE
Benjamin Petersone711caf2008-06-11 16:44:04 +0000687
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100688 util.debug('helping task handler/workers to finish')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000689 cls._help_stuff_finish(inqueue, task_handler, len(pool))
690
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500691 if (not result_handler.is_alive()) and (len(cache) != 0):
692 raise AssertionError(
693 "Cannot have cache with result_hander not alive")
Benjamin Petersone711caf2008-06-11 16:44:04 +0000694
695 result_handler._state = TERMINATE
Pablo Galindo7c994542019-03-16 22:34:24 +0000696 change_notifier.put(None)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000697 outqueue.put(None) # sentinel
698
Antoine Pitrou81dee6b2011-04-11 00:18:59 +0200699 # We must wait for the worker handler to exit before terminating
700 # workers because we don't want workers to be restarted behind our back.
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100701 util.debug('joining worker handler')
Richard Oudkerkf29ec4b2012-06-18 15:54:57 +0100702 if threading.current_thread() is not worker_handler:
703 worker_handler.join()
Antoine Pitrou81dee6b2011-04-11 00:18:59 +0200704
Jesse Noller1f0b6582010-01-27 03:36:01 +0000705 # Terminate workers which haven't already finished.
Benjamin Petersone711caf2008-06-11 16:44:04 +0000706 if pool and hasattr(pool[0], 'terminate'):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100707 util.debug('terminating workers')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000708 for p in pool:
Jesse Noller1f0b6582010-01-27 03:36:01 +0000709 if p.exitcode is None:
710 p.terminate()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000711
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100712 util.debug('joining task handler')
Richard Oudkerkf29ec4b2012-06-18 15:54:57 +0100713 if threading.current_thread() is not task_handler:
714 task_handler.join()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000715
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100716 util.debug('joining result handler')
Richard Oudkerkf29ec4b2012-06-18 15:54:57 +0100717 if threading.current_thread() is not result_handler:
718 result_handler.join()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000719
720 if pool and hasattr(pool[0], 'terminate'):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100721 util.debug('joining pool workers')
Benjamin Petersone711caf2008-06-11 16:44:04 +0000722 for p in pool:
Florent Xicluna998171f2010-03-08 13:32:17 +0000723 if p.is_alive():
Jesse Noller1f0b6582010-01-27 03:36:01 +0000724 # worker has not yet exited
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100725 util.debug('cleaning up worker %d' % p.pid)
Florent Xicluna998171f2010-03-08 13:32:17 +0000726 p.join()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000727
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100728 def __enter__(self):
Victor Stinner08c2ba02018-12-13 02:15:30 +0100729 self._check_running()
Richard Oudkerkd69cfe82012-06-18 17:47:52 +0100730 return self
731
732 def __exit__(self, exc_type, exc_val, exc_tb):
733 self.terminate()
734
Benjamin Petersone711caf2008-06-11 16:44:04 +0000735#
736# Class whose instances are returned by `Pool.apply_async()`
737#
738
739class ApplyResult(object):
740
Pablo Galindo3766f182019-02-11 17:29:00 +0000741 def __init__(self, pool, callback, error_callback):
742 self._pool = pool
Richard Oudkerk692130a2012-05-25 13:26:53 +0100743 self._event = threading.Event()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000744 self._job = next(job_counter)
Pablo Galindo3766f182019-02-11 17:29:00 +0000745 self._cache = pool._cache
Benjamin Petersone711caf2008-06-11 16:44:04 +0000746 self._callback = callback
Ask Solem2afcbf22010-11-09 20:55:52 +0000747 self._error_callback = error_callback
Pablo Galindo3766f182019-02-11 17:29:00 +0000748 self._cache[self._job] = self
Benjamin Petersone711caf2008-06-11 16:44:04 +0000749
750 def ready(self):
Richard Oudkerk692130a2012-05-25 13:26:53 +0100751 return self._event.is_set()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000752
753 def successful(self):
Allen W. Smith, Ph.Dbd73e722017-08-29 17:52:18 -0500754 if not self.ready():
755 raise ValueError("{0!r} not ready".format(self))
Benjamin Petersone711caf2008-06-11 16:44:04 +0000756 return self._success
757
758 def wait(self, timeout=None):
Richard Oudkerk692130a2012-05-25 13:26:53 +0100759 self._event.wait(timeout)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000760
761 def get(self, timeout=None):
762 self.wait(timeout)
Richard Oudkerk692130a2012-05-25 13:26:53 +0100763 if not self.ready():
Benjamin Petersone711caf2008-06-11 16:44:04 +0000764 raise TimeoutError
765 if self._success:
766 return self._value
767 else:
768 raise self._value
769
770 def _set(self, i, obj):
771 self._success, self._value = obj
772 if self._callback and self._success:
773 self._callback(self._value)
Ask Solem2afcbf22010-11-09 20:55:52 +0000774 if self._error_callback and not self._success:
775 self._error_callback(self._value)
Richard Oudkerk692130a2012-05-25 13:26:53 +0100776 self._event.set()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000777 del self._cache[self._job]
Pablo Galindo3766f182019-02-11 17:29:00 +0000778 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000779
Richard Oudkerkdef51ca2013-05-06 12:10:04 +0100780AsyncResult = ApplyResult # create alias -- see #17805
781
Benjamin Petersone711caf2008-06-11 16:44:04 +0000782#
783# Class whose instances are returned by `Pool.map_async()`
784#
785
786class MapResult(ApplyResult):
787
Pablo Galindo3766f182019-02-11 17:29:00 +0000788 def __init__(self, pool, chunksize, length, callback, error_callback):
789 ApplyResult.__init__(self, pool, callback,
Ask Solem2afcbf22010-11-09 20:55:52 +0000790 error_callback=error_callback)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000791 self._success = True
792 self._value = [None] * length
793 self._chunksize = chunksize
794 if chunksize <= 0:
795 self._number_left = 0
Richard Oudkerk692130a2012-05-25 13:26:53 +0100796 self._event.set()
Pablo Galindo3766f182019-02-11 17:29:00 +0000797 del self._cache[self._job]
Benjamin Petersone711caf2008-06-11 16:44:04 +0000798 else:
799 self._number_left = length//chunksize + bool(length % chunksize)
800
801 def _set(self, i, success_result):
Charles-François Natali78f55ff2016-02-10 22:58:18 +0000802 self._number_left -= 1
Benjamin Petersone711caf2008-06-11 16:44:04 +0000803 success, result = success_result
Charles-François Natali78f55ff2016-02-10 22:58:18 +0000804 if success and self._success:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000805 self._value[i*self._chunksize:(i+1)*self._chunksize] = result
Benjamin Petersone711caf2008-06-11 16:44:04 +0000806 if self._number_left == 0:
807 if self._callback:
808 self._callback(self._value)
809 del self._cache[self._job]
Richard Oudkerk692130a2012-05-25 13:26:53 +0100810 self._event.set()
Pablo Galindo3766f182019-02-11 17:29:00 +0000811 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000812 else:
Charles-François Natali78f55ff2016-02-10 22:58:18 +0000813 if not success and self._success:
814 # only store first exception
815 self._success = False
816 self._value = result
817 if self._number_left == 0:
818 # only consider the result ready once all jobs are done
819 if self._error_callback:
820 self._error_callback(self._value)
821 del self._cache[self._job]
822 self._event.set()
Pablo Galindo3766f182019-02-11 17:29:00 +0000823 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000824
825#
826# Class whose instances are returned by `Pool.imap()`
827#
828
829class IMapIterator(object):
830
Pablo Galindo3766f182019-02-11 17:29:00 +0000831 def __init__(self, pool):
832 self._pool = pool
Benjamin Petersone711caf2008-06-11 16:44:04 +0000833 self._cond = threading.Condition(threading.Lock())
834 self._job = next(job_counter)
Pablo Galindo3766f182019-02-11 17:29:00 +0000835 self._cache = pool._cache
Benjamin Petersone711caf2008-06-11 16:44:04 +0000836 self._items = collections.deque()
837 self._index = 0
838 self._length = None
839 self._unsorted = {}
Pablo Galindo3766f182019-02-11 17:29:00 +0000840 self._cache[self._job] = self
Benjamin Petersone711caf2008-06-11 16:44:04 +0000841
842 def __iter__(self):
843 return self
844
845 def next(self, timeout=None):
Charles-François Natalia924fc72014-05-25 14:12:12 +0100846 with self._cond:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000847 try:
848 item = self._items.popleft()
849 except IndexError:
850 if self._index == self._length:
Pablo Galindo3766f182019-02-11 17:29:00 +0000851 self._pool = None
Serhiy Storchaka5affd232017-04-05 09:37:24 +0300852 raise StopIteration from None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000853 self._cond.wait(timeout)
854 try:
855 item = self._items.popleft()
856 except IndexError:
857 if self._index == self._length:
Pablo Galindo3766f182019-02-11 17:29:00 +0000858 self._pool = None
Serhiy Storchaka5affd232017-04-05 09:37:24 +0300859 raise StopIteration from None
860 raise TimeoutError from None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000861
862 success, value = item
863 if success:
864 return value
865 raise value
866
867 __next__ = next # XXX
868
869 def _set(self, i, obj):
Charles-François Natalia924fc72014-05-25 14:12:12 +0100870 with self._cond:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000871 if self._index == i:
872 self._items.append(obj)
873 self._index += 1
874 while self._index in self._unsorted:
875 obj = self._unsorted.pop(self._index)
876 self._items.append(obj)
877 self._index += 1
878 self._cond.notify()
879 else:
880 self._unsorted[i] = obj
881
882 if self._index == self._length:
883 del self._cache[self._job]
Pablo Galindo3766f182019-02-11 17:29:00 +0000884 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000885
886 def _set_length(self, length):
Charles-François Natalia924fc72014-05-25 14:12:12 +0100887 with self._cond:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000888 self._length = length
889 if self._index == self._length:
890 self._cond.notify()
891 del self._cache[self._job]
Pablo Galindo3766f182019-02-11 17:29:00 +0000892 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000893
894#
895# Class whose instances are returned by `Pool.imap_unordered()`
896#
897
898class IMapUnorderedIterator(IMapIterator):
899
900 def _set(self, i, obj):
Charles-François Natalia924fc72014-05-25 14:12:12 +0100901 with self._cond:
Benjamin Petersone711caf2008-06-11 16:44:04 +0000902 self._items.append(obj)
903 self._index += 1
904 self._cond.notify()
905 if self._index == self._length:
906 del self._cache[self._job]
Pablo Galindo3766f182019-02-11 17:29:00 +0000907 self._pool = None
Benjamin Petersone711caf2008-06-11 16:44:04 +0000908
909#
910#
911#
912
913class ThreadPool(Pool):
Richard Oudkerk80a5be12014-03-23 12:30:54 +0000914 _wrap_exception = False
Benjamin Petersone711caf2008-06-11 16:44:04 +0000915
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100916 @staticmethod
Pablo Galindo3766f182019-02-11 17:29:00 +0000917 def Process(ctx, *args, **kwds):
Richard Oudkerk84ed9a62013-08-14 15:35:41 +0100918 from .dummy import Process
919 return Process(*args, **kwds)
Benjamin Petersone711caf2008-06-11 16:44:04 +0000920
921 def __init__(self, processes=None, initializer=None, initargs=()):
922 Pool.__init__(self, processes, initializer, initargs)
923
924 def _setup_queues(self):
Antoine Pitrouab745042018-01-18 10:38:03 +0100925 self._inqueue = queue.SimpleQueue()
926 self._outqueue = queue.SimpleQueue()
Benjamin Petersone711caf2008-06-11 16:44:04 +0000927 self._quick_put = self._inqueue.put
928 self._quick_get = self._outqueue.get
929
Pablo Galindo7c994542019-03-16 22:34:24 +0000930 def _get_sentinels(self):
931 return [self._change_notifier._reader]
932
933 @staticmethod
934 def _get_worker_sentinels(workers):
935 return []
936
Benjamin Petersone711caf2008-06-11 16:44:04 +0000937 @staticmethod
938 def _help_stuff_finish(inqueue, task_handler, size):
Antoine Pitrouab745042018-01-18 10:38:03 +0100939 # drain inqueue, and put sentinels at its head to make workers finish
940 try:
941 while True:
942 inqueue.get(block=False)
943 except queue.Empty:
944 pass
945 for i in range(size):
946 inqueue.put(None)
Pablo Galindo7c994542019-03-16 22:34:24 +0000947
948 def _wait_for_updates(self, sentinels, change_notifier, timeout):
949 time.sleep(timeout)