Skip to content

Commit d99dff7

Browse files
feat: automatic multiprocessing (seung-lab#15)
* fix: in PEP 479, StopIteration is treated differently Prevent StopIteration from being converted into a RuntimeError by using break instead of raising directly. * feat: use pathos for multiprocessing * feat: insert_all with parallel argument * fix: add parallel arguments to insert_all for MockTaskQueue and LocalTaskQueue * docs: show how to use new multiprocessing facility
1 parent 7fe9ee4 commit d99dff7

4 files changed

Lines changed: 145 additions & 102 deletions

File tree

‎README.md‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,35 @@ with ProcessPoolExecutor(max_workers=2) as execute:
212212

213213
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.
214214

215+
```python
216+
# Listing 8: Easy Multiprocessing
217+
218+
import gevent.monkey
219+
gevent.monkey.patch_all(thread=False)
220+
import copy
221+
from taskqueue import GreenTaskQueue
222+
223+
class PrintTaskIterator(object):
224+
def __init__(self, start, end):
225+
self.start = start
226+
self.end = end
227+
def __getitem__(self, slc):
228+
itr = copy.deepcopy(self)
229+
itr.start = self.start + slc.start
230+
itr.end = self.start + slc.stop
231+
return itr
232+
def __len__(self):
233+
return self.end - self.start
234+
def __iter__(self):
235+
for i in range(self.start, self.end):
236+
yield PrintTask(i)
237+
238+
tq = GreenTaskQueue('sqs-queue-name')
239+
tq.insert_all(PrintTaskIterator(0,200), parallel=2)
240+
```
241+
242+
If you design your iterators such that the slice operator works, TaskQueue can
243+
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.
215244

216245
[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.
217246

‎requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
cloud-volume>=0.20.2
22
gevent
33
numpy>=1.13.1
4+
pathos
45
pytest>=3.3.1
56
pbr
67
-e git+https://github.com/seung-lab/tqdm.git#egg=tqdm

‎taskqueue/__init__.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
from .registered_task import RegisteredTask, MockTask, PrintTask
2-
from .taskqueue import TaskQueue, GreenTaskQueue, MockTaskQueue, LocalTaskQueue, upload
2+
from .taskqueue import (
3+
TaskQueue, GreenTaskQueue, MockTaskQueue, LocalTaskQueue,
4+
multiprocess_upload
5+
)
36
from .secrets import (
47
QUEUE_NAME, TEST_QUEUE_NAME, QUEUE_TYPE,
58
PROJECT_NAME, AWS_DEFAULT_REGION

‎taskqueue/taskqueue.py‎

Lines changed: 111 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,10 @@
1111
import time
1212
import traceback
1313

14-
import concurrent.futures
1514
import gevent.pool
1615
import multiprocessing as mp
1716
import numpy as np
17+
import pathos.pools
1818
from tqdm import tqdm
1919

2020
from cloudvolume.threaded_queue import ThreadedQueue
@@ -28,6 +28,8 @@
2828
AWS_DEFAULT_REGION
2929
)
3030

31+
AWS_BATCH_SIZE = 10
32+
3133
def totask(task):
3234
if isinstance(task, RegisteredTask):
3335
return task
@@ -109,6 +111,53 @@ def list(self):
109111
"""
110112
return [ totask(x) for x in self._api.list() ]
111113

114+
def insert_all(self, tasks, delay_seconds=0, total=None, parallel=1):
115+
if total is None:
116+
try:
117+
total = len(tasks)
118+
except TypeError:
119+
pass
120+
121+
if parallel not in (1, False):
122+
multiprocess_upload(self.__class__, self.queue_name, tasks, parallel=parallel)
123+
else:
124+
self._insert_all(tasks, delay_seconds, total)
125+
126+
def _gen_insert_all_tasks(
127+
self, tasks,
128+
batch_size=AWS_BATCH_SIZE, delay_seconds=0, total=None
129+
):
130+
131+
def genbatches(itr):
132+
while True:
133+
batch = []
134+
try:
135+
for i in range(AWS_BATCH_SIZE):
136+
batch.append(next(itr))
137+
except StopIteration:
138+
pass
139+
140+
if len(batch) == 0:
141+
break
142+
143+
yield batch
144+
145+
bodies = (
146+
{
147+
"payload": task.payload(),
148+
"queueName": self._queue_name,
149+
"groupByTag": True,
150+
"tag": task.__class__.__name__
151+
}
152+
153+
for task in tasks
154+
)
155+
156+
def cloud_insertion(batch):
157+
self._api.insert(batch, delay_seconds)
158+
159+
return ( partial(cloud_insertion, batch) for batch in genbatches(bodies) )
160+
112161
def renew_lease(self, task, seconds):
113162
"""Update the duration of a task lease."""
114163
return self._api.renew_lease(task, seconds)
@@ -301,44 +350,12 @@ def cloud_insertion(api):
301350

302351
return self
303352

304-
def insert_all(self, tasks, delay_seconds=0, total=None):
305-
if total is None:
306-
try:
307-
total = len(tasks)
308-
except TypeError:
309-
pass
310-
311-
batch_size = 10
312-
def genbatches(itr):
313-
while True:
314-
batch = []
315-
try:
316-
for i in range(batch_size):
317-
batch.append(next(itr))
318-
except StopIteration:
319-
pass
320-
321-
if len(batch) == 0:
322-
raise StopIteration
323-
324-
yield batch
325-
326-
bodies = (
327-
{
328-
"payload": task.payload(),
329-
"queueName": self._queue_name,
330-
"groupByTag": True,
331-
"tag": task.__class__.__name__
332-
}
333-
334-
for task in tasks
353+
def _insert_all(self, tasks, delay_seconds=0, total=None):
354+
batch_size = AWS_BATCH_SIZE
355+
fns = self._gen_insert_all_tasks(
356+
tasks, batch_size, delay_seconds, total
335357
)
336358

337-
def cloud_insertion(batch):
338-
self._api.insert(batch, delay_seconds)
339-
340-
fns = ( partial(cloud_insertion, batch) for batch in genbatches(bodies) )
341-
342359
schedule_threaded_jobs(
343360
fns=fns,
344361
concurrency=20,
@@ -407,15 +424,16 @@ def check_monkey_patch_status(self):
407424
import gevent.monkey
408425
if not gevent.monkey.is_module_patched("socket"):
409426
print(yellow("""
410-
Switching ThreadedQueue to using green threads. This requires
427+
GreenTaskQueue uses green threads. This requires
411428
monkey patching the standard library to use a cooperative
412429
non-blocking threading model.
413430
414431
Please place the following lines at the beginning of your
415-
program:
432+
program. `thread=False` is there because sometimes this
433+
causes hanging in multiprocessing.
416434
417435
import gevent.monkey
418-
gevent.monkey.patch_all()
436+
gevent.monkey.patch_all(thread=False)
419437
"""))
420438

421439
def insert(self, task, args=[], kwargs={}, delay_seconds=0):
@@ -436,44 +454,12 @@ def cloud_insertion():
436454

437455
return self
438456

439-
def insert_all(self, tasks, delay_seconds=0, total=None):
440-
if total is None:
441-
try:
442-
total = len(tasks)
443-
except TypeError:
444-
pass
445-
446-
batch_size = 10
447-
def genbatches(itr):
448-
while True:
449-
batch = []
450-
try:
451-
for i in range(batch_size):
452-
batch.append(next(itr))
453-
except StopIteration:
454-
pass
455-
456-
if len(batch) == 0:
457-
raise StopIteration
458-
459-
yield batch
460-
461-
bodies = (
462-
{
463-
"payload": task.payload(),
464-
"queueName": self._queue_name,
465-
"groupByTag": True,
466-
"tag": task.__class__.__name__
467-
}
468-
469-
for task in tasks
457+
def _insert_all(self, tasks, delay_seconds=0, total=None):
458+
batch_size = AWS_BATCH_SIZE
459+
fns = self._gen_insert_all_tasks(
460+
tasks, batch_size, delay_seconds, total
470461
)
471462

472-
def cloud_insertion(batch):
473-
self._api.insert(batch, delay_seconds)
474-
475-
fns = ( partial(cloud_insertion, batch) for batch in genbatches(bodies) )
476-
477463
schedule_green_jobs(
478464
fns=fns,
479465
concurrency=20,
@@ -548,7 +534,11 @@ def insert(self, task, args=[], kwargs={}):
548534
task.execute(*args, **kwargs)
549535
del task
550536

551-
def insert_all(self, tasks, args=[], kwargs={}, delay_seconds=0, total=None):
537+
def insert_all(
538+
self, tasks, args=[], kwargs={},
539+
delay_seconds=0, total=None,
540+
parallel=1
541+
):
552542
for task in tasks:
553543
self.insert(task, args=args, kwargs=kwargs)
554544

@@ -587,11 +577,18 @@ def insert(self, task, args=[], kwargs={}):
587577

588578
def insert_all(
589579
self, tasks, args=[], kwargs={},
590-
delay_seconds=0, total=None
580+
delay_seconds=0, total=None,
581+
parallel=None, progress=True
591582
):
583+
584+
if parallel is None:
585+
parallel = self.parallel
586+
592587
for task in tasks:
593588
self.queue.append( (task, args, kwargs) )
594589

590+
self._process(progress)
591+
595592
def wait(self, progress=None):
596593
self._process(progress)
597594
return self
@@ -609,7 +606,7 @@ def _process(self, progress=True):
609606
_task_execute(self.queue.pop(0))
610607
pbar.update()
611608
else:
612-
with concurrent.futures.ProcessPoolExecutor(max_workers=self.parallel) as executor:
609+
with pathos.pools.ProcessPool(self.parallel) as executor:
613610
for _ in executor.map(_task_execute, self.queue):
614611
pbar.update()
615612

@@ -632,38 +629,51 @@ def _task_execute(task_tuple):
632629

633630
def _scatter(sequence, n):
634631
"""Scatters elements of ``sequence`` into ``n`` blocks."""
635-
636-
output = []
637632
chunklen = int(math.ceil(float(len(sequence)) / float(n)))
638633

639-
for i in range(n):
640-
output.append(
641-
copy.deepcopy(sequence[i*chunklen:(i+1)*chunklen])
642-
)
634+
return [
635+
sequence[ i*chunklen : (i+1)*chunklen ] for i in range(n)
636+
]
643637

644-
return output
638+
def soloprocess_upload(QueueClass, queue_name, tasks):
639+
with QueueClass(queue_name) as tq:
640+
tq.insert_all(tasks)
645641

646-
def _upload_tasks_single_process(queue_name, tasks):
647-
with TaskQueue(queue_name=queue_name, queue_server='sqs') as tq:
648-
for task in tasks:
649-
tq.insert(task)
642+
error_queue = mp.Queue()
650643

651-
def upload(queue_name, tasks, parallel=16, progress=False):
652-
if parallel == 0:
644+
def multiprocess_upload(QueueClass, queue_name, tasks, parallel=True):
645+
if parallel is True:
653646
parallel = mp.cpu_count()
654-
elif parallel < 0:
647+
elif parallel <= 0:
655648
raise ValueError("Parallel must be a positive number or zero (all cpus). Got: " + str(parallel))
656649

657-
parallel = max(1, min(parallel, mp.cpu_count()))
658-
659650
if parallel == 1:
660-
_upload_tasks_single_process(queue_name, tasks)
661-
else:
662-
uploadfn = partial(_upload_tasks_single_process, queue_name)
663-
tasks = _scatter(tasks, parallel)
651+
soloprocess_upload(QueueClass, queue_name, tasks)
652+
return
664653

665-
with concurrent.futures.ProcessPoolExecutor(max_workers=parallel) as executor:
666-
executor.map(uploadfn, tasks)
654+
def capturing_soloprocess_upload(*args, **kwargs):
655+
try:
656+
soloprocess_upload(*args, **kwargs)
657+
except Exception as err:
658+
print(err)
659+
error_queue.put(err)
660+
661+
uploadfn = partial(
662+
capturing_soloprocess_upload, QueueClass, queue_name
663+
)
664+
tasks = _scatter(tasks, parallel)
665+
666+
pool = pathos.pools.ProcessPool(parallel)
667+
pool.map(uploadfn, tasks)
668+
669+
if not error_queue.empty():
670+
errors = []
671+
while not error_queue.empty():
672+
err = error_queue.get()
673+
if err is not StopIteration:
674+
errors.append(err)
675+
if len(errors):
676+
raise Exception(errors)
667677

668678

669679

0 commit comments

Comments
 (0)