add fail fast mechanisms for sync api

Author: RichardZhu2

src/pyper/_core/sync_helper/output.py

@@ -21,10 +21,10 @@ def _get_q_out(self, tp: ThreadPool, pp: ProcessPool, *args, **kwargs) -> queue.
         for task, next_task in zip(self.pipeline.tasks, self.pipeline.tasks[1:] + [None]):
             pool = pp if task.multiprocess else tp
             if q_out is None:
-                stage = Producer(task=task, next_task=next_task, q_err=pool.error_queue)
+                stage = Producer(task=task, next_task=next_task, q_err=pool.error_queue, shutdown_event=pool.shutdown_event)
                 stage.start(pool, *args, **kwargs)
             else:
-                stage = ProducerConsumer(q_in=q_out, task=task, next_task=next_task, q_err=pool.error_queue)
+                stage = ProducerConsumer(q_in=q_out, task=task, next_task=next_task, q_err=pool.error_queue, shutdown_event=pool.shutdown_event)
                 stage.start(pool)
             q_out = stage.q_out
 
@@ -45,3 +45,7 @@ def __call__(self, *args, **kwargs):
                 except queue.Empty:
                     tp.raise_error_if_exists()
                     pp.raise_error_if_exists()
+                except (KeyboardInterrupt, SystemExit):
+                    tp.shutdown_event.set()
+                    pp.shutdown_event.set()
+                    raise

src/pyper/_core/sync_helper/stage.py

@@ -9,12 +9,18 @@
 from ..util.sentinel import StopSentinel
 
 if TYPE_CHECKING:
+    from multiprocessing.synchronize import Event as MpEvent
     from ..util.worker_pool import WorkerPool
     from ..task import Task
 
 
 class Producer:
-    def __init__(self, task: Task, next_task: Task, q_err: Union[mp.Queue, queue.Queue]):
+    def __init__(
+            self,
+            task: Task,
+            next_task: Task,
+            q_err: Union[mp.Queue, queue.Queue],
+            shutdown_event: Union[MpEvent, threading.Event]):
         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:
@@ -24,16 +30,16 @@ def __init__(self, task: Task, next_task: Task, q_err: Union[mp.Queue, queue.Que
             else queue.Queue(maxsize=task.throttle)
 
         self._q_err = q_err
+        self._shutdown_event = shutdown_event
         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._shutdown_event.set()
             self._q_err.put(e)
         finally:
             for _ in range(self._n_consumers):
@@ -44,12 +50,21 @@ def start(self, pool: WorkerPool, /, *args, **kwargs):
 
 
 class ProducerConsumer:
-    def __init__(self, q_in: Union[mp.Queue, queue.Queue], task: Task, next_task: Task, q_err: Union[mp.Queue, queue.Queue]):
+    def __init__(
+            self,
+            q_in: Union[mp.Queue, queue.Queue],
+            task: Task,
+            next_task: Task,
+            q_err: Union[mp.Queue, queue.Queue],
+            shutdown_event: Union[MpEvent, threading.Event]):
+        # The output queue is shared between this task and the next. We optimize here by using queue.Queue wherever possible
+        # and only using multiprocess.Queue when the current task or the next task are multiprocessed
         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._shutdown_event = shutdown_event
         self._n_workers = task.concurrency
         self._n_consumers = 1 if next_task is None else next_task.concurrency
         self._dequeue = DequeueFactory(q_in, task)
@@ -80,8 +95,10 @@ def _finish(self):
     def _worker(self):
         try:
             for output in self._dequeue():
-                self._enqueue(output)
+                if not self._shutdown_event.is_set():
+                    self._enqueue(output)
         except Exception as e:
+            self._shutdown_event.set()
             self._q_err.put(e)
         finally:
             self._finish()

src/pyper/_core/util/worker_pool.py

@@ -13,14 +13,13 @@ class WorkerPool:
     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]]
+    worker_type = None
 
     def __init__(self):
         self.error_queue = mp.Queue(1) if self.worker_type is mp.Process else queue.Queue(1)
-        self._workers = []
+        self.shutdown_event = mp.Event() if self.worker_type is mp.Process else threading.Event()
+
+        self._workers: List[Union[mp.Process, threading.Thread]] = []
 
     def __enter__(self):
         return self