3
\$\begingroup\$

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):
        time.sleep(0.02)
        yield i

def square(x):
    time.sleep(0.1)
    return x ** 2

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

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

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):
    pass


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

        Args:
            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
                ThreadPoolExecutor.
            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.
        Returns:
            An iterator equivalent to: map(func, *iterables) but the calls may
            be evaluated out-of-order.

        Raises:
            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
        else:
            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:
                    future_buffer.put(CancelledError())
                    return
                try:
                    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_buffer.put(e)
                    return
                try:
                    future = self.submit(fn, *args)
                except BaseException as e:
                    # E.g., RuntimeError from shut down executor.
                    # Forward the new exception downstream.
                    future_buffer.put(e)
                    return
                future_buffer.put(future)

        # 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):
                        break
                    else:
                        future.cancel()
                raise exc

            # Ensure cleanup happens even if client never starts this generator.
            try:
                yield None
            except GeneratorExit as exc:
                cleanup()
            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
                else:
                    remaining_timeout = end_time - time.time()
                # Reraise new exceptions (errors in the callable fn, TimeOut,
                # GeneratorExit) at map call site, but also cancel upstream.
                try:
                    yield future.result(remaining_timeout)
                except BaseException as exc:
                    cleanup()

        thread = threading.Thread(target=consume_inputs)
        thread.start()
        result = produce_results()
        # Consume the dummy `None` result
        next(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)

Update:

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.

\$\endgroup\$

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.