Skip to content

gh-158677: Fix Pool.imap(buffersize=...) blocking other tasks on the pool #158792

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
oyiakoumis wants to merge 1 commit into from
Closed

gh-158677: Fix Pool.imap(buffersize=...) blocking other tasks on the pool #158792

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 76 additions & 73 deletions Lib/multiprocessing/pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,11 +190,6 @@ def __init__(self, processes=None, initializer=None, initargs=(),
self._ctx = context or get_context()
self._setup_queues()
self._taskqueue = queue.SimpleQueue()
# The _taskqueue_buffersize_semaphores exist to allow calling .release()
# on every active semaphore when the pool is terminating to let task_handler
# wake up to stop. It's a set so that each iterator object can efficiently
# deregister its semaphore when iterator finishes.
self._taskqueue_buffersize_semaphores = set()
# The _change_notifier queue exist to wake up self._handle_workers()
# when the cache (self._cache) is empty or when there is a change in
# the _state variable of the thread that runs _handle_workers.
Expand Down Expand Up @@ -261,8 +256,7 @@ def __init__(self, processes=None, initializer=None, initargs=(),
self, self._terminate_pool,
args=(self._taskqueue, self._inqueue, self._outqueue, self._pool,
self._change_notifier, self._worker_handler, self._task_handler,
self._result_handler, self._cache,
self._taskqueue_buffersize_semaphores),
self._result_handler, self._cache),
exitpriority=15
)
self._state = RUN
Expand Down Expand Up @@ -388,32 +382,14 @@ def starmap_async(self, func, iterable, chunksize=None, callback=None,
return self._map_async(func, iterable, starmapstar, chunksize,
callback, error_callback)

def _guarded_task_generation(self, result_job, func, iterable, sema=None):
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

if sema is None:
for i, x in enumerate(iterable):
yield (result_job, i, func, (x,), {})

else:
enumerated_iter = iter(enumerate(iterable))
while True:
sema.acquire()
if self._state != RUN:
# The pool is closing or terminating; stop submitting
# the still-throttled tasks so the task handler can
# finish instead of blocking here forever.
break
try:
i, x = next(enumerated_iter)
except StopIteration:
break
yield (result_job, i, func, (x,), {})

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,), {})

Expand Down Expand Up @@ -491,27 +467,20 @@ def _imap(self, iterator_cls, func, iterable, chunksize=1,
if buffersize < 1:
raise ValueError("buffersize must be None or > 0")

result = iterator_cls(self, buffersize=buffersize)
result = iterator_cls(self)
if chunksize == 1:
self._taskqueue.put(
(
self._guarded_task_generation(result._job, func, iterable,
result._buffersize_sema),
result._set_length,
)
)
return result
tasks = self._guarded_task_generation(result._job, func, iterable)
items = result
else:
task_batches = Pool._get_tasks(func, iterable, chunksize)
self._taskqueue.put(
(
self._guarded_task_generation(result._job, mapstar,
task_batches,
result._buffersize_sema),
result._set_length,
)
)
return (item for chunk in result for item in chunk)
tasks = self._guarded_task_generation(result._job, mapstar,
task_batches)
items = (item for chunk in result for item in chunk)
if buffersize is None:
self._taskqueue.put((tasks, result._set_length))
else:
result._throttle(tasks, buffersize)
return items

@staticmethod
def _wait_for_updates(sentinels, change_notifier, timeout=None):
Expand Down Expand Up @@ -664,12 +633,14 @@ def close(self):
util.debug('closing pool')
if self._state == RUN:
self._state = CLOSE
# Stop the imap iterators throttled by buffersize from submitting
# more tasks. This has to be done before the task handler is
# told to exit, so that it can still finish them.
for job in self._cache.copy().values():
if isinstance(job, IMapIterator):
job._stop_throttling()
self._worker_handler._state = CLOSE
self._change_notifier.put(None)
# Wake any task generator throttled on a buffersize semaphore so
# it observes the CLOSE state and stops submitting.
for sema in list(self._taskqueue_buffersize_semaphores):
sema.release()

def terminate(self):
util.debug('terminating pool')
Expand Down Expand Up @@ -699,8 +670,7 @@ def _help_stuff_finish(inqueue, task_handler, size):

@classmethod
def _terminate_pool(cls, taskqueue, inqueue, outqueue, pool, change_notifier,
worker_handler, task_handler, result_handler, cache,
taskqueue_buffersize_semaphores):
worker_handler, task_handler, result_handler, cache):
# this is guaranteed to only be called once
util.debug('finalizing pool')

