add multiprocessing functionality for sync api

Author: RichardZhu2

src/pyper/_core/sync_helper/output.py

@@ -5,7 +5,7 @@
 
 from .stage import Producer, ProducerConsumer
 from ..util.sentinel import StopSentinel
-from ..util.thread_pool import ThreadPool
+from ..util.worker_pool import ProcessPool, ThreadPool
 
 if TYPE_CHECKING:
     from ..pipeline import Pipeline
@@ -15,32 +15,33 @@ class PipelineOutput:
     def __init__(self, pipeline: Pipeline):
         self.pipeline = pipeline
 
-    def _get_q_out(self, tp: ThreadPool, *args, **kwargs) -> queue.Queue:
+    def _get_q_out(self, tp: ThreadPool, pp: ProcessPool, *args, **kwargs) -> queue.Queue:
         """Feed forward each stage to the next, returning the output queue of the final stage."""
         q_out = None
         for task, next_task in zip(self.pipeline.tasks, self.pipeline.tasks[1:] + [None]):
-            n_consumers = 1 if next_task is None else next_task.concurrency
+            pool = pp if task.multiprocess else tp
             if q_out is None:
-                stage = Producer(task=self.pipeline.tasks[0], tp=tp, n_consumers=n_consumers)
-                stage.start(*args, **kwargs)
+                stage = Producer(task=task, next_task=next_task, q_err=pool.error_queue)
+                stage.start(pool, *args, **kwargs)
             else:
-                stage = ProducerConsumer(q_in=q_out, task=task, tp=tp, n_consumers=n_consumers)
-                stage.start()
+                stage = ProducerConsumer(q_in=q_out, task=task, next_task=next_task, q_err=pool.error_queue)
+                stage.start(pool)
             q_out = stage.q_out
 
         return q_out
 
     def __call__(self, *args, **kwargs):
-        """Call the pipeline, taking the inputs to the first task, and returning the output from the last task."""
-        with ThreadPool() as tp:
-            q_out = self._get_q_out(tp, *args, **kwargs)
+        """Iterate through the pipeline, taking the inputs to the first task, and yielding each output from the last task."""
+        with ThreadPool() as tp, ProcessPool() as pp:
+            q_out = self._get_q_out(tp, pp, *args, **kwargs)
             while True:
-                tp.raise_error_if_exists()
                 try:
                     # Use the timeout strategy for unblocking main thread without busy waiting
                     if (data := q_out.get(timeout=0.1)) is StopSentinel:
                         tp.raise_error_if_exists()
+                        pp.raise_error_if_exists()
                         break
                     yield data
                 except queue.Empty:
-                    continue
+                    tp.raise_error_if_exists()
+                    pp.raise_error_if_exists()

src/pyper/_core/sync_helper/queue_io.py

@@ -1,27 +1,24 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 from ..util.sentinel import StopSentinel
 
 if TYPE_CHECKING:
+    import multiprocessing as mp
     import queue
     from ..task import Task
 
 
-class Dequeue:
+def DequeueFactory(q_in: Union[mp.Queue, queue.Queue], task: Task):
+    return _JoiningDequeue(q_in=q_in) if task.join \
+        else _SingleDequeue(q_in=q_in)
+
+
+class _Dequeue:
     """Pulls data from an input queue."""
-    def __new__(self, q_in: queue.Queue, task: Task):
-        if task.join:
-            instance = object.__new__(_JoiningDequeue)
-        else:
-            instance = object.__new__(_SingleDequeue)
-        instance.__init__(q_in=q_in, task=task)
-        return instance
-
-    def __init__(self, q_in: queue.Queue, task: Task):
-         self.q_in = q_in
-         self.task = task
+    def __init__(self, q_in: Union[mp.Queue, queue.Queue]):
+        self.q_in = q_in
 
     def _input_stream(self):
         while (data := self.q_in.get()) is not StopSentinel:
@@ -31,41 +28,38 @@ def __call__(self):
         raise NotImplementedError
 
 
