implement async api for multiprocessing

Author: RichardZhu2

src/pyper/_core/async_helper/output.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+from concurrent.futures import ProcessPoolExecutor
 import sys
 from typing import TYPE_CHECKING
 
@@ -21,7 +22,7 @@ class AsyncPipelineOutput:
     def __init__(self, pipeline: AsyncPipeline):
         self.pipeline = pipeline
 
-    def _get_q_out(self, tg: TaskGroup, tp: ThreadPool, pp: ProcessPool, *args, **kwargs) -> asyncio.Queue:
+    def _get_q_out(self, tg: TaskGroup, tp: ThreadPool, pp: ProcessPoolExecutor, *args, **kwargs) -> asyncio.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]):
@@ -39,22 +40,15 @@ def _get_q_out(self, tg: TaskGroup, tp: ThreadPool, pp: ProcessPool, *args, **kw
     async def __call__(self, *args, **kwargs):
         """Call the pipeline, taking the inputs to the first task, and returning the output from the last task."""
         try:
-<<<<<<< HEAD
             # Unify async, threaded, and multiprocessed work by:
             # 1. using TaskGroup to execute asynchronous tasks
             # 2. using ThreadPool to execute threaded synchronous tasks
-            # 3. using ProcessPool to execute multiprocessed synchronous tasks
-            async with asyncio.TaskGroup() as tg, ThreadPool() as tp, ProcessPool as pp:
-                q_out = self._get_q_out(tg, tp, pp, *args, **kwargs)
-=======
-            # We unify async and thread-based concurrency by
-            # 1. using TaskGroup to spin up asynchronous tasks
-            # 2. using ThreadPool to spin up synchronous tasks
-            async with TaskGroup() as tg, ThreadPool() as tp:
-                q_out = self._get_q_out(tg, tp, *args, **kwargs)
->>>>>>> dev
-                while (data := await q_out.get()) is not StopSentinel:
-                    yield data
+            # 3. using ProcessPoolExecutor to execute multiprocessed synchronous tasks
+            async with TaskGroup() as tg:
+                with ThreadPool() as tp, ProcessPoolExecutor() as pp:
+                    q_out = self._get_q_out(tg, tp, pp, *args, **kwargs)
+                    while (data := await q_out.get()) is not StopSentinel:
+                        yield data
         except ExceptionGroup as eg:
-            raise eg.exceptions[0]
+            raise eg.exceptions[0] from None
 

src/pyper/_core/async_helper/stage.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+from concurrent.futures import ProcessPoolExecutor
 import sys
 from typing import TYPE_CHECKING
 
@@ -24,9 +25,9 @@ def __init__(
             task: Task,
             tg: TaskGroup,
             tp: ThreadPool,
-            pp,
+            pp: ProcessPoolExecutor,
             n_consumers: int):
-        self.task = ascynchronize(task, tp)
+        self.task = ascynchronize(task, tp, pp)
         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:
@@ -54,10 +55,10 @@ def __init__(
             task: Task,
             tg: TaskGroup,
             tp: ThreadPool,
-            pp,
+            pp: ProcessPoolExecutor,
             n_consumers: int):
         self.q_in = q_in
-        self.task = ascynchronize(task, tp)
+        self.task = ascynchronize(task, tp, pp)
         self.tg = tg
         self.n_consumers = n_consumers
         self.q_out = asyncio.Queue(maxsize=task.throttle)

src/pyper/_core/util/asynchronize.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
 
 import asyncio
-from concurrent.futures import Future
+from concurrent.futures import ProcessPoolExecutor, Future
 import functools
 
 from .thread_pool import ThreadPool
 from ..task import Task
 
 
-def ascynchronize(task: Task, tp: ThreadPool, pp) -> Task:
+def ascynchronize(task: Task, tp: ThreadPool, pp: ProcessPoolExecutor) -> Task:
     """Unifies async and sync tasks as awaitable futures.
     1. If the task is async already, return it.
     2. Synchronous generators are transformed into asynchronous generators.
-    3. Multiprocessed synchronous functions within a `ProcessPool` are wrapped in `asyncio.wrap_future`.
+    3. Multiprocessed synchronous functions within a `ProcessPoolExecutor` are wrapped in `asyncio.wrap_future`.
     4. Threaded synchronous functions within a `ThreadPool` are wrapped in `asyncio.wrap_future`.
     """
     if task.is_async:
@@ -25,15 +25,9 @@ async def wrapper(*args, **kwargs):
     elif task.multiprocess:
         @functools.wraps(task.func)
         async def wrapper(*args, **kwargs):
-            future = Future()
-            def target(*args, **kwargs):
-                try:
-                    result = task.func(*args, **kwargs)
-                    future.set_result(result)
-                except Exception as e:
-                    future.set_exception(e)
-            pp.submit(target, args=args, kwargs=kwargs)
-            return await asyncio.wrap_future(future)
+            loop = asyncio.get_running_loop()
+            f = functools.partial(task.func, *args, **kwargs)
+            return await loop.run_in_executor(executor=pp, func=f)
     else:
         @functools.wraps(task.func)
         async def wrapper(*args, **kwargs):

src/pyper/_core/util/thread_pool.py

@@ -19,13 +19,6 @@ def __exit__(self, ev, et, tb):
         for thread in self._threads:
             thread.daemon or thread.join()
 
-    async def __aenter__(self):
-        return self
-    
-    async def __aexit__(self, ev, et, tb):
-        for thread in self._threads:
-            thread.daemon or thread.join()
-
     @property
     def has_error(self):
         return not self._error_queue.empty()