class OptionalPool:
def __init__(self, threads=1):
self.pool = Pool(threads) if threads > 1 else None
def map(self, func, iterable, chunksize=None):
return self.pool.map(func, iterable,
chunksize) if self.pool is not None else map(
func, iterable)
def starmap(self, func, iterable, chunksize=None):
return self.pool.starmap(func, iterable, chunksize) if self.pool is not None \
else [func(*params) for params in iterable]
def __enter__(self):
if self.pool is not None:
self.pool.__enter__()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.pool is not None:
self.pool.__exit__(exc_type, exc_val, exc_tb)
def close(self):
if self.pool is not None:
self.pool.close()
class ProcessPool:
def __init__(self, processors: int = None, **kwargs):
if __debug__ and processors is None:
processors = 1
self._pool = Pool(processors, **kwargs)
self._results = list() # type: List[AsyncResult]
def queue_task(self, task: Callable, *args, **kwargs) -> AsyncResult:
res = self._pool.apply_async(task, args=args, kwds=kwargs)
self._results.append(res)
return res
def get_async_results(self) -> List[AsyncResult]:
slot = self._results
self._results = type(self._results)()
return slot
def get_results(self) -> Iterator[Any]:
return (i.get() for i in self.get_async_results())
def check_exceptions(self) -> None:
list(self.get_results())
def __enter__(self):
self._pool.__enter__()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
return self._pool.__exit__(exc_type, exc_val, exc_tb)
def close(self):
return self._pool.close()
def join(self):
return self._pool.join()
class LazyPool:
"""
A version of :class:`multiprocessing.Pool` that limits the number of inputs
being queued at the same time.
At the beginning, ``processes * item_per_process`` inputs are sent
to the pool. Further inputs are only sent if the previous ones have
already been processed. This prevents outputs from being overgenerated
and filling up the memory when the main processes cannot keep up.
Note: ATM only :meth:`imap` is implemented.
Note also that :class:`multiprocessing.Pool` cannot (easily?) be subclassed.
So :class:`LazyPool` is not a subclass of :class:`multiprocessing.Pool`, just
uses the latter. After Python 3.5, :class:`multiprocessing.pool.Pool` is
exposed (or simply exists?), so it's a bit easier.
"""
def __init__(self,
processes=None,
item_per_process=5,
initializer=None,
initargs=(),
maxtasksperchild=None):
self.item_per_process = item_per_process
self.max_items = item_per_process * processes
self.pool = Pool(processes, initializer, initargs, maxtasksperchild)
def __enter__(self):
self.pool.__enter__()
return self
def __exit__(self, *args):
return self.pool.__exit__(*args)
def imap(self, func: Callable,
iterable: Iterable) -> Generator[Tuple[Any, Any], None, None]:
outputs = deque(maxlen=self.max_items)
it = iter(iterable)
# Add the first batch of inputs at the same time
for inp in islice(it, self.max_items):
outputs.append(self.pool.apply_async(func, (inp, )))
# After that, add new input whenever one has been processed. We
# wait for the first result to keep the order of the input and the
# output corpus the same.
for inp in it:
outputs[0].wait()
yield outputs.popleft().get()
outputs.append(self.pool.apply_async(func, (inp, )))
# Just consume the rest
while outputs:
outputs[0].wait()
yield outputs.popleft().get()
def __getattr__(self, name):
return getattr(self.pool, name)