-class _SingleDequeue(Dequeue):
+class _SingleDequeue(_Dequeue):
      def __call__(self):
          for data in self._input_stream():
             yield data
 
 
-class _JoiningDequeue(Dequeue):
+class _JoiningDequeue(_Dequeue):
      def __call__(self):
         yield self._input_stream()
 
 
-class Enqueue:
+def EnqueueFactory(q_out: Union[mp.Queue, queue.Queue], task: Task):
+    return _BranchingEnqueue(q_out=q_out, task=task) if task.is_gen \
+        else _SingleEnqueue(q_out=q_out, task=task)
+
+
+class _Enqueue:
     """Puts output from a task onto an output queue."""
-    def __new__(cls, q_out: queue.Queue, task: Task):
-        if task.is_gen:
-            instance = object.__new__(_BranchingEnqueue)
-        else:
-            instance = object.__new__(_SingleEnqueue)
-        instance.__init__(q_out=q_out, task=task)
-        return instance
-
-    def __init__(self, q_out: queue.Queue, task: Task):
-         self.q_out = q_out
-         self.task = task
+    def __init__(self, q_out: Union[mp.Queue, queue.Queue], task: Task):
+        self.q_out = q_out
+        self.task = task
 
     def __call__(self, *args, **kwargs):
         raise NotImplementedError
 
 
-class _SingleEnqueue(Enqueue):        
+class _SingleEnqueue(_Enqueue):        
      def __call__(self, *args, **kwargs):
         self.q_out.put(self.task.func(*args, **kwargs))
 
 
-class _BranchingEnqueue(Enqueue):
+class _BranchingEnqueue(_Enqueue):
      def __call__(self, *args, **kwargs):
-         for output in self.task.func(*args, **kwargs):
+        for output in self.task.func(*args, **kwargs):
             self.q_out.put(output)

src/pyper/_core/sync_helper/stage.py

@@ -1,69 +1,91 @@
 from __future__ import annotations
 
+import multiprocessing as mp
 import queue
 import threading
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
-from .queue_io import Dequeue, Enqueue
+from .queue_io import DequeueFactory, EnqueueFactory
 from ..util.sentinel import StopSentinel
 
 if TYPE_CHECKING:
-    from ..util.thread_pool import ThreadPool
+    from ..util.worker_pool import WorkerPool
     from ..task import Task
 
 
 class Producer:
-    def __init__(self, task: Task, tp: ThreadPool, n_consumers: int):
-        self.task = task
+    def __init__(self, task: Task, next_task: Task, q_err: Union[mp.Queue, queue.Queue]):
         if task.concurrency > 1:
             raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot have concurrency greater than 1")
         if task.join:
             raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot join previous results")
-        self.tp = tp
-        self.n_consumers = n_consumers
-        self.q_out = queue.Queue(maxsize=task.throttle)
+        self.q_out = mp.Queue(maxsize=task.throttle) \
+            if task.multiprocess or (next_task is not None and next_task.multiprocess) \
+            else queue.Queue(maxsize=task.throttle)
 
-        self._enqueue = Enqueue(self.q_out, self.task)
+        self._q_err = q_err
+        self._n_workers = task.concurrency
+        self._n_consumers = 1 if next_task is None else next_task.concurrency
+        self._enqueue = EnqueueFactory(self.q_out, task)
+
+        self._multiprocess = task.multiprocess
 
     def _worker(self, *args, **kwargs):
         try:
             self._enqueue(*args, **kwargs)
         except Exception as e:
-            self.tp.put_error(e)
+            self._q_err.put(e)
         finally:
-            for _ in range(self.n_consumers):
+            for _ in range(self._n_consumers):
                 self.q_out.put(StopSentinel)
 
-    def start(self, *args, **kwargs):
-        self.tp.submit(self._worker, args, kwargs, daemon=self.task.daemon)
+    def start(self, pool: WorkerPool, /, *args, **kwargs):
+        pool.submit(self._worker, *args, **kwargs)
 
 
 class ProducerConsumer:
