Callbacks should complete immediately since otherwise the thread which
handles the results will get blocked.
- .. method:: imap(func, iterable[, chunksize])
+ .. method:: imap(func, iterable, chunksize=1, *, buffersize=None)
A lazier version of :meth:`.map`.
``next(timeout)`` will raise :exc:`multiprocessing.TimeoutError` if the
result cannot be returned within *timeout* seconds.
- .. method:: imap_unordered(func, iterable[, chunksize])
+ The *iterable* is collected immediately rather than lazily, unless a
+ *buffersize* is specified to limit the number of submitted tasks whose
+ results have not yet been yielded. If the buffer is full, iteration over
+ the *iterables* pauses until a result is yielded from the buffer.
+ To fully utilize pool's capacity when using this feature,
+ set *buffersize* at least to the number of processes in pool
+ (to consume *iterable* as you go), or even higher
+ (to prefetch the next ``N=buffersize-processes`` arguments).
+
+ .. versionchanged:: next
+ Added the *buffersize* parameter.
+
+ .. method:: imap_unordered(func, iterable, chunksize=1, *, buffersize=None)
The same as :meth:`imap` except that the ordering of the results from the
returned iterator should be considered arbitrary. (Only when there is
only one worker process is the order guaranteed to be "correct".)
+ .. versionchanged:: next
+ Added the *buffersize* parameter.
+
.. method:: starmap(func, iterable[, chunksize])
Like :meth:`~multiprocessing.pool.Pool.map` except that the
(Contributed by Jeff Epler in :gh:`150534`.)
+multiprocessing
+---------------
+
+* Add the optional ``buffersize`` parameter to
+ :meth:`multiprocessing.pool.Pool.imap` and
+ :meth:`multiprocessing.pool.Pool.imap_unordered` to limit the number of
+ submitted tasks whose results have not yet been yielded. If the buffer is
+ full, iteration over the *iterables* pauses until a result is yielded from
+ the buffer. To fully utilize pool's capacity when using this feature, set
+ *buffersize* at least to the number of processes in pool (to consume
+ *iterable* as you go), or even higher (to prefetch the next
+ ``N=buffersize-processes`` arguments).
+
+ (Contributed by Oleksandr Baltian in :gh:`136871`.)
+
+
os
--
Removed
-=======
+========
annotationlib
-------------
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.
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._result_handler, self._cache,
+ self._taskqueue_buffersize_semaphores),
exitpriority=15
)
self._state = RUN
return self._map_async(func, iterable, starmapstar, chunksize,
callback, error_callback)
- def _guarded_task_generation(self, result_job, func, iterable):
+ def _guarded_task_generation(self, result_job, func, iterable, sema=None):
'''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,), {})
+
+ 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()
+ try:
+ i, x = next(enumerated_iter)
+ except StopIteration:
+ break
+ 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):
+ def imap(self, func, iterable, chunksize=1, *, buffersize=None):
'''
Equivalent of `map()` -- can be MUCH slower than `Pool.map()`.
'''
- self._check_running()
- if chunksize == 1:
- result = IMapIterator(self)
- self._taskqueue.put(
- (
- self._guarded_task_generation(result._job, func, iterable),
- result._set_length
- ))
- return result
- else:
- if chunksize < 1:
- raise ValueError(
- "Chunksize must be 1+, not {0:n}".format(
- chunksize))
- task_batches = Pool._get_tasks(func, iterable, chunksize)
- result = IMapIterator(self)
- 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)
+ return self._imap(IMapIterator, func, iterable, chunksize,
+ buffersize=buffersize)
- def imap_unordered(self, func, iterable, chunksize=1):
+ def imap_unordered(self, func, iterable, chunksize=1, *, buffersize=None):
'''
Like `imap()` method but ordering of results is arbitrary.
'''
- self._check_running()
- if chunksize == 1:
- result = IMapUnorderedIterator(self)
- self._taskqueue.put(
- (
- self._guarded_task_generation(result._job, func, iterable),
- result._set_length
- ))
- return result
- else:
- if chunksize < 1:
- raise ValueError(
- "Chunksize must be 1+, not {0!r}".format(chunksize))
- task_batches = Pool._get_tasks(func, iterable, chunksize)
- result = IMapUnorderedIterator(self)
- 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)
+ return self._imap(IMapUnorderedIterator, func, iterable, chunksize,
+ buffersize=buffersize)
def apply_async(self, func, args=(), kwds={}, callback=None,
error_callback=None):
)
return result
+ def _imap(self, iterator_cls, func, iterable, chunksize=1,
+ *, buffersize=None):
+ self._check_running()
+ if chunksize < 1:
+ raise ValueError(
+ f"Chunksize must be 1+, not {chunksize}"
+ )
+ if buffersize is not None:
+ if not isinstance(buffersize, int):
+ raise TypeError("buffersize must be an integer or None")
+ if buffersize < 1:
+ raise ValueError("buffersize must be None or > 0")
+
+ result = iterator_cls(self, buffersize=buffersize)
+ if chunksize == 1:
+ self._taskqueue.put(
+ (
+ self._guarded_task_generation(result._job, func, iterable,
+ result._buffersize_sema),
+ result._set_length,
+ )
+ )
+ return 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)
+
@staticmethod
def _wait_for_updates(sentinels, change_notifier, timeout=None):
wait(sentinels, timeout=timeout)
@classmethod
def _terminate_pool(cls, taskqueue, inqueue, outqueue, pool, change_notifier,
- worker_handler, task_handler, result_handler, cache):
+ worker_handler, task_handler, result_handler, cache,
+ taskqueue_buffersize_semaphores):
# this is guaranteed to only be called once
util.debug('finalizing pool')
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))
class IMapIterator(object):
- def __init__(self, pool):
+ def __init__(self, pool, *, buffersize=None):
self._pool = pool
self._cond = threading.Condition(threading.Lock())
self._job = next(job_counter)
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)
def __iter__(self):
return self
item = self._items.popleft()
except IndexError:
if self._index == self._length:
- self._pool = None
- raise StopIteration from None
+ self._stop_iterator()
self._cond.wait(timeout)
try:
item = self._items.popleft()
except IndexError:
if self._index == self._length:
- self._pool = None
- raise StopIteration from None
+ self._stop_iterator()
raise TimeoutError from None
+ if self._buffersize_sema is not None:
+ self._buffersize_sema.release()
+
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 _set(self, i, obj):
class _TestPool(BaseTestCase):
+ _POOL_SIZE = 4
+
@classmethod
def setUpClass(cls):
with warnings_helper.ignore_fork_in_thread_deprecation_warnings():
super().setUpClass()
- cls.pool = cls.Pool(4)
+ cls.pool = cls.Pool(cls._POOL_SIZE)
@classmethod
def tearDownClass(cls):
p.terminate()
p.join()
- def test_imap(self):
- it = self.pool.imap(sqr, list(range(10)))
- self.assertEqual(list(it), list(map(sqr, list(range(10)))))
-
- it = self.pool.imap(sqr, list(range(10)))
+ @support.subTests('buffersize', (
+ None,
+ 1,
+ _POOL_SIZE,
+ _POOL_SIZE * 2,
+ ))
+ def test_imap(self, buffersize):
+ iterable = range(10)
+ if self.TYPE != "threads":
+ iterable = list(iterable)
+ it = self.pool.imap(sqr, iterable, buffersize=buffersize)
for i in range(10):
- self.assertEqual(next(it), i*i)
+ self.assertEqual(next(it), i * i)
+ self.assertRaises(StopIteration, it.__next__)
+ # again, verify that it's truly exhausted
self.assertRaises(StopIteration, it.__next__)
- it = self.pool.imap(sqr, list(range(1000)), chunksize=100)
+ @support.subTests(('chunksize', 'buffersize'), (
+ (100, None),
+ (100, _POOL_SIZE),
+ ))
+ def test_imap_with_chunksize(self, chunksize, buffersize):
+ iterable = range(1000)
+ if self.TYPE != "threads":
+ iterable = list(iterable)
+ it = self.pool.imap(sqr, iterable, chunksize=chunksize, buffersize=buffersize)
for i in range(1000):
- self.assertEqual(next(it), i*i)
+ self.assertEqual(next(it), i * i)
+ self.assertRaises(StopIteration, it.__next__)
+ # again, verify that it's truly exhausted
self.assertRaises(StopIteration, it.__next__)
def test_imap_handle_iterable_exception(self):
self.assertEqual(next(it), i*i)
self.assertRaises(SayWhenError, it.__next__)
- def test_imap_unordered(self):
- it = self.pool.imap_unordered(sqr, list(range(10)))
+ @support.subTests('buffersize', (
+ None,
+ 1,
+ _POOL_SIZE,
+ _POOL_SIZE * 2,
+ ))
+ def test_imap_unordered(self, buffersize):
+ iterable = range(10)
+ if self.TYPE != "threads":
+ iterable = list(iterable)
+ it = self.pool.imap(sqr, iterable, buffersize=buffersize)
self.assertEqual(sorted(it), list(map(sqr, list(range(10)))))
- it = self.pool.imap_unordered(sqr, list(range(1000)), chunksize=100)
+ @support.subTests(('chunksize', 'buffersize'), (
+ (100, None),
+ (100, _POOL_SIZE),
+ ))
+ def test_imap_unordered_with_chunksize(self, chunksize, buffersize):
+ iterable = range(1000)
+ if self.TYPE != "threads":
+ iterable = list(iterable)
+ it = self.pool.imap_unordered(sqr, iterable, chunksize=chunksize,
+ buffersize=buffersize)
self.assertEqual(sorted(it), list(map(sqr, list(range(1000)))))
def test_imap_unordered_handle_iterable_exception(self):
self.assertIn(value, expected_values)
expected_values.remove(value)
+ @support.subTests('method_name', ("imap", "imap_unordered"))
+ @support.subTests(('buffersize', 'expected_exception', 'expected_regex'), (
+ ("foo", TypeError, "buffersize must be an integer or None"),
+ (2.0, TypeError, "buffersize must be an integer or None"),
+ (0, ValueError, "buffersize must be None or > 0"),
+ (-1, ValueError, "buffersize must be None or > 0"),
+ ))
+ def test_imap_and_imap_unordered_with_buffersize_type_validation(
+ self, method_name, buffersize, expected_exception, expected_regex
+ ):
+ method = getattr(self.pool, method_name)
+ with self.assertRaisesRegex(expected_exception, expected_regex):
+ method(str, range(4), buffersize=buffersize)
+
+ @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
+ @support.subTests('method_name', ("imap", "imap_unordered"))
+ def test_imap_and_imap_unordered_when_buffer_is_full(self, method_name):
+ if self.TYPE != "threads":
+ self.skipTest("test not appropriate for {}".format(self.TYPE))
+
+ processes = 4
+ p = self.Pool(processes)
+ last_produced_task_arg = Value("i")
+
+ def produce_args():
+ for arg in itertools.count(1):
+ last_produced_task_arg.value = arg
+ yield arg
+
+ method = getattr(p, method_name)
+ it = method(functools.partial(sqr, wait=0.2), produce_args())
+
+ time.sleep(0.2)
+ # `iterable` could've been advanced only `processes` times,
+ # but in fact it advances further (`> processes`) because of
+ # not waiting for workers or user code to catch up.
+ self.assertGreater(last_produced_task_arg.value, processes)
+
+ next(it)
+ time.sleep(0.2)
+ self.assertGreater(last_produced_task_arg.value, processes + 1)
+
+ next(it)
+ time.sleep(0.2)
+ self.assertGreater(last_produced_task_arg.value, processes + 2)
+
+ p.terminate()
+ p.join()
+
+ @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
+ @support.subTests('method_name', ("imap", "imap_unordered"))
+ def test_imap_and_imap_unordered_with_buffersize_when_buffer_is_full(
+ self, method_name
+ ):
+ if self.TYPE != "threads":
+ self.skipTest("test not appropriate for {}".format(self.TYPE))
+
+ processes = 4
+ p = self.Pool(processes)
+ last_produced_task_arg = Value("i")
+
+ def produce_args():
+ for arg in itertools.count(1):
+ last_produced_task_arg.value = arg
+ yield arg
+
+ method = getattr(p, method_name)
+ it = method(functools.partial(sqr, wait=0.2), produce_args(),
+ buffersize=processes)
+
+ time.sleep(0.2)
+ self.assertEqual(last_produced_task_arg.value, processes)
+
+ next(it)
+ time.sleep(0.2)
+ self.assertEqual(last_produced_task_arg.value, processes + 1)
+
+ next(it)
+ time.sleep(0.2)
+ self.assertEqual(last_produced_task_arg.value, processes + 2)
+
+ p.terminate()
+ p.join()
+
+ @support.subTests('method_name', ("imap", "imap_unordered"))
+ def test_imap_and_imap_unordered_with_buffersize_on_empty_iterable(
+ self, method_name
+ ):
+ method = getattr(self.pool, method_name)
+ res = method(str, [], buffersize=2)
+ self.assertIsNone(next(res, None))
+ self.assertIsNone(next(res, None))
+
+ @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
+ def test_imap_with_buffersize_on_infinite_iterable(self):
+ if self.TYPE != "threads":
+ self.skipTest("test not appropriate for {}".format(self.TYPE))
+
+ p = self.Pool(4)
+ res = p.imap(str, itertools.count(), buffersize=2)
+
+ self.assertEqual(next(res, None), "0")
+ self.assertEqual(next(res, None), "1")
+ self.assertEqual(next(res, None), "2")
+
+ p.terminate()
+ p.join()
+
+ @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
+ def test_imap_unordered_with_buffersize_on_infinite_iterable(self):
+ if self.TYPE != "threads":
+ self.skipTest("test not appropriate for {}".format(self.TYPE))
+
+ p = self.Pool(4)
+ res = p.imap_unordered(str, itertools.count(), buffersize=2)
+
+ # (4, 5, ...) can also be submitted to the pool, so assert just 3 unique results
+ first_three_results = [next(res, None) for _ in range(3)]
+ self.assertEqual(len(first_three_results), 3)
+ self.assertEqual(len(set(first_three_results)), 3)
+
+ p.terminate()
+ p.join()
+
@warnings_helper.ignore_fork_in_thread_deprecation_warnings()
def test_make_pool(self):
expected_error = (RemoteError if self.TYPE == 'manager'
--- /dev/null
+Add the optional ``buffersize`` parameter to
+:meth:`multiprocessing.pool.Pool.imap` and
+:meth:`multiprocessing.pool.Pool.imap_unordered` to limit the number of
+submitted tasks whose results have not yet been yielded. If the buffer is
+full, iteration over the *iterables* pauses until a result is yielded from
+the buffer. To fully utilize pool's capacity when using this feature, set
+*buffersize* at least to the number of processes in pool (to consume
+*iterable* as you go), or even higher (to prefetch the next
+``N=buffersize-processes`` arguments).
projects / thirdparty / Python / cpython.git / commitdiff
