"""Parallel processing tools to help speed up certain tasks like data
preprocessing.
Authors
* Sylvain de Langen 2023
"""
import itertools
import multiprocessing
from collections import deque
from concurrent.futures import Executor, ProcessPoolExecutor
from threading import Condition
from typing import Any, Callable, Iterable, Optional
from tqdm.auto import tqdm
def _chunk_process_wrapper(fn, chunk):
return list(map(fn, chunk))
[docs]
class CancelFuturesOnExit:
"""Context manager that .cancel()s all elements of a list upon exit.
This is used to abort futures faster when raising an exception."""
def __init__(self, future_list):
self.future_list = future_list
def __enter__(self):
pass
def __exit__(self, _type, _value, _traceback):
for future in self.future_list:
future.cancel()
class _ParallelMapper:
"""Internal class for `parallel_map`, arguments match the constructor's."""
def __init__(
self,
fn: Callable[[Any], Any],
source: Iterable[Any],
process_count: int,
chunk_size: int,
queue_size: int,
executor: Optional[Executor],
progress_bar: bool,
progress_bar_kwargs: dict,
):
self.future_chunks = deque()
self.cv = Condition()
self.just_finished_count = 0
"""Number of jobs that were just done processing, guarded by
`self.cv`."""
self.remote_exception = None
"""Set by a worker when it encounters an exception, guarded by
`self.cv`."""
self.fn = fn
self.source = source
self.process_count = process_count
self.chunk_size = chunk_size
self.queue_size = queue_size
self.executor = executor
self.known_len = len(source) if hasattr(source, "__len__") else None
self.source_it = iter(source)
self.depleted_source = False
if progress_bar:
tqdm_final_kwargs = {"total": self.known_len}
tqdm_final_kwargs.update(progress_bar_kwargs)
self.pbar = tqdm(**tqdm_final_kwargs)
else:
self.pbar = None
def run(self):
"""Spins up an executor (if none were provided), then yields all
processed chunks in order."""
with CancelFuturesOnExit(self.future_chunks):
if self.executor is not None:
# just use the executor we were provided
yield from self._map_all()
else:
# start and shut down a process pool executor -- ok for
# long-running tasks
with ProcessPoolExecutor(
max_workers=self.process_count
) as pool:
self.executor = pool
yield from self._map_all()
def _bump_processed_count(self, future):
"""Notifies the main thread of the finished job, bumping the number of
jobs it should requeue. Updates the progress bar based on the returned
chunk length.
Arguments
---------
future: concurrent.futures.Future
A future holding a processed chunk (of type `list`).
"""
if future.cancelled():
# the scheduler wants us to stop or something else happened, give up
return
future_exception = future.exception()
# wake up dispatcher thread to refill the queue
with self.cv:
if future_exception is not None:
# signal to the main thread that it should raise
self.remote_exception = future_exception
self.just_finished_count += 1
self.cv.notify()
if future_exception is None:
# update progress bar with the length of the output as the progress
# bar is over element count, not chunk count.
if self.pbar is not None:
self.pbar.update(len(future.result()))
def _enqueue_job(self):
"""Pulls a chunk from the source iterable and submits it to the
pool; must be run from the main thread.
Returns
-------
`True` if any job was submitted (that is, if there was any chunk
left to process), `False` otherwise.
"""
# immediately deplete the input stream of chunk_size elems (or less)
chunk = list(itertools.islice(self.source_it, self.chunk_size))
# empty chunk? then we finished iterating over the input stream
if len(chunk) == 0:
self.depleted_source = True
return False
future = self.executor.submit(_chunk_process_wrapper, self.fn, chunk)
future.add_done_callback(self._bump_processed_count)
self.future_chunks.append(future)
return True
def _map_all(self):
"""Performs all the parallel mapping logic.
Arguments
---------
executor: concurrent.futures.Executor
The executor to `.invoke` all jobs on. The executor is NOT shut down
at the end of processing.
"""
# initial queue fill
for _ in range(self.queue_size):
if not self._enqueue_job():
break
# consume & requeue logic
while (not self.depleted_source) or (len(self.future_chunks) != 0):
with self.cv:
# if `cv.notify` was called by a worker _after_ the `with cv`
# block last iteration, then `just_finished_count` would be
# incremented, but this `cv.wait` would not wake up -- skip it.
while self.just_finished_count == 0:
# wait to be woken up by a worker thread, which could mean:
# - that a chunk was processed: try to yield any
# - that a call failed with an exception: raise it
# - nothing; it could be a spurious CV wakeup: keep looping
self.cv.wait()
if self.remote_exception is not None:
raise self.remote_exception
# store the amount to requeue, avoiding data races
to_queue_count = self.just_finished_count
self.just_finished_count = 0
# try to enqueue as many jobs as there were just finished.
# when the input is finished, the queue will not be refilled.
for _ in range(to_queue_count):
if not self._enqueue_job():
break
# yield from left to right as long as there is enough ready
# e.g. | done | done | !done | done | !done | !done
# would yield from the first two. we might deplete the entire queue
# at that point, the `depleted_source` loop check is needed as such.
while len(self.future_chunks) != 0 and self.future_chunks[0].done():
yield from self.future_chunks.popleft().result()
if self.pbar is not None:
self.pbar.close()
[docs]
def parallel_map(
fn: Callable[[Any], Any],
source: Iterable[Any],
process_count: int = multiprocessing.cpu_count(),
chunk_size: int = 8,
queue_size: int = 128,
executor: Optional[Executor] = None,
progress_bar: bool = True,
progress_bar_kwargs: dict = {"smoothing": 0.02},
):
"""Maps iterable items with a function, processing chunks of items in
parallel with multiple processes and displaying progress with tqdm.
Processed elements will always be returned in the original, correct order.
Unlike `ProcessPoolExecutor.map`, elements are produced AND consumed lazily.
Arguments
---------
fn
The function that is called for every element in the source list.
The output is an iterator over the source list after fn(elem) is called.
source: Iterable
Iterator whose elements are passed through the mapping function.
process_count: int
The number of processes to spawn. Ignored if a custom executor is
provided.
For CPU-bound tasks, it is generally not useful to exceed logical core
count.
For IO-bound tasks, it may make sense to as to limit the amount of time
spent in iowait.
chunk_size: int
How many elements are fed to the worker processes at once. A value of 8
is generally fine. Low values may increase overhead and reduce CPU
occupancy.
queue_size: int
Number of chunks to be waited for on the main process at a time.
Low values increase the chance of the queue being starved, forcing
workers to idle.
Very high values may cause high memory usage, especially if the source
iterable yields large objects.
executor: Optional[Executor]
Allows providing an existing executor (preferably a
ProcessPoolExecutor). If None (the default), a process pool will be
spawned for this mapping task and will be shut down after.
progress_bar: bool
Whether to show a tqdm progress bar.
progress_bar_kwargs: dict
A dict of keyword arguments that is forwarded to tqdm when
`progress_bar == True`. Allows overriding the defaults or e.g.
specifying `total` when it cannot be inferred from the source iterable.
"""
mapper = _ParallelMapper(
fn,
source,
process_count,
chunk_size,
queue_size,
executor,
progress_bar,
progress_bar_kwargs,
)
yield from mapper.run()