# Implementing non-blocking Executor.map

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

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

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 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
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()

result = produce_results()
# Consume the dummy None result
next(result)
return result

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.