start mp implementation

Author: RichardZhu2

src/pyper/_core/async_helper/output.py

@@ -15,16 +15,16 @@ class AsyncPipelineOutput:
     def __init__(self, pipeline: AsyncPipeline):
         self.pipeline = pipeline
 
-    def _get_q_out(self, tg: asyncio.TaskGroup, tp: ThreadPool ,*args, **kwargs) -> asyncio.Queue:
+    def _get_q_out(self, tg: asyncio.TaskGroup, tp: ThreadPool, pp: ProcessPool, *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]):
             n_consumers = 1 if next_task is None else next_task.concurrency
             if q_out is None:
-                stage = AsyncProducer(task=self.pipeline.tasks[0], tg=tg, tp=tp, n_consumers=n_consumers)
+                stage = AsyncProducer(task=self.pipeline.tasks[0], tg=tg, tp=tp, pp=pp, n_consumers=n_consumers)
                 stage.start(*args, **kwargs)
             else:
-                stage = AsyncProducerConsumer(q_in=q_out, task=task, tg=tg, tp=tp, n_consumers=n_consumers)
+                stage = AsyncProducerConsumer(q_in=q_out, task=task, tg=tg, tp=tp, pp=pp, n_consumers=n_consumers)
                 stage.start()
             q_out = stage.q_out
 
@@ -33,11 +33,12 @@ def _get_q_out(self, tg: asyncio.TaskGroup, tp: ThreadPool ,*args, **kwargs) ->
     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:
-            # 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 asyncio.TaskGroup() as tg, ThreadPool() as tp:
-                q_out = self._get_q_out(tg, tp, *args, **kwargs)
+            # 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)
                 while (data := await q_out.get()) is not StopSentinel:
                     yield data
         except ExceptionGroup as eg:

src/pyper/_core/async_helper/stage.py

@@ -13,7 +13,13 @@
 
 
 class AsyncProducer:
-    def __init__(self, task: Task, tg: asyncio.TaskGroup, tp: ThreadPool, n_consumers: int):
+    def __init__(
+            self,
+            task: Task,
+            tg: asyncio.TaskGroup,
+            tp: ThreadPool,
+            pp,
+            n_consumers: int):
         self.task = ascynchronize(task, tp)
         if task.concurrency > 1:
             raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot have concurrency greater than 1")
@@ -36,7 +42,14 @@ def start(self, *args, **kwargs):
 
 
 class AsyncProducerConsumer:
-    def __init__(self, q_in: asyncio.Queue, task: Task, tg: asyncio.TaskGroup, tp: ThreadPool, n_consumers: int):
+    def __init__(
+            self,
+            q_in: asyncio.Queue,
+            task: Task,
+            tg: asyncio.TaskGroup,
+            tp: ThreadPool,
+            pp,
+            n_consumers: int):
         self.q_in = q_in
         self.task = ascynchronize(task, tp)
         self.tg = tg

src/pyper/_core/decorators.py