-    def __init__(self, q_in: queue.Queue, task: Task, tp: ThreadPool, n_consumers: int):
-        self.q_in = q_in
-        self.task = task
-        self.tp = tp
-        self.n_consumers = n_consumers
-        self.q_out = queue.Queue(maxsize=task.throttle)
+    def __init__(self, q_in: Union[mp.Queue, queue.Queue], task: Task, next_task: Task, q_err: Union[mp.Queue, queue.Queue]):
+        self.q_out = mp.Queue(maxsize=task.throttle) \
+            if task.multiprocess or (next_task is not None and next_task.multiprocess) \
+            else queue.Queue(maxsize=task.throttle)
+        
+        self._q_err = q_err
+        self._n_workers = task.concurrency
+        self._n_consumers = 1 if next_task is None else next_task.concurrency
+        self._dequeue = DequeueFactory(q_in, task)
+        self._enqueue = EnqueueFactory(self.q_out, task)
+
+        self._multiprocess = task.multiprocess
+        if self._multiprocess:
+            self._workers_done = mp.Value('i', 0)
+            self._lock = self._workers_done.get_lock()
+        else:
+            self._workers_done = 0
+            self._lock = threading.Lock()
+
+    def _increment_workers_done(self):
+        with self._lock:
+            if self._multiprocess:
+                self._workers_done.value += 1
+                return self._workers_done.value
+            else:
+                self._workers_done += 1
+                return self._workers_done
+
+    def _finish(self):
+        if self._increment_workers_done() == self._n_workers:
+            for _ in range(self._n_consumers):
+                self.q_out.put(StopSentinel)
 
-        self._workers_done = 0
-        self._workers_done_lock = threading.Lock()
-        self._dequeue = Dequeue(self.q_in, self.task)
-        self._enqueue = Enqueue(self.q_out, self.task)
-    
     def _worker(self):
         try:
             for output in self._dequeue():
                 self._enqueue(output)
         except Exception as e:
-            self.tp.put_error(e)
+            self._q_err.put(e)
         finally:
-            with self._workers_done_lock:
-                self._workers_done += 1
-                if self._workers_done == self.task.concurrency:
-                    for _ in range(self.n_consumers):
-                        self.q_out.put(StopSentinel)
-
-    def start(self):
-        for _ in range(self.task.concurrency):
-            self.tp.submit(self._worker, daemon=self.task.daemon)
+            self._finish()
+            
+    def start(self, pool: WorkerPool, /):
+        for _ in range(self._n_workers):
+            pool.submit(self._worker)

src/pyper/_core/util/thread_pool.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import multiprocessing as mp
+import queue
+import threading
+from typing import List, Union
+
+
+class WorkerPool:
+    """A context for fine-grained thread/process management and error handling.
+
+    1. Spins up thread/process workers and maintains a reference to each
+    2. Provides a mechanism to capture and propagate errors to the main thread/process
+    3. Ensures safe tear-down of all workers 
+    """
+    worker_type: type = None
+
+    error_queue: Union[mp.Queue, queue.Queue]
+    _workers: List[Union[mp.Process, threading.Thread]]
+
+    def __init__(self):
+        self.error_queue = mp.Queue(1) if self.worker_type is mp.Process else queue.Queue(1)
+        self._workers = []
+
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, et, ev, tb):
+        for worker in self._workers:
+            worker.join()
+
+    @property
+    def has_error(self):
+        return not self.error_queue.empty()
+    
+    def get_error(self) -> Exception:
+        return self.error_queue.get()
+
+    def put_error(self, e: Exception):
+        self.error_queue.put(e)
+
+    def raise_error_if_exists(self):
+        if self.has_error:
+            raise self.get_error() from None
+
+    def submit(self, func, /, *args, **kwargs):
+        w = self.worker_type(target=func, args=args, kwargs=kwargs)
+        w.start()
+        self._workers.append(w)
+
+
+class ThreadPool(WorkerPool):
+    worker_type = threading.Thread
+
+
+class ProcessPool(WorkerPool):
+    worker_type = mp.Process