william-silversmith · Feb 22, 2019 · Feb 21, 2019 · Feb 21, 2019 · Feb 21, 2019 · Feb 21, 2019
diff --git 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
@@ -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
@@ -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
@@ -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)