@@ -25,84 +25,84 @@ class task:
     """
     @t.overload
     def __new__(
-        cls,
-        func: None = None,
-        /,
-        *,
-        join: bool = False,
-        concurrency: int = 1,
-        throttle: int = 0,
-        daemon: bool = False,
-        bind: _ArgsKwargs = None
-    ) -> t.Type[task]: ...
+            cls,
+            func: None = None,
+            /,
+            *,
+            join: bool = False,
+            concurrency: int = 1,
+            throttle: int = 0,
+            daemon: bool = False,
+            multiprocess: bool = False,
+            bind: _ArgsKwargs = None) -> t.Type[task]: ...
 
     @t.overload
     def __new__(
-        cls,
-        func: t.Callable[_P, t.Awaitable[_R]],
-        /,
-        *,
-        join: bool = False,
-        concurrency: int = 1,
-        throttle: int = 0,
-        daemon: bool = False,
-        bind: _ArgsKwargs = None
-    ) -> AsyncPipeline[_P, _R]: ...
+            cls,
+            func: t.Callable[_P, t.Awaitable[_R]],
+            /,
+            *,
+            join: bool = False,
+            concurrency: int = 1,
+            throttle: int = 0,
+            daemon: bool = False,
+            multiprocess: bool = False,
+            bind: _ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
 
     @t.overload
     def __new__(
-        cls,
-        func: t.Callable[_P, t.AsyncGenerator[_R]],
-        /,
-        *,
-        join: bool = False,
-        concurrency: int = 1,
-        throttle: int = 0,
-        daemon: bool = False,
-        bind: _ArgsKwargs = None
-    ) -> AsyncPipeline[_P, _R]: ...
-    
+            cls,
+            func: t.Callable[_P, t.AsyncGenerator[_R]],
+            /,
+            *,
+            join: bool = False,
+            concurrency: int = 1,
+            throttle: int = 0,
+            daemon: bool = False,
+            multiprocess: bool = False,
+            bind: _ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
+        
     @t.overload
     def __new__(
-        cls,
-        func: t.Callable[_P, t.Generator[_R]],
-        /,
-        *,
-        join: bool = False,
-        concurrency: int = 1,
-        throttle: int = 0,
-        daemon: bool = False,
-        bind: _ArgsKwargs = None
-    ) -> Pipeline[_P, _R]: ...
+            cls,
+            func: t.Callable[_P, t.Generator[_R]],
+            /,
+            *,
+            join: bool = False,
+            concurrency: int = 1,
+            throttle: int = 0,
+            daemon: bool = False,
+            multiprocess: bool = False,
+            bind: _ArgsKwargs = None) -> Pipeline[_P, _R]: ...
 
     @t.overload
     def __new__(
-        cls,
-        func: t.Callable[_P, _R],
-        /,
-        *,
-        join: bool = False,
-        concurrency: int = 1,
-        throttle: int = 0,
-        daemon: bool = False,
-        bind: _ArgsKwargs = None
-    ) -> Pipeline[_P, _R]: ...
+            cls,
+            func: t.Callable[_P, _R],
+            /,
+            *,
+            join: bool = False,
+            concurrency: int = 1,
+            throttle: int = 0,
+            daemon: bool = False,
+            multiprocess: bool = False,
+            bind: _ArgsKwargs = None) -> Pipeline[_P, _R]: ...
 
     def __new__(
-        cls,
-        func: t.Optional[t.Callable] = None,
-        /,
-        *,
-        join: bool = False,
-        concurrency: int = 1,
-        throttle: int = 0,
-        daemon: bool = False,
-        bind: _ArgsKwargs = None
-    ):
+            cls,
+            func: t.Optional[t.Callable] = None,
+            /,
+            *,
+            join: bool = False,
+            concurrency: int = 1,
+            throttle: int = 0,
+            daemon: bool = False,
+            multiprocess: bool = False,
+            bind: _ArgsKwargs = None):
         # Classic decorator trick: @task() means func is None, @task without parentheses means func is passed. 
         if func is None:
-            return functools.partial(cls, join=join, concurrency=concurrency, throttle=throttle, daemon=daemon, bind=bind)
-        return Pipeline([Task(func=func, join=join, concurrency=concurrency, throttle=throttle, daemon=daemon, bind=bind)])
+            return functools.partial(cls, join=join, concurrency=concurrency, throttle=throttle, daemon=daemon, multiprocess=multiprocess, bind=bind)
+        return Pipeline([Task(func=func, join=join, concurrency=concurrency, throttle=throttle, daemon=daemon, multiprocess=multiprocess, bind=bind)])
 
     @staticmethod
     def bind(*args, **kwargs) -> _ArgsKwargs:

src/pyper/_core/task.py

@@ -8,17 +8,26 @@
 class Task:
     """The representation of a function within a Pipeline."""
 
-    __slots__ = "is_gen", "is_async", "func", "join", "concurrency", "throttle", "daemon"
+    __slots__ = (
+        "is_gen",
+        "is_async",
+        "func",
+        "join",
+        "concurrency",
+        "throttle",
+        "daemon",
+        "multiprocess"
+    )
 
     def __init__(
-        self,
-        func: Callable,
-        join: bool = False,
-        concurrency: int = 1,
-        throttle: int = 0,
-        daemon: bool = False,
-        bind: Optional[Tuple[Tuple, Dict]] = None
-    ):
+            self,
+            func: Callable,
+            join: bool = False,
+            concurrency: int = 1,
+            throttle: int = 0,
+            daemon: bool = False,
+            multiprocess: bool = False,
+            bind: Optional[Tuple[Tuple, Dict]] = None):
         if not isinstance(concurrency, int):
             raise TypeError("concurrency must be an integer")
         if concurrency < 1:
@@ -29,6 +38,8 @@ def __init__(
             raise ValueError("throttle cannot be less than 0")
         if not callable(func):
             raise TypeError("A task must be a callable object")
+        if daemon and multiprocess:
+            raise ValueError("daemon and multiprocess cannot both be True")
 
         self.is_gen = inspect.isgeneratorfunction(func) \
             or inspect.isasyncgenfunction(func) \
@@ -41,9 +52,12 @@ def __init__(
 
         if self.is_async and daemon:
             raise ValueError("daemon cannot be True for an async task")
+        if self.is_async and multiprocess:
+            raise ValueError("multiprocess cannot be True for an async task")
 
         self.func = func if bind is None else functools.partial(func, *bind[0], **bind[1])
         self.join = join
         self.concurrency = concurrency
         self.throttle = throttle
         self.daemon = daemon
+        self.multiprocess = multiprocess

src/pyper/_core/util/asynchronize.py

@@ -8,20 +8,32 @@
 from ..task import Task
 
 
-def ascynchronize(task: Task, tp: ThreadPool) -> Task:
-    """Unifies async and sync functions within an `AsyncPipeline`.
-
-    Synchronous functions within a `ThreadPool` are transformed into asynchronous functions via `asyncio.wrap_future`.
-    Synchronous generators are transformed into asynchronous generators.
+def ascynchronize(task: Task, tp: ThreadPool, pp) -> 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`.
+    4. Threaded synchronous functions within a `ThreadPool` are wrapped in `asyncio.wrap_future`.
     """
     if task.is_async:
-        return task
-    
+        return task  
     if task.is_gen:
         @functools.wraps(task.func)
         async def wrapper(*args, **kwargs):
             for output in task.func(*args, **kwargs):
                 yield output
+    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)
     else:
         @functools.wraps(task.func)
         async def wrapper(*args, **kwargs):
@@ -34,7 +46,6 @@ def target(*args, **kwargs):
                     future.set_exception(e)
             tp.submit(target, args=args, kwargs=kwargs, daemon=task.daemon)
             return await asyncio.wrap_future(future)
-    
     return Task(
         func=wrapper,
         join=task.join,