bpo-28699: fix abnormal behaviour of pools in multiprocessing.pool (GH-693)
an exception raised at the very first of an iterable would cause pools behave abnormally
(swallow the exception or hang)
diff --git a/Lib/multiprocessing/pool.py b/Lib/multiprocessing/pool.py
index ae8cec4..a545f3c 100644
--- a/Lib/multiprocessing/pool.py
+++ b/Lib/multiprocessing/pool.py
@@ -118,7 +118,7 @@
try:
result = (True, func(*args, **kwds))
except Exception as e:
- if wrap_exception:
+ if wrap_exception and func is not _helper_reraises_exception:
e = ExceptionWithTraceback(e, e.__traceback__)
result = (False, e)
try:
@@ -133,6 +133,10 @@
completed += 1
util.debug('worker exiting after %d tasks' % completed)
+def _helper_reraises_exception(ex):
+ 'Pickle-able helper function for use by _guarded_task_generation.'
+ raise ex
+
#
# Class representing a process pool
#
@@ -277,6 +281,17 @@
return self._map_async(func, iterable, starmapstar, chunksize,
callback, error_callback)
+ def _guarded_task_generation(self, result_job, func, iterable):
+ '''Provides a generator of tasks for imap and imap_unordered with
+ appropriate handling for iterables which throw exceptions during
+ iteration.'''
+ try:
+ i = -1
+ for i, x in enumerate(iterable):
+ yield (result_job, i, func, (x,), {})
+ except Exception as e:
+ yield (result_job, i+1, _helper_reraises_exception, (e,), {})
+
def imap(self, func, iterable, chunksize=1):
'''
Equivalent of `map()` -- can be MUCH slower than `Pool.map()`.
@@ -285,15 +300,23 @@
raise ValueError("Pool not running")
if chunksize == 1:
result = IMapIterator(self._cache)
- self._taskqueue.put((((result._job, i, func, (x,), {})
- for i, x in enumerate(iterable)), result._set_length))
+ self._taskqueue.put(
+ (
+ self._guarded_task_generation(result._job, func, iterable),
+ result._set_length
+ ))
return result
else:
assert chunksize > 1
task_batches = Pool._get_tasks(func, iterable, chunksize)
result = IMapIterator(self._cache)
- self._taskqueue.put((((result._job, i, mapstar, (x,), {})
- for i, x in enumerate(task_batches)), result._set_length))
+ self._taskqueue.put(
+ (
+ self._guarded_task_generation(result._job,
+ mapstar,
+ task_batches),
+ result._set_length
+ ))
return (item for chunk in result for item in chunk)
def imap_unordered(self, func, iterable, chunksize=1):
@@ -304,15 +327,23 @@
raise ValueError("Pool not running")
if chunksize == 1:
result = IMapUnorderedIterator(self._cache)
- self._taskqueue.put((((result._job, i, func, (x,), {})
- for i, x in enumerate(iterable)), result._set_length))
+ self._taskqueue.put(
+ (
+ self._guarded_task_generation(result._job, func, iterable),
+ result._set_length
+ ))
return result
else:
assert chunksize > 1
task_batches = Pool._get_tasks(func, iterable, chunksize)
result = IMapUnorderedIterator(self._cache)
- self._taskqueue.put((((result._job, i, mapstar, (x,), {})
- for i, x in enumerate(task_batches)), result._set_length))
+ self._taskqueue.put(
+ (
+ self._guarded_task_generation(result._job,
+ mapstar,
+ task_batches),
+ result._set_length
+ ))
return (item for chunk in result for item in chunk)
def apply_async(self, func, args=(), kwds={}, callback=None,
@@ -323,7 +354,7 @@
if self._state != RUN:
raise ValueError("Pool not running")
result = ApplyResult(self._cache, callback, error_callback)
- self._taskqueue.put(([(result._job, None, func, args, kwds)], None))
+ self._taskqueue.put(([(result._job, 0, func, args, kwds)], None))
return result
def map_async(self, func, iterable, chunksize=None, callback=None,
@@ -354,8 +385,14 @@
task_batches = Pool._get_tasks(func, iterable, chunksize)
result = MapResult(self._cache, chunksize, len(iterable), callback,
error_callback=error_callback)
- self._taskqueue.put((((result._job, i, mapper, (x,), {})
- for i, x in enumerate(task_batches)), None))
+ self._taskqueue.put(
+ (
+ self._guarded_task_generation(result._job,
+ mapper,
+ task_batches),
+ None
+ )
+ )
return result
@staticmethod
@@ -377,33 +414,27 @@
for taskseq, set_length in iter(taskqueue.get, None):
task = None
- i = -1
try:
- for i, task in enumerate(taskseq):
+ # iterating taskseq cannot fail
+ for task in taskseq:
if thread._state:
util.debug('task handler found thread._state != RUN')
break
try:
put(task)
except Exception as e:
- job, ind = task[:2]
+ job, idx = task[:2]
try:
- cache[job]._set(ind, (False, e))
+ cache[job]._set(idx, (False, e))
except KeyError:
pass
else:
if set_length:
util.debug('doing set_length()')
- set_length(i+1)
+ idx = task[1] if task else -1
+ set_length(idx + 1)
continue
break
- except Exception as ex:
- job, ind = task[:2] if task else (0, 0)
- if job in cache:
- cache[job]._set(ind + 1, (False, ex))
- if set_length:
- util.debug('doing set_length()')
- set_length(i+1)
finally:
task = taskseq = job = None
else: