|
jpayne@69
|
1 #
|
|
jpayne@69
|
2 # Module providing the `Pool` class for managing a process pool
|
|
jpayne@69
|
3 #
|
|
jpayne@69
|
4 # multiprocessing/pool.py
|
|
jpayne@69
|
5 #
|
|
jpayne@69
|
6 # Copyright (c) 2006-2008, R Oudkerk
|
|
jpayne@69
|
7 # Licensed to PSF under a Contributor Agreement.
|
|
jpayne@69
|
8 #
|
|
jpayne@69
|
9
|
|
jpayne@69
|
10 __all__ = ['Pool', 'ThreadPool']
|
|
jpayne@69
|
11
|
|
jpayne@69
|
12 #
|
|
jpayne@69
|
13 # Imports
|
|
jpayne@69
|
14 #
|
|
jpayne@69
|
15
|
|
jpayne@69
|
16 import collections
|
|
jpayne@69
|
17 import itertools
|
|
jpayne@69
|
18 import os
|
|
jpayne@69
|
19 import queue
|
|
jpayne@69
|
20 import threading
|
|
jpayne@69
|
21 import time
|
|
jpayne@69
|
22 import traceback
|
|
jpayne@69
|
23 import warnings
|
|
jpayne@69
|
24 from queue import Empty
|
|
jpayne@69
|
25
|
|
jpayne@69
|
26 # If threading is available then ThreadPool should be provided. Therefore
|
|
jpayne@69
|
27 # we avoid top-level imports which are liable to fail on some systems.
|
|
jpayne@69
|
28 from . import util
|
|
jpayne@69
|
29 from . import get_context, TimeoutError
|
|
jpayne@69
|
30 from .connection import wait
|
|
jpayne@69
|
31
|
|
jpayne@69
|
32 #
|
|
jpayne@69
|
33 # Constants representing the state of a pool
|
|
jpayne@69
|
34 #
|
|
jpayne@69
|
35
|
|
jpayne@69
|
36 INIT = "INIT"
|
|
jpayne@69
|
37 RUN = "RUN"
|
|
jpayne@69
|
38 CLOSE = "CLOSE"
|
|
jpayne@69
|
39 TERMINATE = "TERMINATE"
|
|
jpayne@69
|
40
|
|
jpayne@69
|
41 #
|
|
jpayne@69
|
42 # Miscellaneous
|
|
jpayne@69
|
43 #
|
|
jpayne@69
|
44
|
|
jpayne@69
|
45 job_counter = itertools.count()
|
|
jpayne@69
|
46
|
|
jpayne@69
|
47 def mapstar(args):
|
|
jpayne@69
|
48 return list(map(*args))
|
|
jpayne@69
|
49
|
|
jpayne@69
|
50 def starmapstar(args):
|
|
jpayne@69
|
51 return list(itertools.starmap(args[0], args[1]))
|
|
jpayne@69
|
52
|
|
jpayne@69
|
53 #
|
|
jpayne@69
|
54 # Hack to embed stringification of remote traceback in local traceback
|
|
jpayne@69
|
55 #
|
|
jpayne@69
|
56
|
|
jpayne@69
|
57 class RemoteTraceback(Exception):
|
|
jpayne@69
|
58 def __init__(self, tb):
|
|
jpayne@69
|
59 self.tb = tb
|
|
jpayne@69
|
60 def __str__(self):
|
|
jpayne@69
|
61 return self.tb
|
|
jpayne@69
|
62
|
|
jpayne@69
|
63 class ExceptionWithTraceback:
|
|
jpayne@69
|
64 def __init__(self, exc, tb):
|
|
jpayne@69
|
65 tb = traceback.format_exception(type(exc), exc, tb)
|
|
jpayne@69
|
66 tb = ''.join(tb)
|
|
jpayne@69
|
67 self.exc = exc
|
|
jpayne@69
|
68 self.tb = '\n"""\n%s"""' % tb
|
|
jpayne@69
|
69 def __reduce__(self):
|
|
jpayne@69
|
70 return rebuild_exc, (self.exc, self.tb)
|
|
jpayne@69
|
71
|
|
jpayne@69
|
72 def rebuild_exc(exc, tb):
|
|
jpayne@69
|
73 exc.__cause__ = RemoteTraceback(tb)
|
|
jpayne@69
|
74 return exc
|
|
jpayne@69
|
75
|
|
jpayne@69
|
76 #
|
|
jpayne@69
|
77 # Code run by worker processes
|
|
jpayne@69
|
78 #
|
|
jpayne@69
|
79
|
|
jpayne@69
|
80 class MaybeEncodingError(Exception):
|
|
jpayne@69
|
81 """Wraps possible unpickleable errors, so they can be
|
|
jpayne@69
|
82 safely sent through the socket."""
|
|
jpayne@69
|
83
|
|
jpayne@69
|
84 def __init__(self, exc, value):
|
|
jpayne@69
|
85 self.exc = repr(exc)
|
|
jpayne@69
|
86 self.value = repr(value)
|
|
jpayne@69
|
87 super(MaybeEncodingError, self).__init__(self.exc, self.value)
|
|
jpayne@69
|
88
|
|
jpayne@69
|
89 def __str__(self):
|
|
jpayne@69
|
90 return "Error sending result: '%s'. Reason: '%s'" % (self.value,
|
|
jpayne@69
|
91 self.exc)
|
|
jpayne@69
|
92
|
|
jpayne@69
|
93 def __repr__(self):
|
|
jpayne@69
|
94 return "<%s: %s>" % (self.__class__.__name__, self)
|
|
jpayne@69
|
95
|
|
jpayne@69
|
96
|
|
jpayne@69
|
97 def worker(inqueue, outqueue, initializer=None, initargs=(), maxtasks=None,
|
|
jpayne@69
|
98 wrap_exception=False):
|
|
jpayne@69
|
99 if (maxtasks is not None) and not (isinstance(maxtasks, int)
|
|
jpayne@69
|
100 and maxtasks >= 1):
|
|
jpayne@69
|
101 raise AssertionError("Maxtasks {!r} is not valid".format(maxtasks))
|
|
jpayne@69
|
102 put = outqueue.put
|
|
jpayne@69
|
103 get = inqueue.get
|
|
jpayne@69
|
104 if hasattr(inqueue, '_writer'):
|
|
jpayne@69
|
105 inqueue._writer.close()
|
|
jpayne@69
|
106 outqueue._reader.close()
|
|
jpayne@69
|
107
|
|
jpayne@69
|
108 if initializer is not None:
|
|
jpayne@69
|
109 initializer(*initargs)
|
|
jpayne@69
|
110
|
|
jpayne@69
|
111 completed = 0
|
|
jpayne@69
|
112 while maxtasks is None or (maxtasks and completed < maxtasks):
|
|
jpayne@69
|
113 try:
|
|
jpayne@69
|
114 task = get()
|
|
jpayne@69
|
115 except (EOFError, OSError):
|
|
jpayne@69
|
116 util.debug('worker got EOFError or OSError -- exiting')
|
|
jpayne@69
|
117 break
|
|
jpayne@69
|
118
|
|
jpayne@69
|
119 if task is None:
|
|
jpayne@69
|
120 util.debug('worker got sentinel -- exiting')
|
|
jpayne@69
|
121 break
|
|
jpayne@69
|
122
|
|
jpayne@69
|
123 job, i, func, args, kwds = task
|
|
jpayne@69
|
124 try:
|
|
jpayne@69
|
125 result = (True, func(*args, **kwds))
|
|
jpayne@69
|
126 except Exception as e:
|
|
jpayne@69
|
127 if wrap_exception and func is not _helper_reraises_exception:
|
|
jpayne@69
|
128 e = ExceptionWithTraceback(e, e.__traceback__)
|
|
jpayne@69
|
129 result = (False, e)
|
|
jpayne@69
|
130 try:
|
|
jpayne@69
|
131 put((job, i, result))
|
|
jpayne@69
|
132 except Exception as e:
|
|
jpayne@69
|
133 wrapped = MaybeEncodingError(e, result[1])
|
|
jpayne@69
|
134 util.debug("Possible encoding error while sending result: %s" % (
|
|
jpayne@69
|
135 wrapped))
|
|
jpayne@69
|
136 put((job, i, (False, wrapped)))
|
|
jpayne@69
|
137
|
|
jpayne@69
|
138 task = job = result = func = args = kwds = None
|
|
jpayne@69
|
139 completed += 1
|
|
jpayne@69
|
140 util.debug('worker exiting after %d tasks' % completed)
|
|
jpayne@69
|
141
|
|
jpayne@69
|
142 def _helper_reraises_exception(ex):
|
|
jpayne@69
|
143 'Pickle-able helper function for use by _guarded_task_generation.'
|
|
jpayne@69
|
144 raise ex
|
|
jpayne@69
|
145
|
|
jpayne@69
|
146 #
|
|
jpayne@69
|
147 # Class representing a process pool
|
|
jpayne@69
|
148 #
|
|
jpayne@69
|
149
|
|
jpayne@69
|
150 class _PoolCache(dict):
|
|
jpayne@69
|
151 """
|
|
jpayne@69
|
152 Class that implements a cache for the Pool class that will notify
|
|
jpayne@69
|
153 the pool management threads every time the cache is emptied. The
|
|
jpayne@69
|
154 notification is done by the use of a queue that is provided when
|
|
jpayne@69
|
155 instantiating the cache.
|
|
jpayne@69
|
156 """
|
|
jpayne@69
|
157 def __init__(self, /, *args, notifier=None, **kwds):
|
|
jpayne@69
|
158 self.notifier = notifier
|
|
jpayne@69
|
159 super().__init__(*args, **kwds)
|
|
jpayne@69
|
160
|
|
jpayne@69
|
161 def __delitem__(self, item):
|
|
jpayne@69
|
162 super().__delitem__(item)
|
|
jpayne@69
|
163
|
|
jpayne@69
|
164 # Notify that the cache is empty. This is important because the
|
|
jpayne@69
|
165 # pool keeps maintaining workers until the cache gets drained. This
|
|
jpayne@69
|
166 # eliminates a race condition in which a task is finished after the
|
|
jpayne@69
|
167 # the pool's _handle_workers method has enter another iteration of the
|
|
jpayne@69
|
168 # loop. In this situation, the only event that can wake up the pool
|
|
jpayne@69
|
169 # is the cache to be emptied (no more tasks available).
|
|
jpayne@69
|
170 if not self:
|
|
jpayne@69
|
171 self.notifier.put(None)
|
|
jpayne@69
|
172
|
|
jpayne@69
|
173 class Pool(object):
|
|
jpayne@69
|
174 '''
|
|
jpayne@69
|
175 Class which supports an async version of applying functions to arguments.
|
|
jpayne@69
|
176 '''
|
|
jpayne@69
|
177 _wrap_exception = True
|
|
jpayne@69
|
178
|
|
jpayne@69
|
179 @staticmethod
|
|
jpayne@69
|
180 def Process(ctx, *args, **kwds):
|
|
jpayne@69
|
181 return ctx.Process(*args, **kwds)
|
|
jpayne@69
|
182
|
|
jpayne@69
|
183 def __init__(self, processes=None, initializer=None, initargs=(),
|
|
jpayne@69
|
184 maxtasksperchild=None, context=None):
|
|
jpayne@69
|
185 # Attributes initialized early to make sure that they exist in
|
|
jpayne@69
|
186 # __del__() if __init__() raises an exception
|
|
jpayne@69
|
187 self._pool = []
|
|
jpayne@69
|
188 self._state = INIT
|
|
jpayne@69
|
189
|
|
jpayne@69
|
190 self._ctx = context or get_context()
|
|
jpayne@69
|
191 self._setup_queues()
|
|
jpayne@69
|
192 self._taskqueue = queue.SimpleQueue()
|
|
jpayne@69
|
193 # The _change_notifier queue exist to wake up self._handle_workers()
|
|
jpayne@69
|
194 # when the cache (self._cache) is empty or when there is a change in
|
|
jpayne@69
|
195 # the _state variable of the thread that runs _handle_workers.
|
|
jpayne@69
|
196 self._change_notifier = self._ctx.SimpleQueue()
|
|
jpayne@69
|
197 self._cache = _PoolCache(notifier=self._change_notifier)
|
|
jpayne@69
|
198 self._maxtasksperchild = maxtasksperchild
|
|
jpayne@69
|
199 self._initializer = initializer
|
|
jpayne@69
|
200 self._initargs = initargs
|
|
jpayne@69
|
201
|
|
jpayne@69
|
202 if processes is None:
|
|
jpayne@69
|
203 processes = os.cpu_count() or 1
|
|
jpayne@69
|
204 if processes < 1:
|
|
jpayne@69
|
205 raise ValueError("Number of processes must be at least 1")
|
|
jpayne@69
|
206
|
|
jpayne@69
|
207 if initializer is not None and not callable(initializer):
|
|
jpayne@69
|
208 raise TypeError('initializer must be a callable')
|
|
jpayne@69
|
209
|
|
jpayne@69
|
210 self._processes = processes
|
|
jpayne@69
|
211 try:
|
|
jpayne@69
|
212 self._repopulate_pool()
|
|
jpayne@69
|
213 except Exception:
|
|
jpayne@69
|
214 for p in self._pool:
|
|
jpayne@69
|
215 if p.exitcode is None:
|
|
jpayne@69
|
216 p.terminate()
|
|
jpayne@69
|
217 for p in self._pool:
|
|
jpayne@69
|
218 p.join()
|
|
jpayne@69
|
219 raise
|
|
jpayne@69
|
220
|
|
jpayne@69
|
221 sentinels = self._get_sentinels()
|
|
jpayne@69
|
222
|
|
jpayne@69
|
223 self._worker_handler = threading.Thread(
|
|
jpayne@69
|
224 target=Pool._handle_workers,
|
|
jpayne@69
|
225 args=(self._cache, self._taskqueue, self._ctx, self.Process,
|
|
jpayne@69
|
226 self._processes, self._pool, self._inqueue, self._outqueue,
|
|
jpayne@69
|
227 self._initializer, self._initargs, self._maxtasksperchild,
|
|
jpayne@69
|
228 self._wrap_exception, sentinels, self._change_notifier)
|
|
jpayne@69
|
229 )
|
|
jpayne@69
|
230 self._worker_handler.daemon = True
|
|
jpayne@69
|
231 self._worker_handler._state = RUN
|
|
jpayne@69
|
232 self._worker_handler.start()
|
|
jpayne@69
|
233
|
|
jpayne@69
|
234
|
|
jpayne@69
|
235 self._task_handler = threading.Thread(
|
|
jpayne@69
|
236 target=Pool._handle_tasks,
|
|
jpayne@69
|
237 args=(self._taskqueue, self._quick_put, self._outqueue,
|
|
jpayne@69
|
238 self._pool, self._cache)
|
|
jpayne@69
|
239 )
|
|
jpayne@69
|
240 self._task_handler.daemon = True
|
|
jpayne@69
|
241 self._task_handler._state = RUN
|
|
jpayne@69
|
242 self._task_handler.start()
|
|
jpayne@69
|
243
|
|
jpayne@69
|
244 self._result_handler = threading.Thread(
|
|
jpayne@69
|
245 target=Pool._handle_results,
|
|
jpayne@69
|
246 args=(self._outqueue, self._quick_get, self._cache)
|
|
jpayne@69
|
247 )
|
|
jpayne@69
|
248 self._result_handler.daemon = True
|
|
jpayne@69
|
249 self._result_handler._state = RUN
|
|
jpayne@69
|
250 self._result_handler.start()
|
|
jpayne@69
|
251
|
|
jpayne@69
|
252 self._terminate = util.Finalize(
|
|
jpayne@69
|
253 self, self._terminate_pool,
|
|
jpayne@69
|
254 args=(self._taskqueue, self._inqueue, self._outqueue, self._pool,
|
|
jpayne@69
|
255 self._change_notifier, self._worker_handler, self._task_handler,
|
|
jpayne@69
|
256 self._result_handler, self._cache),
|
|
jpayne@69
|
257 exitpriority=15
|
|
jpayne@69
|
258 )
|
|
jpayne@69
|
259 self._state = RUN
|
|
jpayne@69
|
260
|
|
jpayne@69
|
261 # Copy globals as function locals to make sure that they are available
|
|
jpayne@69
|
262 # during Python shutdown when the Pool is destroyed.
|
|
jpayne@69
|
263 def __del__(self, _warn=warnings.warn, RUN=RUN):
|
|
jpayne@69
|
264 if self._state == RUN:
|
|
jpayne@69
|
265 _warn(f"unclosed running multiprocessing pool {self!r}",
|
|
jpayne@69
|
266 ResourceWarning, source=self)
|
|
jpayne@69
|
267 if getattr(self, '_change_notifier', None) is not None:
|
|
jpayne@69
|
268 self._change_notifier.put(None)
|
|
jpayne@69
|
269
|
|
jpayne@69
|
270 def __repr__(self):
|
|
jpayne@69
|
271 cls = self.__class__
|
|
jpayne@69
|
272 return (f'<{cls.__module__}.{cls.__qualname__} '
|
|
jpayne@69
|
273 f'state={self._state} '
|
|
jpayne@69
|
274 f'pool_size={len(self._pool)}>')
|
|
jpayne@69
|
275
|
|
jpayne@69
|
276 def _get_sentinels(self):
|
|
jpayne@69
|
277 task_queue_sentinels = [self._outqueue._reader]
|
|
jpayne@69
|
278 self_notifier_sentinels = [self._change_notifier._reader]
|
|
jpayne@69
|
279 return [*task_queue_sentinels, *self_notifier_sentinels]
|
|
jpayne@69
|
280
|
|
jpayne@69
|
281 @staticmethod
|
|
jpayne@69
|
282 def _get_worker_sentinels(workers):
|
|
jpayne@69
|
283 return [worker.sentinel for worker in
|
|
jpayne@69
|
284 workers if hasattr(worker, "sentinel")]
|
|
jpayne@69
|
285
|
|
jpayne@69
|
286 @staticmethod
|
|
jpayne@69
|
287 def _join_exited_workers(pool):
|
|
jpayne@69
|
288 """Cleanup after any worker processes which have exited due to reaching
|
|
jpayne@69
|
289 their specified lifetime. Returns True if any workers were cleaned up.
|
|
jpayne@69
|
290 """
|
|
jpayne@69
|
291 cleaned = False
|
|
jpayne@69
|
292 for i in reversed(range(len(pool))):
|
|
jpayne@69
|
293 worker = pool[i]
|
|
jpayne@69
|
294 if worker.exitcode is not None:
|
|
jpayne@69
|
295 # worker exited
|
|
jpayne@69
|
296 util.debug('cleaning up worker %d' % i)
|
|
jpayne@69
|
297 worker.join()
|
|
jpayne@69
|
298 cleaned = True
|
|
jpayne@69
|
299 del pool[i]
|
|
jpayne@69
|
300 return cleaned
|
|
jpayne@69
|
301
|
|
jpayne@69
|
302 def _repopulate_pool(self):
|
|
jpayne@69
|
303 return self._repopulate_pool_static(self._ctx, self.Process,
|
|
jpayne@69
|
304 self._processes,
|
|
jpayne@69
|
305 self._pool, self._inqueue,
|
|
jpayne@69
|
306 self._outqueue, self._initializer,
|
|
jpayne@69
|
307 self._initargs,
|
|
jpayne@69
|
308 self._maxtasksperchild,
|
|
jpayne@69
|
309 self._wrap_exception)
|
|
jpayne@69
|
310
|
|
jpayne@69
|
311 @staticmethod
|
|
jpayne@69
|
312 def _repopulate_pool_static(ctx, Process, processes, pool, inqueue,
|
|
jpayne@69
|
313 outqueue, initializer, initargs,
|
|
jpayne@69
|
314 maxtasksperchild, wrap_exception):
|
|
jpayne@69
|
315 """Bring the number of pool processes up to the specified number,
|
|
jpayne@69
|
316 for use after reaping workers which have exited.
|
|
jpayne@69
|
317 """
|
|
jpayne@69
|
318 for i in range(processes - len(pool)):
|
|
jpayne@69
|
319 w = Process(ctx, target=worker,
|
|
jpayne@69
|
320 args=(inqueue, outqueue,
|
|
jpayne@69
|
321 initializer,
|
|
jpayne@69
|
322 initargs, maxtasksperchild,
|
|
jpayne@69
|
323 wrap_exception))
|
|
jpayne@69
|
324 w.name = w.name.replace('Process', 'PoolWorker')
|
|
jpayne@69
|
325 w.daemon = True
|
|
jpayne@69
|
326 w.start()
|
|
jpayne@69
|
327 pool.append(w)
|
|
jpayne@69
|
328 util.debug('added worker')
|
|
jpayne@69
|
329
|
|
jpayne@69
|
330 @staticmethod
|
|
jpayne@69
|
331 def _maintain_pool(ctx, Process, processes, pool, inqueue, outqueue,
|
|
jpayne@69
|
332 initializer, initargs, maxtasksperchild,
|
|
jpayne@69
|
333 wrap_exception):
|
|
jpayne@69
|
334 """Clean up any exited workers and start replacements for them.
|
|
jpayne@69
|
335 """
|
|
jpayne@69
|
336 if Pool._join_exited_workers(pool):
|
|
jpayne@69
|
337 Pool._repopulate_pool_static(ctx, Process, processes, pool,
|
|
jpayne@69
|
338 inqueue, outqueue, initializer,
|
|
jpayne@69
|
339 initargs, maxtasksperchild,
|
|
jpayne@69
|
340 wrap_exception)
|
|
jpayne@69
|
341
|
|
jpayne@69
|
342 def _setup_queues(self):
|
|
jpayne@69
|
343 self._inqueue = self._ctx.SimpleQueue()
|
|
jpayne@69
|
344 self._outqueue = self._ctx.SimpleQueue()
|
|
jpayne@69
|
345 self._quick_put = self._inqueue._writer.send
|
|
jpayne@69
|
346 self._quick_get = self._outqueue._reader.recv
|
|
jpayne@69
|
347
|
|
jpayne@69
|
348 def _check_running(self):
|
|
jpayne@69
|
349 if self._state != RUN:
|
|
jpayne@69
|
350 raise ValueError("Pool not running")
|
|
jpayne@69
|
351
|
|
jpayne@69
|
352 def apply(self, func, args=(), kwds={}):
|
|
jpayne@69
|
353 '''
|
|
jpayne@69
|
354 Equivalent of `func(*args, **kwds)`.
|
|
jpayne@69
|
355 Pool must be running.
|
|
jpayne@69
|
356 '''
|
|
jpayne@69
|
357 return self.apply_async(func, args, kwds).get()
|
|
jpayne@69
|
358
|
|
jpayne@69
|
359 def map(self, func, iterable, chunksize=None):
|
|
jpayne@69
|
360 '''
|
|
jpayne@69
|
361 Apply `func` to each element in `iterable`, collecting the results
|
|
jpayne@69
|
362 in a list that is returned.
|
|
jpayne@69
|
363 '''
|
|
jpayne@69
|
364 return self._map_async(func, iterable, mapstar, chunksize).get()
|
|
jpayne@69
|
365
|
|
jpayne@69
|
366 def starmap(self, func, iterable, chunksize=None):
|
|
jpayne@69
|
367 '''
|
|
jpayne@69
|
368 Like `map()` method but the elements of the `iterable` are expected to
|
|
jpayne@69
|
369 be iterables as well and will be unpacked as arguments. Hence
|
|
jpayne@69
|
370 `func` and (a, b) becomes func(a, b).
|
|
jpayne@69
|
371 '''
|
|
jpayne@69
|
372 return self._map_async(func, iterable, starmapstar, chunksize).get()
|
|
jpayne@69
|
373
|
|
jpayne@69
|
374 def starmap_async(self, func, iterable, chunksize=None, callback=None,
|
|
jpayne@69
|
375 error_callback=None):
|
|
jpayne@69
|
376 '''
|
|
jpayne@69
|
377 Asynchronous version of `starmap()` method.
|
|
jpayne@69
|
378 '''
|
|
jpayne@69
|
379 return self._map_async(func, iterable, starmapstar, chunksize,
|
|
jpayne@69
|
380 callback, error_callback)
|
|
jpayne@69
|
381
|
|
jpayne@69
|
382 def _guarded_task_generation(self, result_job, func, iterable):
|
|
jpayne@69
|
383 '''Provides a generator of tasks for imap and imap_unordered with
|
|
jpayne@69
|
384 appropriate handling for iterables which throw exceptions during
|
|
jpayne@69
|
385 iteration.'''
|
|
jpayne@69
|
386 try:
|
|
jpayne@69
|
387 i = -1
|
|
jpayne@69
|
388 for i, x in enumerate(iterable):
|
|
jpayne@69
|
389 yield (result_job, i, func, (x,), {})
|
|
jpayne@69
|
390 except Exception as e:
|
|
jpayne@69
|
391 yield (result_job, i+1, _helper_reraises_exception, (e,), {})
|
|
jpayne@69
|
392
|
|
jpayne@69
|
393 def imap(self, func, iterable, chunksize=1):
|
|
jpayne@69
|
394 '''
|
|
jpayne@69
|
395 Equivalent of `map()` -- can be MUCH slower than `Pool.map()`.
|
|
jpayne@69
|
396 '''
|
|
jpayne@69
|
397 self._check_running()
|
|
jpayne@69
|
398 if chunksize == 1:
|
|
jpayne@69
|
399 result = IMapIterator(self)
|
|
jpayne@69
|
400 self._taskqueue.put(
|
|
jpayne@69
|
401 (
|
|
jpayne@69
|
402 self._guarded_task_generation(result._job, func, iterable),
|
|
jpayne@69
|
403 result._set_length
|
|
jpayne@69
|
404 ))
|
|
jpayne@69
|
405 return result
|
|
jpayne@69
|
406 else:
|
|
jpayne@69
|
407 if chunksize < 1:
|
|
jpayne@69
|
408 raise ValueError(
|
|
jpayne@69
|
409 "Chunksize must be 1+, not {0:n}".format(
|
|
jpayne@69
|
410 chunksize))
|
|
jpayne@69
|
411 task_batches = Pool._get_tasks(func, iterable, chunksize)
|
|
jpayne@69
|
412 result = IMapIterator(self)
|
|
jpayne@69
|
413 self._taskqueue.put(
|
|
jpayne@69
|
414 (
|
|
jpayne@69
|
415 self._guarded_task_generation(result._job,
|
|
jpayne@69
|
416 mapstar,
|
|
jpayne@69
|
417 task_batches),
|
|
jpayne@69
|
418 result._set_length
|
|
jpayne@69
|
419 ))
|
|
jpayne@69
|
420 return (item for chunk in result for item in chunk)
|
|
jpayne@69
|
421
|
|
jpayne@69
|
422 def imap_unordered(self, func, iterable, chunksize=1):
|
|
jpayne@69
|
423 '''
|
|
jpayne@69
|
424 Like `imap()` method but ordering of results is arbitrary.
|
|
jpayne@69
|
425 '''
|
|
jpayne@69
|
426 self._check_running()
|
|
jpayne@69
|
427 if chunksize == 1:
|
|
jpayne@69
|
428 result = IMapUnorderedIterator(self)
|
|
jpayne@69
|
429 self._taskqueue.put(
|
|
jpayne@69
|
430 (
|
|
jpayne@69
|
431 self._guarded_task_generation(result._job, func, iterable),
|
|
jpayne@69
|
432 result._set_length
|
|
jpayne@69
|
433 ))
|
|
jpayne@69
|
434 return result
|
|
jpayne@69
|
435 else:
|
|
jpayne@69
|
436 if chunksize < 1:
|
|
jpayne@69
|
437 raise ValueError(
|
|
jpayne@69
|
438 "Chunksize must be 1+, not {0!r}".format(chunksize))
|
|
jpayne@69
|
439 task_batches = Pool._get_tasks(func, iterable, chunksize)
|
|
jpayne@69
|
440 result = IMapUnorderedIterator(self)
|
|
jpayne@69
|
441 self._taskqueue.put(
|
|
jpayne@69
|
442 (
|
|
jpayne@69
|
443 self._guarded_task_generation(result._job,
|
|
jpayne@69
|
444 mapstar,
|
|
jpayne@69
|
445 task_batches),
|
|
jpayne@69
|
446 result._set_length
|
|
jpayne@69
|
447 ))
|
|
jpayne@69
|
448 return (item for chunk in result for item in chunk)
|
|
jpayne@69
|
449
|
|
jpayne@69
|
450 def apply_async(self, func, args=(), kwds={}, callback=None,
|
|
jpayne@69
|
451 error_callback=None):
|
|
jpayne@69
|
452 '''
|
|
jpayne@69
|
453 Asynchronous version of `apply()` method.
|
|
jpayne@69
|
454 '''
|
|
jpayne@69
|
455 self._check_running()
|
|
jpayne@69
|
456 result = ApplyResult(self, callback, error_callback)
|
|
jpayne@69
|
457 self._taskqueue.put(([(result._job, 0, func, args, kwds)], None))
|
|
jpayne@69
|
458 return result
|
|
jpayne@69
|
459
|
|
jpayne@69
|
460 def map_async(self, func, iterable, chunksize=None, callback=None,
|
|
jpayne@69
|
461 error_callback=None):
|
|
jpayne@69
|
462 '''
|
|
jpayne@69
|
463 Asynchronous version of `map()` method.
|
|
jpayne@69
|
464 '''
|
|
jpayne@69
|
465 return self._map_async(func, iterable, mapstar, chunksize, callback,
|
|
jpayne@69
|
466 error_callback)
|
|
jpayne@69
|
467
|
|
jpayne@69
|
468 def _map_async(self, func, iterable, mapper, chunksize=None, callback=None,
|
|
jpayne@69
|
469 error_callback=None):
|
|
jpayne@69
|
470 '''
|
|
jpayne@69
|
471 Helper function to implement map, starmap and their async counterparts.
|
|
jpayne@69
|
472 '''
|
|
jpayne@69
|
473 self._check_running()
|
|
jpayne@69
|
474 if not hasattr(iterable, '__len__'):
|
|
jpayne@69
|
475 iterable = list(iterable)
|
|
jpayne@69
|
476
|
|
jpayne@69
|
477 if chunksize is None:
|
|
jpayne@69
|
478 chunksize, extra = divmod(len(iterable), len(self._pool) * 4)
|
|
jpayne@69
|
479 if extra:
|
|
jpayne@69
|
480 chunksize += 1
|
|
jpayne@69
|
481 if len(iterable) == 0:
|
|
jpayne@69
|
482 chunksize = 0
|
|
jpayne@69
|
483
|
|
jpayne@69
|
484 task_batches = Pool._get_tasks(func, iterable, chunksize)
|
|
jpayne@69
|
485 result = MapResult(self, chunksize, len(iterable), callback,
|
|
jpayne@69
|
486 error_callback=error_callback)
|
|
jpayne@69
|
487 self._taskqueue.put(
|
|
jpayne@69
|
488 (
|
|
jpayne@69
|
489 self._guarded_task_generation(result._job,
|
|
jpayne@69
|
490 mapper,
|
|
jpayne@69
|
491 task_batches),
|
|
jpayne@69
|
492 None
|
|
jpayne@69
|
493 )
|
|
jpayne@69
|
494 )
|
|
jpayne@69
|
495 return result
|
|
jpayne@69
|
496
|
|
jpayne@69
|
497 @staticmethod
|
|
jpayne@69
|
498 def _wait_for_updates(sentinels, change_notifier, timeout=None):
|
|
jpayne@69
|
499 wait(sentinels, timeout=timeout)
|
|
jpayne@69
|
500 while not change_notifier.empty():
|
|
jpayne@69
|
501 change_notifier.get()
|
|
jpayne@69
|
502
|
|
jpayne@69
|
503 @classmethod
|
|
jpayne@69
|
504 def _handle_workers(cls, cache, taskqueue, ctx, Process, processes,
|
|
jpayne@69
|
505 pool, inqueue, outqueue, initializer, initargs,
|
|
jpayne@69
|
506 maxtasksperchild, wrap_exception, sentinels,
|
|
jpayne@69
|
507 change_notifier):
|
|
jpayne@69
|
508 thread = threading.current_thread()
|
|
jpayne@69
|
509
|
|
jpayne@69
|
510 # Keep maintaining workers until the cache gets drained, unless the pool
|
|
jpayne@69
|
511 # is terminated.
|
|
jpayne@69
|
512 while thread._state == RUN or (cache and thread._state != TERMINATE):
|
|
jpayne@69
|
513 cls._maintain_pool(ctx, Process, processes, pool, inqueue,
|
|
jpayne@69
|
514 outqueue, initializer, initargs,
|
|
jpayne@69
|
515 maxtasksperchild, wrap_exception)
|
|
jpayne@69
|
516
|
|
jpayne@69
|
517 current_sentinels = [*cls._get_worker_sentinels(pool), *sentinels]
|
|
jpayne@69
|
518
|
|
jpayne@69
|
519 cls._wait_for_updates(current_sentinels, change_notifier)
|
|
jpayne@69
|
520 # send sentinel to stop workers
|
|
jpayne@69
|
521 taskqueue.put(None)
|
|
jpayne@69
|
522 util.debug('worker handler exiting')
|
|
jpayne@69
|
523
|
|
jpayne@69
|
524 @staticmethod
|
|
jpayne@69
|
525 def _handle_tasks(taskqueue, put, outqueue, pool, cache):
|
|
jpayne@69
|
526 thread = threading.current_thread()
|
|
jpayne@69
|
527
|
|
jpayne@69
|
528 for taskseq, set_length in iter(taskqueue.get, None):
|
|
jpayne@69
|
529 task = None
|
|
jpayne@69
|
530 try:
|
|
jpayne@69
|
531 # iterating taskseq cannot fail
|
|
jpayne@69
|
532 for task in taskseq:
|
|
jpayne@69
|
533 if thread._state != RUN:
|
|
jpayne@69
|
534 util.debug('task handler found thread._state != RUN')
|
|
jpayne@69
|
535 break
|
|
jpayne@69
|
536 try:
|
|
jpayne@69
|
537 put(task)
|
|
jpayne@69
|
538 except Exception as e:
|
|
jpayne@69
|
539 job, idx = task[:2]
|
|
jpayne@69
|
540 try:
|
|
jpayne@69
|
541 cache[job]._set(idx, (False, e))
|
|
jpayne@69
|
542 except KeyError:
|
|
jpayne@69
|
543 pass
|
|
jpayne@69
|
544 else:
|
|
jpayne@69
|
545 if set_length:
|
|
jpayne@69
|
546 util.debug('doing set_length()')
|
|
jpayne@69
|
547 idx = task[1] if task else -1
|
|
jpayne@69
|
548 set_length(idx + 1)
|
|
jpayne@69
|
549 continue
|
|
jpayne@69
|
550 break
|
|
jpayne@69
|
551 finally:
|
|
jpayne@69
|
552 task = taskseq = job = None
|
|
jpayne@69
|
553 else:
|
|
jpayne@69
|
554 util.debug('task handler got sentinel')
|
|
jpayne@69
|
555
|
|
jpayne@69
|
556 try:
|
|
jpayne@69
|
557 # tell result handler to finish when cache is empty
|
|
jpayne@69
|
558 util.debug('task handler sending sentinel to result handler')
|
|
jpayne@69
|
559 outqueue.put(None)
|
|
jpayne@69
|
560
|
|
jpayne@69
|
561 # tell workers there is no more work
|
|
jpayne@69
|
562 util.debug('task handler sending sentinel to workers')
|
|
jpayne@69
|
563 for p in pool:
|
|
jpayne@69
|
564 put(None)
|
|
jpayne@69
|
565 except OSError:
|
|
jpayne@69
|
566 util.debug('task handler got OSError when sending sentinels')
|
|
jpayne@69
|
567
|
|
jpayne@69
|
568 util.debug('task handler exiting')
|
|
jpayne@69
|
569
|
|
jpayne@69
|
570 @staticmethod
|
|
jpayne@69
|
571 def _handle_results(outqueue, get, cache):
|
|
jpayne@69
|
572 thread = threading.current_thread()
|
|
jpayne@69
|
573
|
|
jpayne@69
|
574 while 1:
|
|
jpayne@69
|
575 try:
|
|
jpayne@69
|
576 task = get()
|
|
jpayne@69
|
577 except (OSError, EOFError):
|
|
jpayne@69
|
578 util.debug('result handler got EOFError/OSError -- exiting')
|
|
jpayne@69
|
579 return
|
|
jpayne@69
|
580
|
|
jpayne@69
|
581 if thread._state != RUN:
|
|
jpayne@69
|
582 assert thread._state == TERMINATE, "Thread not in TERMINATE"
|
|
jpayne@69
|
583 util.debug('result handler found thread._state=TERMINATE')
|
|
jpayne@69
|
584 break
|
|
jpayne@69
|
585
|
|
jpayne@69
|
586 if task is None:
|
|
jpayne@69
|
587 util.debug('result handler got sentinel')
|
|
jpayne@69
|
588 break
|
|
jpayne@69
|
589
|
|
jpayne@69
|
590 job, i, obj = task
|
|
jpayne@69
|
591 try:
|
|
jpayne@69
|
592 cache[job]._set(i, obj)
|
|
jpayne@69
|
593 except KeyError:
|
|
jpayne@69
|
594 pass
|
|
jpayne@69
|
595 task = job = obj = None
|
|
jpayne@69
|
596
|
|
jpayne@69
|
597 while cache and thread._state != TERMINATE:
|
|
jpayne@69
|
598 try:
|
|
jpayne@69
|
599 task = get()
|
|
jpayne@69
|
600 except (OSError, EOFError):
|
|
jpayne@69
|
601 util.debug('result handler got EOFError/OSError -- exiting')
|
|
jpayne@69
|
602 return
|
|
jpayne@69
|
603
|
|
jpayne@69
|
604 if task is None:
|
|
jpayne@69
|
605 util.debug('result handler ignoring extra sentinel')
|
|
jpayne@69
|
606 continue
|
|
jpayne@69
|
607 job, i, obj = task
|
|
jpayne@69
|
608 try:
|
|
jpayne@69
|
609 cache[job]._set(i, obj)
|
|
jpayne@69
|
610 except KeyError:
|
|
jpayne@69
|
611 pass
|
|
jpayne@69
|
612 task = job = obj = None
|
|
jpayne@69
|
613
|
|
jpayne@69
|
614 if hasattr(outqueue, '_reader'):
|
|
jpayne@69
|
615 util.debug('ensuring that outqueue is not full')
|
|
jpayne@69
|
616 # If we don't make room available in outqueue then
|
|
jpayne@69
|
617 # attempts to add the sentinel (None) to outqueue may
|
|
jpayne@69
|
618 # block. There is guaranteed to be no more than 2 sentinels.
|
|
jpayne@69
|
619 try:
|
|
jpayne@69
|
620 for i in range(10):
|
|
jpayne@69
|
621 if not outqueue._reader.poll():
|
|
jpayne@69
|
622 break
|
|
jpayne@69
|
623 get()
|
|
jpayne@69
|
624 except (OSError, EOFError):
|
|
jpayne@69
|
625 pass
|
|
jpayne@69
|
626
|
|
jpayne@69
|
627 util.debug('result handler exiting: len(cache)=%s, thread._state=%s',
|
|
jpayne@69
|
628 len(cache), thread._state)
|
|
jpayne@69
|
629
|
|
jpayne@69
|
630 @staticmethod
|
|
jpayne@69
|
631 def _get_tasks(func, it, size):
|
|
jpayne@69
|
632 it = iter(it)
|
|
jpayne@69
|
633 while 1:
|
|
jpayne@69
|
634 x = tuple(itertools.islice(it, size))
|
|
jpayne@69
|
635 if not x:
|
|
jpayne@69
|
636 return
|
|
jpayne@69
|
637 yield (func, x)
|
|
jpayne@69
|
638
|
|
jpayne@69
|
639 def __reduce__(self):
|
|
jpayne@69
|
640 raise NotImplementedError(
|
|
jpayne@69
|
641 'pool objects cannot be passed between processes or pickled'
|
|
jpayne@69
|
642 )
|
|
jpayne@69
|
643
|
|
jpayne@69
|
644 def close(self):
|
|
jpayne@69
|
645 util.debug('closing pool')
|
|
jpayne@69
|
646 if self._state == RUN:
|
|
jpayne@69
|
647 self._state = CLOSE
|
|
jpayne@69
|
648 self._worker_handler._state = CLOSE
|
|
jpayne@69
|
649 self._change_notifier.put(None)
|
|
jpayne@69
|
650
|
|
jpayne@69
|
651 def terminate(self):
|
|
jpayne@69
|
652 util.debug('terminating pool')
|
|
jpayne@69
|
653 self._state = TERMINATE
|
|
jpayne@69
|
654 self._worker_handler._state = TERMINATE
|
|
jpayne@69
|
655 self._change_notifier.put(None)
|
|
jpayne@69
|
656 self._terminate()
|
|
jpayne@69
|
657
|
|
jpayne@69
|
658 def join(self):
|
|
jpayne@69
|
659 util.debug('joining pool')
|
|
jpayne@69
|
660 if self._state == RUN:
|
|
jpayne@69
|
661 raise ValueError("Pool is still running")
|
|
jpayne@69
|
662 elif self._state not in (CLOSE, TERMINATE):
|
|
jpayne@69
|
663 raise ValueError("In unknown state")
|
|
jpayne@69
|
664 self._worker_handler.join()
|
|
jpayne@69
|
665 self._task_handler.join()
|
|
jpayne@69
|
666 self._result_handler.join()
|
|
jpayne@69
|
667 for p in self._pool:
|
|
jpayne@69
|
668 p.join()
|
|
jpayne@69
|
669
|
|
jpayne@69
|
670 @staticmethod
|
|
jpayne@69
|
671 def _help_stuff_finish(inqueue, task_handler, size):
|
|
jpayne@69
|
672 # task_handler may be blocked trying to put items on inqueue
|
|
jpayne@69
|
673 util.debug('removing tasks from inqueue until task handler finished')
|
|
jpayne@69
|
674 inqueue._rlock.acquire()
|
|
jpayne@69
|
675 while task_handler.is_alive() and inqueue._reader.poll():
|
|
jpayne@69
|
676 inqueue._reader.recv()
|
|
jpayne@69
|
677 time.sleep(0)
|
|
jpayne@69
|
678
|
|
jpayne@69
|
679 @classmethod
|
|
jpayne@69
|
680 def _terminate_pool(cls, taskqueue, inqueue, outqueue, pool, change_notifier,
|
|
jpayne@69
|
681 worker_handler, task_handler, result_handler, cache):
|
|
jpayne@69
|
682 # this is guaranteed to only be called once
|
|
jpayne@69
|
683 util.debug('finalizing pool')
|
|
jpayne@69
|
684
|
|
jpayne@69
|
685 worker_handler._state = TERMINATE
|
|
jpayne@69
|
686 task_handler._state = TERMINATE
|
|
jpayne@69
|
687
|
|
jpayne@69
|
688 util.debug('helping task handler/workers to finish')
|
|
jpayne@69
|
689 cls._help_stuff_finish(inqueue, task_handler, len(pool))
|
|
jpayne@69
|
690
|
|
jpayne@69
|
691 if (not result_handler.is_alive()) and (len(cache) != 0):
|
|
jpayne@69
|
692 raise AssertionError(
|
|
jpayne@69
|
693 "Cannot have cache with result_hander not alive")
|
|
jpayne@69
|
694
|
|
jpayne@69
|
695 result_handler._state = TERMINATE
|
|
jpayne@69
|
696 change_notifier.put(None)
|
|
jpayne@69
|
697 outqueue.put(None) # sentinel
|
|
jpayne@69
|
698
|
|
jpayne@69
|
699 # We must wait for the worker handler to exit before terminating
|
|
jpayne@69
|
700 # workers because we don't want workers to be restarted behind our back.
|
|
jpayne@69
|
701 util.debug('joining worker handler')
|
|
jpayne@69
|
702 if threading.current_thread() is not worker_handler:
|
|
jpayne@69
|
703 worker_handler.join()
|
|
jpayne@69
|
704
|
|
jpayne@69
|
705 # Terminate workers which haven't already finished.
|
|
jpayne@69
|
706 if pool and hasattr(pool[0], 'terminate'):
|
|
jpayne@69
|
707 util.debug('terminating workers')
|
|
jpayne@69
|
708 for p in pool:
|
|
jpayne@69
|
709 if p.exitcode is None:
|
|
jpayne@69
|
710 p.terminate()
|
|
jpayne@69
|
711
|
|
jpayne@69
|
712 util.debug('joining task handler')
|
|
jpayne@69
|
713 if threading.current_thread() is not task_handler:
|
|
jpayne@69
|
714 task_handler.join()
|
|
jpayne@69
|
715
|
|
jpayne@69
|
716 util.debug('joining result handler')
|
|
jpayne@69
|
717 if threading.current_thread() is not result_handler:
|
|
jpayne@69
|
718 result_handler.join()
|
|
jpayne@69
|
719
|
|
jpayne@69
|
720 if pool and hasattr(pool[0], 'terminate'):
|
|
jpayne@69
|
721 util.debug('joining pool workers')
|
|
jpayne@69
|
722 for p in pool:
|
|
jpayne@69
|
723 if p.is_alive():
|
|
jpayne@69
|
724 # worker has not yet exited
|
|
jpayne@69
|
725 util.debug('cleaning up worker %d' % p.pid)
|
|
jpayne@69
|
726 p.join()
|
|
jpayne@69
|
727
|
|
jpayne@69
|
728 def __enter__(self):
|
|
jpayne@69
|
729 self._check_running()
|
|
jpayne@69
|
730 return self
|
|
jpayne@69
|
731
|
|
jpayne@69
|
732 def __exit__(self, exc_type, exc_val, exc_tb):
|
|
jpayne@69
|
733 self.terminate()
|
|
jpayne@69
|
734
|
|
jpayne@69
|
735 #
|
|
jpayne@69
|
736 # Class whose instances are returned by `Pool.apply_async()`
|
|
jpayne@69
|
737 #
|
|
jpayne@69
|
738
|
|
jpayne@69
|
739 class ApplyResult(object):
|
|
jpayne@69
|
740
|
|
jpayne@69
|
741 def __init__(self, pool, callback, error_callback):
|
|
jpayne@69
|
742 self._pool = pool
|
|
jpayne@69
|
743 self._event = threading.Event()
|
|
jpayne@69
|
744 self._job = next(job_counter)
|
|
jpayne@69
|
745 self._cache = pool._cache
|
|
jpayne@69
|
746 self._callback = callback
|
|
jpayne@69
|
747 self._error_callback = error_callback
|
|
jpayne@69
|
748 self._cache[self._job] = self
|
|
jpayne@69
|
749
|
|
jpayne@69
|
750 def ready(self):
|
|
jpayne@69
|
751 return self._event.is_set()
|
|
jpayne@69
|
752
|
|
jpayne@69
|
753 def successful(self):
|
|
jpayne@69
|
754 if not self.ready():
|
|
jpayne@69
|
755 raise ValueError("{0!r} not ready".format(self))
|
|
jpayne@69
|
756 return self._success
|
|
jpayne@69
|
757
|
|
jpayne@69
|
758 def wait(self, timeout=None):
|
|
jpayne@69
|
759 self._event.wait(timeout)
|
|
jpayne@69
|
760
|
|
jpayne@69
|
761 def get(self, timeout=None):
|
|
jpayne@69
|
762 self.wait(timeout)
|
|
jpayne@69
|
763 if not self.ready():
|
|
jpayne@69
|
764 raise TimeoutError
|
|
jpayne@69
|
765 if self._success:
|
|
jpayne@69
|
766 return self._value
|
|
jpayne@69
|
767 else:
|
|
jpayne@69
|
768 raise self._value
|
|
jpayne@69
|
769
|
|
jpayne@69
|
770 def _set(self, i, obj):
|
|
jpayne@69
|
771 self._success, self._value = obj
|
|
jpayne@69
|
772 if self._callback and self._success:
|
|
jpayne@69
|
773 self._callback(self._value)
|
|
jpayne@69
|
774 if self._error_callback and not self._success:
|
|
jpayne@69
|
775 self._error_callback(self._value)
|
|
jpayne@69
|
776 self._event.set()
|
|
jpayne@69
|
777 del self._cache[self._job]
|
|
jpayne@69
|
778 self._pool = None
|
|
jpayne@69
|
779
|
|
jpayne@69
|
780 AsyncResult = ApplyResult # create alias -- see #17805
|
|
jpayne@69
|
781
|
|
jpayne@69
|
782 #
|
|
jpayne@69
|
783 # Class whose instances are returned by `Pool.map_async()`
|
|
jpayne@69
|
784 #
|
|
jpayne@69
|
785
|
|
jpayne@69
|
786 class MapResult(ApplyResult):
|
|
jpayne@69
|
787
|
|
jpayne@69
|
788 def __init__(self, pool, chunksize, length, callback, error_callback):
|
|
jpayne@69
|
789 ApplyResult.__init__(self, pool, callback,
|
|
jpayne@69
|
790 error_callback=error_callback)
|
|
jpayne@69
|
791 self._success = True
|
|
jpayne@69
|
792 self._value = [None] * length
|
|
jpayne@69
|
793 self._chunksize = chunksize
|
|
jpayne@69
|
794 if chunksize <= 0:
|
|
jpayne@69
|
795 self._number_left = 0
|
|
jpayne@69
|
796 self._event.set()
|
|
jpayne@69
|
797 del self._cache[self._job]
|
|
jpayne@69
|
798 else:
|
|
jpayne@69
|
799 self._number_left = length//chunksize + bool(length % chunksize)
|
|
jpayne@69
|
800
|
|
jpayne@69
|
801 def _set(self, i, success_result):
|
|
jpayne@69
|
802 self._number_left -= 1
|
|
jpayne@69
|
803 success, result = success_result
|
|
jpayne@69
|
804 if success and self._success:
|
|
jpayne@69
|
805 self._value[i*self._chunksize:(i+1)*self._chunksize] = result
|
|
jpayne@69
|
806 if self._number_left == 0:
|
|
jpayne@69
|
807 if self._callback:
|
|
jpayne@69
|
808 self._callback(self._value)
|
|
jpayne@69
|
809 del self._cache[self._job]
|
|
jpayne@69
|
810 self._event.set()
|
|
jpayne@69
|
811 self._pool = None
|
|
jpayne@69
|
812 else:
|
|
jpayne@69
|
813 if not success and self._success:
|
|
jpayne@69
|
814 # only store first exception
|
|
jpayne@69
|
815 self._success = False
|
|
jpayne@69
|
816 self._value = result
|
|
jpayne@69
|
817 if self._number_left == 0:
|
|
jpayne@69
|
818 # only consider the result ready once all jobs are done
|
|
jpayne@69
|
819 if self._error_callback:
|
|
jpayne@69
|
820 self._error_callback(self._value)
|
|
jpayne@69
|
821 del self._cache[self._job]
|
|
jpayne@69
|
822 self._event.set()
|
|
jpayne@69
|
823 self._pool = None
|
|
jpayne@69
|
824
|
|
jpayne@69
|
825 #
|
|
jpayne@69
|
826 # Class whose instances are returned by `Pool.imap()`
|
|
jpayne@69
|
827 #
|
|
jpayne@69
|
828
|
|
jpayne@69
|
829 class IMapIterator(object):
|
|
jpayne@69
|
830
|
|
jpayne@69
|
831 def __init__(self, pool):
|
|
jpayne@69
|
832 self._pool = pool
|
|
jpayne@69
|
833 self._cond = threading.Condition(threading.Lock())
|
|
jpayne@69
|
834 self._job = next(job_counter)
|
|
jpayne@69
|
835 self._cache = pool._cache
|
|
jpayne@69
|
836 self._items = collections.deque()
|
|
jpayne@69
|
837 self._index = 0
|
|
jpayne@69
|
838 self._length = None
|
|
jpayne@69
|
839 self._unsorted = {}
|
|
jpayne@69
|
840 self._cache[self._job] = self
|
|
jpayne@69
|
841
|
|
jpayne@69
|
842 def __iter__(self):
|
|
jpayne@69
|
843 return self
|
|
jpayne@69
|
844
|
|
jpayne@69
|
845 def next(self, timeout=None):
|
|
jpayne@69
|
846 with self._cond:
|
|
jpayne@69
|
847 try:
|
|
jpayne@69
|
848 item = self._items.popleft()
|
|
jpayne@69
|
849 except IndexError:
|
|
jpayne@69
|
850 if self._index == self._length:
|
|
jpayne@69
|
851 self._pool = None
|
|
jpayne@69
|
852 raise StopIteration from None
|
|
jpayne@69
|
853 self._cond.wait(timeout)
|
|
jpayne@69
|
854 try:
|
|
jpayne@69
|
855 item = self._items.popleft()
|
|
jpayne@69
|
856 except IndexError:
|
|
jpayne@69
|
857 if self._index == self._length:
|
|
jpayne@69
|
858 self._pool = None
|
|
jpayne@69
|
859 raise StopIteration from None
|
|
jpayne@69
|
860 raise TimeoutError from None
|
|
jpayne@69
|
861
|
|
jpayne@69
|
862 success, value = item
|
|
jpayne@69
|
863 if success:
|
|
jpayne@69
|
864 return value
|
|
jpayne@69
|
865 raise value
|
|
jpayne@69
|
866
|
|
jpayne@69
|
867 __next__ = next # XXX
|
|
jpayne@69
|
868
|
|
jpayne@69
|
869 def _set(self, i, obj):
|
|
jpayne@69
|
870 with self._cond:
|
|
jpayne@69
|
871 if self._index == i:
|
|
jpayne@69
|
872 self._items.append(obj)
|
|
jpayne@69
|
873 self._index += 1
|
|
jpayne@69
|
874 while self._index in self._unsorted:
|
|
jpayne@69
|
875 obj = self._unsorted.pop(self._index)
|
|
jpayne@69
|
876 self._items.append(obj)
|
|
jpayne@69
|
877 self._index += 1
|
|
jpayne@69
|
878 self._cond.notify()
|
|
jpayne@69
|
879 else:
|
|
jpayne@69
|
880 self._unsorted[i] = obj
|
|
jpayne@69
|
881
|
|
jpayne@69
|
882 if self._index == self._length:
|
|
jpayne@69
|
883 del self._cache[self._job]
|
|
jpayne@69
|
884 self._pool = None
|
|
jpayne@69
|
885
|
|
jpayne@69
|
886 def _set_length(self, length):
|
|
jpayne@69
|
887 with self._cond:
|
|
jpayne@69
|
888 self._length = length
|
|
jpayne@69
|
889 if self._index == self._length:
|
|
jpayne@69
|
890 self._cond.notify()
|
|
jpayne@69
|
891 del self._cache[self._job]
|
|
jpayne@69
|
892 self._pool = None
|
|
jpayne@69
|
893
|
|
jpayne@69
|
894 #
|
|
jpayne@69
|
895 # Class whose instances are returned by `Pool.imap_unordered()`
|
|
jpayne@69
|
896 #
|
|
jpayne@69
|
897
|
|
jpayne@69
|
898 class IMapUnorderedIterator(IMapIterator):
|
|
jpayne@69
|
899
|
|
jpayne@69
|
900 def _set(self, i, obj):
|
|
jpayne@69
|
901 with self._cond:
|
|
jpayne@69
|
902 self._items.append(obj)
|
|
jpayne@69
|
903 self._index += 1
|
|
jpayne@69
|
904 self._cond.notify()
|
|
jpayne@69
|
905 if self._index == self._length:
|
|
jpayne@69
|
906 del self._cache[self._job]
|
|
jpayne@69
|
907 self._pool = None
|
|
jpayne@69
|
908
|
|
jpayne@69
|
909 #
|
|
jpayne@69
|
910 #
|
|
jpayne@69
|
911 #
|
|
jpayne@69
|
912
|
|
jpayne@69
|
913 class ThreadPool(Pool):
|
|
jpayne@69
|
914 _wrap_exception = False
|
|
jpayne@69
|
915
|
|
jpayne@69
|
916 @staticmethod
|
|
jpayne@69
|
917 def Process(ctx, *args, **kwds):
|
|
jpayne@69
|
918 from .dummy import Process
|
|
jpayne@69
|
919 return Process(*args, **kwds)
|
|
jpayne@69
|
920
|
|
jpayne@69
|
921 def __init__(self, processes=None, initializer=None, initargs=()):
|
|
jpayne@69
|
922 Pool.__init__(self, processes, initializer, initargs)
|
|
jpayne@69
|
923
|
|
jpayne@69
|
924 def _setup_queues(self):
|
|
jpayne@69
|
925 self._inqueue = queue.SimpleQueue()
|
|
jpayne@69
|
926 self._outqueue = queue.SimpleQueue()
|
|
jpayne@69
|
927 self._quick_put = self._inqueue.put
|
|
jpayne@69
|
928 self._quick_get = self._outqueue.get
|
|
jpayne@69
|
929
|
|
jpayne@69
|
930 def _get_sentinels(self):
|
|
jpayne@69
|
931 return [self._change_notifier._reader]
|
|
jpayne@69
|
932
|
|
jpayne@69
|
933 @staticmethod
|
|
jpayne@69
|
934 def _get_worker_sentinels(workers):
|
|
jpayne@69
|
935 return []
|
|
jpayne@69
|
936
|
|
jpayne@69
|
937 @staticmethod
|
|
jpayne@69
|
938 def _help_stuff_finish(inqueue, task_handler, size):
|
|
jpayne@69
|
939 # drain inqueue, and put sentinels at its head to make workers finish
|
|
jpayne@69
|
940 try:
|
|
jpayne@69
|
941 while True:
|
|
jpayne@69
|
942 inqueue.get(block=False)
|
|
jpayne@69
|
943 except queue.Empty:
|
|
jpayne@69
|
944 pass
|
|
jpayne@69
|
945 for i in range(size):
|
|
jpayne@69
|
946 inqueue.put(None)
|
|
jpayne@69
|
947
|
|
jpayne@69
|
948 def _wait_for_updates(self, sentinels, change_notifier, timeout):
|
|
jpayne@69
|
949 time.sleep(timeout)
|
