gh-64192: Add *buffersize* to imap()/imap_unordered() in multiprocessing.pool (GH-136871)

The new argument allows consuming the input iterator lazily, potentially
saving memory, and allowing long/infinite iterators.

It mirrors the *buffersize* argument to `concurrent.futures.Executor.map`
that was added in 3.14.

Co-authored-by: Oleksandr Baltian <oleksandr.baltian@maklai.com.ua>
Co-authored-by: Petr Viktorin <encukou@gmail.com>
Co-authored-by: Sasha Baltian <sasha.baltian@hyperexponential.com>

Author: 4 people

Doc/library/multiprocessing.rst

@@ -2536,7 +2536,7 @@ with the :class:`Pool` class.
       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`.
 
@@ -2550,12 +2550,27 @@ with the :class:`Pool` class.
       ``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

Doc/whatsnew/3.16.rst

@@ -376,6 +376,22 @@ math
   (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
 --
 
@@ -606,7 +622,7 @@ module_name
 
 
 Removed
-=======
+========
 
 annotationlib
 -------------

Lib/multiprocessing/pool.py

@@ -190,6 +190,11 @@ 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.
@@ -256,7 +261,8 @@ 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._result_handler, self._cache,
+                  self._taskqueue_buffersize_semaphores),
             exitpriority=15
             )
         self._state = RUN
@@ -382,73 +388,43 @@ 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):
+    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):
@@ -497,6 +473,41 @@ def _map_async(self, func, iterable, mapper, chunksize=None, 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)
@@ -679,7 +690,8 @@ 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):
+                        worker_handler, task_handler, result_handler, cache,
+                        taskqueue_buffersize_semaphores):
         # this is guaranteed to only be called once
         util.debug('finalizing pool')
 
@@ -690,6 +702,10 @@ 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))
@@ -836,7 +852,7 @@ def _set(self, i, success_result):
 
 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)
@@ -846,6 +862,11 @@ def __init__(self, pool):
         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
@@ -856,22 +877,30 @@ def next(self, timeout=None):
                 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):