diff --git a/README.md b/README.md index b96f331..f36ed17 100644 --- a/README.md +++ b/README.md @@ -212,6 +212,35 @@ with ProcessPoolExecutor(max_workers=2) as execute: If you insist on wanting to pass generators to your subprocesses, you can use iterators instead. The construction above allows us to write the generator call up front, pass only a few primatives through the pickling process, and transparently call the generator on the other side. We can even support the `len()` function which is not available for generators. +```python +# Listing 8: Easy Multiprocessing + +import gevent.monkey +gevent.monkey.patch_all(thread=False) +import copy +from taskqueue import GreenTaskQueue + +class PrintTaskIterator(object): + def __init__(self, start, end): + self.start = start + self.end = end + def __getitem__(self, slc): + itr = copy.deepcopy(self) + itr.start = self.start + slc.start + itr.end = self.start + slc.stop + return itr + def __len__(self): + return self.end - self.start + def __iter__(self): + for i in range(self.start, self.end): + yield PrintTask(i) + +tq = GreenTaskQueue('sqs-queue-name') +tq.insert_all(PrintTaskIterator(0,200), parallel=2) +``` + +If you design your iterators such that the slice operator works, TaskQueue can +automatically resection the iterator such that it can be fed to multiple processes. Notably, we don't return `PrintTaskIterator(self.start+slc.start, self.start+slc.stop)` because it triggers an endless recursion during pickling. However, the runtime copy implementation above sidesteps this issue. Internally, `PrintTaskIterator(0,200)` will be turned into `[ PrintTaskIterator(0,100), PrintTaskIterator(100,200) ]`. We also perform tracking of exceptions raised by child processes in a queue. `gevent.monkey.patch_all(thread=False)` was necessary to avoid multiprocess hanging. [1] You can't pass generators in CPython but [you can pass iterators](https://stackoverflow.com/questions/1939015/singleton-python-generator-or-pickle-a-python-generator/1939493#1939493). You can pass generators if you use Pypy or Stackless Python. diff --git a/requirements.txt b/requirements.txt index d9c83f9..497459d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ cloud-volume>=0.20.2 gevent numpy>=1.13.1 +pathos pytest>=3.3.1 pbr -e git+https://github.com/seung-lab/tqdm.git#egg=tqdm diff --git a/taskqueue/__init__.py b/taskqueue/__init__.py index 3ac2200..e2ccab8 100644 --- a/taskqueue/__init__.py +++ b/taskqueue/__init__.py @@ -1,5 +1,8 @@ from .registered_task import RegisteredTask, MockTask, PrintTask -from .taskqueue import TaskQueue, GreenTaskQueue, MockTaskQueue, LocalTaskQueue, upload +from .taskqueue import ( + TaskQueue, GreenTaskQueue, MockTaskQueue, LocalTaskQueue, + multiprocess_upload +) from .secrets import ( QUEUE_NAME, TEST_QUEUE_NAME, QUEUE_TYPE, PROJECT_NAME, AWS_DEFAULT_REGION diff --git a/taskqueue/taskqueue.py b/taskqueue/taskqueue.py index da6bb10..4e6a356 100644 --- a/taskqueue/taskqueue.py +++ b/taskqueue/taskqueue.py @@ -11,10 +11,10 @@ import time import traceback -import concurrent.futures import gevent.pool import multiprocessing as mp import numpy as np +import pathos.pools from tqdm import tqdm from cloudvolume.threaded_queue import ThreadedQueue @@ -28,6 +28,8 @@ AWS_DEFAULT_REGION ) +AWS_BATCH_SIZE = 10 + def totask(task): if isinstance(task, RegisteredTask): return task @@ -109,6 +111,53 @@ def list(self): """ return [ totask(x) for x in self._api.list() ] + def insert_all(self, tasks, delay_seconds=0, total=None, parallel=1): + if total is None: + try: + total = len(tasks) + except TypeError: + pass + + if parallel not in (1, False): + multiprocess_upload(self.__class__, self.queue_name, tasks, parallel=parallel) + else: + self._insert_all(tasks, delay_seconds, total) + + def _gen_insert_all_tasks( + self, tasks, + batch_size=AWS_BATCH_SIZE, delay_seconds=0, total=None + ): + + def genbatches(itr): + while True: + batch = [] + try: + for i in range(AWS_BATCH_SIZE): + batch.append(next(itr)) + except StopIteration: + pass + + if len(batch) == 0: + break + + yield batch + + bodies = ( + { + "payload": task.payload(), + "queueName": self._queue_name, + "groupByTag": True, + "tag": task.__class__.__name__ + } + + for task in tasks + ) + + def cloud_insertion(batch): + self._api.insert(batch, delay_seconds) + + return ( partial(cloud_insertion, batch) for batch in genbatches(bodies) ) + def renew_lease(self, task, seconds): """Update the duration of a task lease.""" return self._api.renew_lease(task, seconds) @@ -301,44 +350,12 @@ def cloud_insertion(api): return self - def insert_all(self, tasks, delay_seconds=0, total=None): - if total is None: - try: - total = len(tasks) - except TypeError: - pass - - batch_size = 10 - def genbatches(itr): - while True: - batch = [] - try: - for i in range(batch_size): - batch.append(next(itr)) - except StopIteration: - pass - - if len(batch) == 0: - raise StopIteration - - yield batch - - bodies = ( - { - "payload": task.payload(), - "queueName": self._queue_name, - "groupByTag": True, - "tag": task.__class__.__name__ - } - - for task in tasks + def _insert_all(self, tasks, delay_seconds=0, total=None): + batch_size = AWS_BATCH_SIZE + fns = self._gen_insert_all_tasks( + tasks, batch_size, delay_seconds, total ) - def cloud_insertion(batch): - self._api.insert(batch, delay_seconds) - - fns = ( partial(cloud_insertion, batch) for batch in genbatches(bodies) ) - schedule_threaded_jobs( fns=fns, concurrency=20, @@ -407,15 +424,16 @@ def check_monkey_patch_status(self): import gevent.monkey if not gevent.monkey.is_module_patched("socket"): print(yellow(""" - Switching ThreadedQueue to using green threads. This requires + GreenTaskQueue uses green threads. This requires monkey patching the standard library to use a cooperative non-blocking threading model. Please place the following lines at the beginning of your - program: + program. `thread=False` is there because sometimes this + causes hanging in multiprocessing. import gevent.monkey - gevent.monkey.patch_all() + gevent.monkey.patch_all(thread=False) """)) def insert(self, task, args=[], kwargs={}, delay_seconds=0): @@ -436,44 +454,12 @@ def cloud_insertion(): return self - def insert_all(self, tasks, delay_seconds=0, total=None): - if total is None: - try: - total = len(tasks) - except TypeError: - pass - - batch_size = 10 - def genbatches(itr): - while True: - batch = [] - try: - for i in range(batch_size): - batch.append(next(itr)) - except StopIteration: - pass - - if len(batch) == 0: - raise StopIteration - - yield batch - - bodies = ( - { - "payload": task.payload(), - "queueName": self._queue_name, - "groupByTag": True, - "tag": task.__class__.__name__ - } - - for task in tasks + def _insert_all(self, tasks, delay_seconds=0, total=None): + batch_size = AWS_BATCH_SIZE + fns = self._gen_insert_all_tasks( + tasks, batch_size, delay_seconds, total ) - def cloud_insertion(batch): - self._api.insert(batch, delay_seconds) - - fns = ( partial(cloud_insertion, batch) for batch in genbatches(bodies) ) - schedule_green_jobs( fns=fns, concurrency=20, @@ -548,7 +534,11 @@ def insert(self, task, args=[], kwargs={}): task.execute(*args, **kwargs) del task - def insert_all(self, tasks, args=[], kwargs={}, delay_seconds=0, total=None): + def insert_all( + self, tasks, args=[], kwargs={}, + delay_seconds=0, total=None, + parallel=1 + ): for task in tasks: self.insert(task, args=args, kwargs=kwargs) @@ -587,11 +577,18 @@ def insert(self, task, args=[], kwargs={}): def insert_all( self, tasks, args=[], kwargs={}, - delay_seconds=0, total=None + delay_seconds=0, total=None, + parallel=None, progress=True ): + + if parallel is None: + parallel = self.parallel + for task in tasks: self.queue.append( (task, args, kwargs) ) + self._process(progress) + def wait(self, progress=None): self._process(progress) return self @@ -609,7 +606,7 @@ def _process(self, progress=True): _task_execute(self.queue.pop(0)) pbar.update() else: - with concurrent.futures.ProcessPoolExecutor(max_workers=self.parallel) as executor: + with pathos.pools.ProcessPool(self.parallel) as executor: for _ in executor.map(_task_execute, self.queue): pbar.update() @@ -632,38 +629,51 @@ def _task_execute(task_tuple): def _scatter(sequence, n): """Scatters elements of ``sequence`` into ``n`` blocks.""" - - output = [] chunklen = int(math.ceil(float(len(sequence)) / float(n))) - for i in range(n): - output.append( - copy.deepcopy(sequence[i*chunklen:(i+1)*chunklen]) - ) + return [ + sequence[ i*chunklen : (i+1)*chunklen ] for i in range(n) + ] - return output +def soloprocess_upload(QueueClass, queue_name, tasks): + with QueueClass(queue_name) as tq: + tq.insert_all(tasks) -def _upload_tasks_single_process(queue_name, tasks): - with TaskQueue(queue_name=queue_name, queue_server='sqs') as tq: - for task in tasks: - tq.insert(task) +error_queue = mp.Queue() -def upload(queue_name, tasks, parallel=16, progress=False): - if parallel == 0: +def multiprocess_upload(QueueClass, queue_name, tasks, parallel=True): + if parallel is True: parallel = mp.cpu_count() - elif parallel < 0: + elif parallel <= 0: raise ValueError("Parallel must be a positive number or zero (all cpus). Got: " + str(parallel)) - parallel = max(1, min(parallel, mp.cpu_count())) - if parallel == 1: - _upload_tasks_single_process(queue_name, tasks) - else: - uploadfn = partial(_upload_tasks_single_process, queue_name) - tasks = _scatter(tasks, parallel) + soloprocess_upload(QueueClass, queue_name, tasks) + return - with concurrent.futures.ProcessPoolExecutor(max_workers=parallel) as executor: - executor.map(uploadfn, tasks) + def capturing_soloprocess_upload(*args, **kwargs): + try: + soloprocess_upload(*args, **kwargs) + except Exception as err: + print(err) + error_queue.put(err) + + uploadfn = partial( + capturing_soloprocess_upload, QueueClass, queue_name + ) + tasks = _scatter(tasks, parallel) + + pool = pathos.pools.ProcessPool(parallel) + pool.map(uploadfn, tasks) + + if not error_queue.empty(): + errors = [] + while not error_queue.empty(): + err = error_queue.get() + if err is not StopIteration: + errors.append(err) + if len(errors): + raise Exception(errors)