tentative sync refactor to poolexecutors

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, shutdown_event=pool.shutdown_event)
+                stage = Producer(task=task, next_task=next_task, manager=pp.manager, 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, shutdown_event=pool.shutdown_event)
+                stage = ProducerConsumer(q_in=q_out, task=task, next_task=next_task, manager=pp.manager, shutdown_event=pool.shutdown_event)
                 stage.start(pool)
             q_out = stage.q_out
 
@@ -37,15 +37,8 @@ def __call__(self, *args, **kwargs):
             while True:
                 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()
+                    if (data := q_out.get()) is StopSentinel:
                         break
                     yield data
-                except queue.Empty:
-                    tp.raise_error_if_exists()
-                    pp.raise_error_if_exists()
                 except (KeyboardInterrupt, SystemExit): # pragma: no cover
-                    tp.shutdown_event.set()
-                    pp.shutdown_event.set()
                     raise

src/pyper/_core/sync_helper/stage.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
-import multiprocessing as mp
 import queue
+import threading
 from typing import TYPE_CHECKING, Union
 
 from .queue_io import DequeueFactory, EnqueueFactory
 from ..util.sentinel import StopSentinel
-from ..util.value import ThreadingValue
+from ..util.counter import MultiprocessingCounter, ThreadingCounter
 
 if TYPE_CHECKING:
-    from multiprocessing.synchronize import Event as ProcessEvent
-    from threading import Event as ThreadEvent
+    from multiprocessing.managers import SyncManager
+    import multiprocessing.queues as mpq
+    import multiprocessing.synchronize as mpsync
     from ..util.worker_pool import WorkerPool
     from ..task import Task
 
@@ -20,17 +21,16 @@ def __init__(
             self,
             task: Task,
             next_task: Task,
-            q_err: Union[mp.Queue, queue.Queue],
-            shutdown_event: Union[ProcessEvent, ThreadEvent]):
+            manager: SyncManager,
+            shutdown_event: Union[mpsync.Event, threading.Event]):
         if task.workers > 1:
             raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot have more than 1 worker")
         if task.join:
             raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot join previous results")
-        self.q_out = mp.Queue(maxsize=task.throttle) \
+        self.q_out = manager.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.workers
         self._n_consumers = 1 if next_task is None else next_task.workers
@@ -39,9 +39,9 @@ def __init__(
     def _worker(self, *args, **kwargs):
         try:
             self._enqueue(*args, **kwargs)
-        except Exception as e:
+        except Exception:
             self._shutdown_event.set()
-            self._q_err.put(e)
+            raise
         finally:
             for _ in range(self._n_consumers):
                 self.q_out.put(StopSentinel)
@@ -53,45 +53,37 @@ def start(self, pool: WorkerPool, /, *args, **kwargs):
 class ProducerConsumer:
     def __init__(
             self,
-            q_in: Union[mp.Queue, queue.Queue],
+            q_in: Union[mpq.Queue, queue.Queue],
             task: Task,
             next_task: Task,
-            q_err: Union[mp.Queue, queue.Queue],
-            shutdown_event: Union[ProcessEvent, ThreadEvent]):
+            manager: SyncManager,
+            shutdown_event: Union[mpsync.Event, 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) \
+        # and only using a multiprocess Queue when the current task or the next task are multiprocessed
+        self.q_out = manager.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.workers
         self._n_consumers = 1 if next_task is None else next_task.workers
         self._dequeue = DequeueFactory(q_in, task)
         self._enqueue = EnqueueFactory(self.q_out, task)
-        self._workers_done = mp.Value('i', 0) if task.multiprocess else ThreadingValue(0)
-
-    def _increment_workers_done(self):
-        with self._workers_done.get_lock():
-            self._workers_done.value += 1
-            return self._workers_done.value
-
-    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 = MultiprocessingCounter(0, manager) if task.multiprocess else ThreadingCounter(0)
 
     def _worker(self):
         try:
             for output in self._dequeue():
                 if not self._shutdown_event.is_set():
                     self._enqueue(output)
-        except Exception as e:
+        except Exception:
             self._shutdown_event.set()
-            self._q_err.put(e)
+            raise
         finally:
-            self._finish()
+            if self._workers_done.increment() == self._n_workers:
+                for _ in range(self._n_consumers):
+                    self.q_out.put(StopSentinel)
+
 
     def start(self, pool: WorkerPool, /):
         for _ in range(self._n_workers):

src/pyper/_core/util/counter.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from multiprocessing.managers import SyncManager
+import threading
+
+
+class MultiprocessingCounter:
+    """Utility class to manage process based access to an integer."""
+    def __init__(self, value: int, manager: SyncManager):
+        self._manager_value = manager.Value('i', value)
+        self._lock = manager.Lock()
+    
+    def increment(self):
+        with self._lock:
+            self._manager_value.value += 1
+            return self._manager_value.value
+
+
+class ThreadingCounter:
+    """Utility class to manage thread based access to an integer."""
+    def __init__(self, value: int):
+        self._value = value
+        self._lock = threading.Lock()
+    
+    def increment(self):
+        with self._lock:
+            self._value += 1
+            return self._value

src/pyper/_core/util/value.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import concurrent.futures as cf
 import multiprocessing as mp
-import queue
+import multiprocessing.synchronize as mpsync
 import threading
 from typing import List, Union
 
@@ -13,41 +14,40 @@ 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 = None
-
-    def __init__(self):
-        self.error_queue = mp.Queue() if self.worker_type is mp.Process else queue.Queue()
-        self.shutdown_event = mp.Event() if self.worker_type is mp.Process else threading.Event()
-
-        self._workers: List[Union[mp.Process, threading.Thread]] = []
+    shutdown_event: Union[mpsync.Event, threading.Event]
+    _executor: Union[cf.ProcessPoolExecutor, cf.ThreadPoolExecutor]
+    _futures: List[cf.Future]
 
     def __enter__(self):
+        self._executor.__enter__()
         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()
+        self._executor.__exit__(et, ev, tb)
+        for future in self._futures:
+            future.result()
 
-    def raise_error_if_exists(self):
-        if self.has_error:
-            raise self.get_error() from None
+    def create_queue(self):
+        raise NotImplementedError
 
     def submit(self, func, /, *args, **kwargs):
-        w = self.worker_type(target=func, args=args, kwargs=kwargs)
-        w.start()
-        self._workers.append(w)
+        future = self._executor.submit(func, *args, **kwargs)
+        self._futures.append(future)
+        return future
 
 
 class ThreadPool(WorkerPool):
-    worker_type = threading.Thread
+    def __init__(self):
+        self.shutdown_event = threading.Event()
+
+        self._executor = cf.ThreadPoolExecutor()
+        self._futures = []
 
 
 class ProcessPool(WorkerPool):
-    worker_type = mp.Process
+    def __init__(self):
+        self.manager = mp.Manager()
+        self.shutdown_event = self.manager.Event()
+
+        self._executor = cf.ProcessPoolExecutor()
+        self._futures = []