I wanted to extend concurrent.futures.Executor to make the map method non-blocking. It seems to work fine, but I would be very interested in feedback about the general approach, implementation, and code quality. The entire thing is on Github, but below is all the important code:

Usage example:

import time
from itertools import count, islice

from streamexecutors import StreamThreadPoolExecutor

def produce():
    for i in range(100):
        yield i

def square(x):
    return x ** 2

def print10(iterable):
    print(list(islice(iterable, 10)))

squares = StreamThreadPoolExecutor().map(square, produce())

And the implementation:

import time
from queue import Queue
from concurrent.futures import Executor, ThreadPoolExecutor, ProcessPoolExecutor
from concurrent.futures.process import _get_chunks, _process_chunk
from functools import partial
import sys
import threading
import itertools

class CancelledError(Exception):

class StreamExecutor(Executor):
    def map(self, fn, *iterables, timeout=None, chunksize=1, buffer_size=10000):
        """Returns an iterator equivalent to map(fn, iter).

            fn: A callable that will take as many arguments as there are
                passed iterables.
            timeout: The maximum number of seconds to wait. If None, then there
                is no limit on the wait time.
            chunksize: The size of the chunks the iterable will be broken into
                before being passed to a child process. This argument is only
                used by ProcessPoolExecutor; it is ignored by
            buffer_size: The maximum number of input items that may be
                stored at once; default is a small buffer; 0 for no limit. The
                drawback of using a large buffer is the possibility of wasted
                computation and memory (in case not all input is needed), as
                well as higher peak memory usage.
            An iterator equivalent to: map(func, *iterables) but the calls may
            be evaluated out-of-order.

            TimeoutError: If the entire result iterator could not be generated
                before the given timeout.
            Exception: If fn(*args) raises for any values.
        if not callable(fn):
            raise TypeError('fn argument must be a callable')
        if timeout is None:
            end_time = None
            end_time = timeout + time.time()

        if buffer_size is None:
            buffer_size = -1
        elif buffer_size <= 0:
            raise ValueError('buffer_size must be a positive number')

        iterators = [iter(iterable) for iterable in iterables]

        # Set to True to gracefully terminate all producers
        cancel = False

        # Deadlocks on the two queues are avoided using the following rule.
        # The writer guarantees to place a sentinel value into the buffer
        # before exiting, and to write nothing after that; the reader
        # guarantees to read the queue until it encounters a sentinel value
        # and to stop reading after that. Any value of type BaseException is
        # treated as a sentinel.
        future_buffer = Queue(maxsize=buffer_size)

        # This function will run in a separate thread.
        def consume_inputs():
            while True:
                if cancel:
                    args = [next(iterator) for iterator in iterators]
                except BaseException as e:
                    # StopIteration represents exhausted input; any other
                    # exception is due to an error in the input generator. We
                    # forward the exception downstream so it can be raised
                    # when client iterates through the result of map.
                    future = self.submit(fn, *args)
                except BaseException as e:
                    # E.g., RuntimeError from shut down executor.
                    # Forward the new exception downstream.

        # This function will run in the main thread.
        def produce_results():
            def cleanup():
                nonlocal cancel
                cancel = True
                while True:
                    future = future_buffer.get()
                    if isinstance(future, BaseException):
                raise exc

            # Ensure cleanup happens even if client never starts this generator.
                yield None
            except GeneratorExit as exc:
            while True:
                future = future_buffer.get()
                if isinstance(future, BaseException):
                    # Reraise upstream exceptions at the map call site.
                    raise future
                if end_time is None:
                    remaining_timeout = None
                    remaining_timeout = end_time - time.time()
                # Reraise new exceptions (errors in the callable fn, TimeOut,
                # GeneratorExit) at map call site, but also cancel upstream.
                    yield future.result(remaining_timeout)
                except BaseException as exc:

        thread = threading.Thread(target=consume_inputs)
        result = produce_results()
        # Consume the dummy `None` result
        return result

class StreamThreadPoolExecutor(StreamExecutor, ThreadPoolExecutor): ...

class StreamProcessPoolExecutor(StreamExecutor, ProcessPoolExecutor):
    def map(self, fn, *iterables, timeout=None, chunksize=1, buffer_size=10000):
        if buffer_size is not None:
            buffer_size //= max(1, chunksize)
        if chunksize < 1:
            raise ValueError("chunksize must be >= 1.")
        results = super().map(partial(_process_chunk, fn),
                              _get_chunks(*iterables, chunksize=chunksize),
                              timeout=timeout, buffer_size=buffer_size)
        return itertools.chain.from_iterable(results)


I failed to fix an issue of the process hanging in some cases of main thread termination. I had to rewrite this code somewhat to ping the main thread to check if it's alive. Not sure if I should copy the updated version from Github to this post, since it might end up being too many edits. I guess I'll leave the original code, since I'm still interested to see if it's good: I prefer my original approach to the periodic ping.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.