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