Expand All @@ -711,10 +681,6 @@ def _terminate_pool(cls, taskqueue, inqueue, outqueue, pool, change_notifier,
change_notifier.put(None)

task_handler._state = TERMINATE
# Release all semaphores to wake up task_handler to stop.
for buffersize_sema in tuple(taskqueue_buffersize_semaphores):
buffersize_sema.release()
taskqueue_buffersize_semaphores.discard(buffersize_sema)

util.debug('helping task handler/workers to finish')
cls._help_stuff_finish(inqueue, task_handler, len(pool))
Expand Down Expand Up @@ -893,7 +859,7 @@ def _set(self, i, success_result):

class IMapIterator(object):

def __init__(self, pool, *, buffersize=None):
def __init__(self, pool):
self._pool = pool
self._cond = threading.Condition(threading.Lock())
self._job = next(job_counter)
Expand All @@ -903,11 +869,14 @@ def __init__(self, pool, *, buffersize=None):
self._length = None
self._unsorted = {}
self._cache[self._job] = self
if buffersize is None:
self._buffersize_sema = None
else:
self._buffersize_sema = threading.Semaphore(buffersize)
self._pool._taskqueue_buffersize_semaphores.add(self._buffersize_sema)
# Only used when the tasks are throttled by buffersize.
self._taskqueue = pool._taskqueue
self._throttle_lock = threading.Lock()
self._tasks = None
self._free_slots = None
self._submitted = 0
self._parked = False
self._stopped = False

def __iter__(self):
return self
Expand All @@ -918,32 +887,66 @@ def next(self, timeout=None):
item = self._items.popleft()
except IndexError:
if self._index == self._length:
self._stop_iterator()
self._pool = None
raise StopIteration from None
self._cond.wait(timeout)
try:
item = self._items.popleft()
except IndexError:
if self._index == self._length:
self._stop_iterator()
self._pool = None
raise StopIteration from None
raise TimeoutError from None

if self._buffersize_sema is not None:
self._buffersize_sema.release()
if self._free_slots is not None:
# A result left the buffer: there is room for one more task.
with self._throttle_lock:
self._free_slots += 1
self._resume_throttled_tasks()

success, value = item
if success:
return value
raise value

def _stop_iterator(self):
if self._pool is not None:
# `self._pool` could be set to `None` in previous `.next()` calls
self._pool._taskqueue_buffersize_semaphores.discard(self._buffersize_sema)
self._pool = None
raise StopIteration from None

__next__ = next # XXX

def _throttle(self, tasks, buffersize):
self._tasks = tasks
self._free_slots = buffersize
self._taskqueue.put((self._throttled_tasks(), None))

def _throttled_tasks(self):
# Runs in the task handler thread. It must not wait there for the
# buffer to have room, since that would hold up every other task of
# the pool. It returns instead, and is put back on the task queue
# once there is room again.
while True:
with self._throttle_lock:
if self._stopped:
break
if not self._free_slots:
self._parked = True
return
self._free_slots -= 1
task = next(self._tasks, None)
if task is None:
break
self._submitted += 1
yield task
self._set_length(self._submitted)

def _resume_throttled_tasks(self):
if self._parked:
self._parked = False
self._taskqueue.put((self._throttled_tasks(), None))

def _stop_throttling(self):
if self._free_slots is not None:
with self._throttle_lock:
self._stopped = True
self._resume_throttled_tasks()

def _set(self, i, obj):
with self._cond:
if self._index == i:
Expand Down
42 changes: 42 additions & 0 deletions Lib/test/_test_multiprocessing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3250,6 +3250,48 @@ def finalize():
t.join(support.SHORT_TIMEOUT)
self.assertTrue(finished.is_set(), "close()/join() deadlocked")

@support.subTests('method_name', ("imap", "imap_unordered"))
def test_imap_with_buffersize_does_not_block_other_tasks(
self, method_name
):
# gh-158677: a buffersize iterator waiting to be consumed must not
# prevent other tasks from being submitted to the pool.
p = self.Pool(2)
self.addCleanup(p.join)
self.addCleanup(p.terminate)
method = getattr(p, method_name)
a = method(sqr, range(10), buffersize=2)
b = method(sqr, range(10, 20), buffersize=2)

res = p.apply_async(sqr, (7,))
self.assertEqual(res.get(timeout=support.SHORT_TIMEOUT), 49)

results_a, results_b = zip(*zip(a, b))
self.assertEqual(sorted(results_a), list(map(sqr, range(10))))
self.assertEqual(sorted(results_b), list(map(sqr, range(10, 20))))

@support.subTests('method_name', ("imap", "imap_unordered"))
def test_imap_with_buffersize_handle_iterable_exception(
self, method_name
):
if self.TYPE == 'manager':
self.skipTest('test not appropriate for {}'.format(self.TYPE))

# An exception raised by the iterable is delivered by next(), at
# its position in the results.
method = getattr(self.pool, method_name)
it = method(sqr, exception_throwing_generator(1, -1), buffersize=2)
self.assertRaises(SayWhenError, it.__next__)
self.assertRaises(StopIteration, it.__next__)

it = method(sqr, exception_throwing_generator(10, 3), buffersize=2)
results = []
with self.assertRaises(SayWhenError):
for _ in range(4):
results.append(next(it))
results.extend(it)
self.assertEqual(sorted(results), list(map(sqr, range(3))))

@support.subTests('method_name', ("imap", "imap_unordered"))
def test_imap_and_imap_unordered_with_buffersize_on_empty_iterable(
self, method_name
Expand Down
Loading