Merge pull request #15 from pyper-dev/dev

Safer multiprocessing implementation

Author: pyper-dev

docs/src/docs/ApiReference/task.md

@@ -233,6 +233,11 @@ An asynchronous task cannot set `multiprocessing` as `True`
 
 See some [considerations](../UserGuide/AdvancedConcepts#cpu-bound-work) for when to set this parameter.
 
+Note, also, that normal Python multiprocessing restrictions apply:
+
+* Only [picklable](https://docs.python.org/3/library/pickle.html#module-pickle) functions can be multiprocessed, which excludes certain types of functions like lambdas and closures.
+* Arguments and return values of multiprocessed tasks must also be picklable, which excludes objects like file handles, connections, and (on Windows) generators.
+
 {: .text-beta}
 ### `bind`
 

docs/src/docs/UserGuide/AdvancedConcepts.md

@@ -39,7 +39,7 @@ IO-bound tasks benefit from both concurrent and parallel execution.
 However, to avoid the overhead costs of creating processes, it is generally preferable to use either threading or async code.
 
 {: .info}
-Threads incur a higher overhead cost compared to async coroutines, but are suitable if your application prefers or requires a synchronous implementation
+Threads incur a higher overhead cost compared to async coroutines, but are suitable if the function / application prefers or requires a synchronous implementation
 
 Note that asynchronous functions need to `await` or `yield` something in order to benefit from concurrency.
 Any long-running call in an async task which does not yield execution will prevent other tasks from making progress:
@@ -87,7 +87,7 @@ def long_computation(data: int):
     return data
 ```
 
-Note, however, that processes incur a very high overhead cost (performance in creation and memory in maintaining inter-process communication). Specific cases should be benchmarked to fine-tune the task parameters for your program / your machine.
+Note, however, that processes incur a very high overhead cost (performance cost in creation and memory cost in inter-process communication). Specific cases should be benchmarked to fine-tune the task parameters for your program / your machine.
 
 ### Summary
 
@@ -101,7 +101,7 @@ Note, however, that processes incur a very high overhead cost (performance in cr
 {: .text-green-200}
 **Key Considerations:**
 
-* If a task is doing extremely expensive CPU-bound work, define it synchronously and set `multiprocess=True`
+* If a task is doing expensive CPU-bound work, define it synchronously and set `multiprocess=True`
 * If a task is doing expensive IO-bound work, consider implementing it asynchronously, or use threads
 * Do _not_ put expensive, blocking work in an async task, as this clogs up the async event loop
 
@@ -111,7 +111,7 @@ Note, however, that processes incur a very high overhead cost (performance in cr
 
 Writing clean code is partly about defining functions with single, clear responsibilities.
 
-In Pyper specifically, it is especially important to separate out different types of work into different tasks if we want to optimize their performance. For example, consider a task which performs an IO-bound network request along with a CPU-bound function to parse the data.
+In Pyper, it is especially important to separate out different types of work into different tasks if we want to optimize their performance. For example, consider a task which performs an IO-bound network request along with a CPU-bound function to parse the data.
 
 ```python
 # Bad -- functions not separated
@@ -165,10 +165,10 @@ When defining a pipeline, these additional arguments are plugged into tasks usin
 async def main():
     async with ClientSession("http://localhost:8000/api") as session:
         user_data_pipeline = (
-            task(list_user_ids, branch=True, bind=task.bind(session=session))
+            task(list_user_ids, branch=True)
             | task(fetch_user_data, workers=10, bind=task.bind(session=session))
         )
-        async for output in user_data_pipeline():
+        async for output in user_data_pipeline(session):
             print(output)
 ```
 
@@ -208,4 +208,90 @@ async def main():
             > copy_to_db
         )
         await run()
-```
+```
+
+## Generators
+
+### Usage
+
+Generators in Python are a mechanism for _lazy execution_, whereby results in an iterable are returned one by one (via underlying calls to `__next__`) instead of within a data structure, like a `list`, which requires all of its elements to be allocated in memory.
+
+Using generators is an indispensible approach for processing large volumes of data in a memory-friendly way. We can define generator functions by using the `yield` keyword within a normal `def` block:
+
+```python
+import typing
+from pyper import task
+
+# Okay
+@task(branch=True)
+def generate_values_lazily() -> typing.Iterable[dict]:
+    for i in range(10_000_000):
+        yield {"data": i}
+
+# Bad -- this creates 10 million values in memory
+# Subsequent tasks also cannot start executing until the entire list is created
+@task(branch=True)
+def create_values_in_list() -> typing.List[dict]:
+    return [{"data": i} for i in range(10_000_000)]
+```
+
+{: .info}
+Generator `functions` return immediately. They return `generator` objects, which are iterable
+
+Using the `branch` task parameter in Pyper allows generators to generate multiple outputs, which get picked up by subsequent tasks as soon as the data is available.
+
+Using a generator function without `branch=True` is also possible; this just means the task submits `generator` objects as output, instead of each generated value.
+
+```python
+from pyper import task
+
+def get_data():
+    yield 1
+    yield 2
+    yield 3
+
+if __name__ == "__main__":
+    branched_pipeline = task(get_data, branch=True)
+    for output in branched_pipeline():
+        print(output)
+        # Prints:
+        # 1
+        # 2
+        # 3
+
+    non_branched_pipeline = task(get_data)
+    for output in non_branched_pipeline():
+        print(output)
+        # Prints:
+        # <generator object get_data at ...>
+```
+
+### Limitations
+
+Implementing generator objects in a pipeline can also come with some caveats that are important to keep in mind.
+
+{: .text-green-200}
+**Synchronous Generators with Asynchronous Code**
+
+Synchronous generators in an `AsyncPipeline` do not benefit from threading or multiprocessing.
+
+This is because, in order to be scheduled in an async event loop, each synchronous task is run by a thread/process, and then wrapped in an `asyncio.Task`.
+
+Generator functions, which return _immediately_, do most of their work outside of the thread/process and this synchronous work will therefore not benefit from multiple workers in an async context.
+
+The alternatives are to:
+
+1. Use a synchronous generator anyway (if its performance is unlikely to be a bottleneck)
+
+2. Use a normal synchronous function, and return an iterable data structure (if memory is unlikely to be a bottleneck)
+
+3. Use an async generator (if an async implementation of the function is appropriate)
+
+{: .text-green-200}
+**Multiprocessing and Pickling**
+
+In Python, anything that goes into and comes out of a process must be picklable.
+
+On Windows, generator objects cannot be pickled, so cannot be passed as inputs and outputs when multiprocessing.
+
+Note that, for example, using `branch=True` to pass individual outputs from a generator into a multiprocessed task is still fine, because the task input would not be a `generator` object.

docs/src/docs/UserGuide/ComposingPipelines.md

@@ -78,7 +78,7 @@ if __name__ == "__main__":
     writer(pipeline(limit=10))  # Run
 ```
 
-The `>` operator (again inspired by UNIX syntax) is used to pipe a `Pipeline` into a consumer function (any callable that takes a data stream) returning simply a function that handles the 'run' operation. This is syntactic sugar for the `Pipeline.consume` method.
+The `>` operator (again inspired by UNIX syntax) is used to pipe a `Pipeline` into a consumer function (any callable that takes an `Iterable` of inputs) returning simply a function that handles the 'run' operation. This is syntactic sugar for the `Pipeline.consume` method.
 ```python
 if __name__ == "__main__":
     run = step1 | step2 > JsonFileWriter("data.json")
@@ -105,26 +105,26 @@ For example, let's say we have a theoretical pipeline which takes `(source: str)
 
 ```python
 download_files_from_source = (
-    task(list_files, branch=True)
-    | task(download_file, workers=20)
-    | task(decrypt_file, workers=5, multiprocess=True)
+    task(list_files, branch=True)  # Return a list of file info
+    | task(download_file, workers=20)  # Return a filepath
+    | task(decrypt_file, workers=5, multiprocess=True)  # Return a filepath
 )
 ```
 
 This is a function which generates multiple outputs per source. But we may wish to process _batches of filepaths_ downstream, after waiting for a single source to finish downloading. This means a piping approach, where we pass each _individual_ filepath along to subsequent tasks, won't work.
 
-Instead, we can define a function to create a list of filepaths as `download_files_from_source > list`. This is now a composable function which can be used in an outer pipeline.
+Instead, we can define `download_files_from_source` as a task within an outer pipeline, which is as simple as wrapping it in `task` like we would with any other function.
 
 ```python
 download_and_merge_files = (
-    task(get_sources, branch=True)
-    | task(download_files_from_source > list)
-    | task(merge_files, workers=5, multiprocess=True)
+    task(get_sources, branch=True)  # Return a list of sources
+    | task(download_files_from_source)  # Return a batch of filepaths (as a generator)
+    | task(sync_files, workers=5)  # Do something with each batch
 )
 ```
 
-* `download_files_from source > list` takes a source as input, downloads all files, and creates a list of filepaths as output.
-* `merge_files` takes a list of filepaths as input.
+* `download_files_from_source` takes a source as input, and returns a generator of filepaths (note that we are _not_ setting `branch=True`; a batch of filepaths is being passed along per source)
+* `sync_files` takes each batch of filepaths as input, and works on them concurrently
 
 ## Asynchronous Code
 

docs/src/index.md

@@ -52,7 +52,7 @@ It is designed with the following goals in mind:
 * **Error Handling**: Data flows fail fast, even in long-running threads, and propagate their errors cleanly
 * **Complex Data Flows**: Data pipelines support branching/joining data flows, as well as sharing contexts/resources between tasks
 
-In addition, Pyper provides an extensible way to write code that can be integrated with other frameworks like those aforementioned.
+In addition, Pyper enables developers to write code in an extensible way that can be integrated naturally with other frameworks like those aforementioned.
 
 ## Installation
 

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
 ]
 
 [project.urls]
+Homepage = "https://pyper-dev.github.io/pyper/"
 Documentation = "https://pyper-dev.github.io/pyper/"
 Repository = "https://github.com/pyper-dev/pyper"
 Issues = "https://github.com/pyper-dev/pyper/issues"

src/pyper/_core/async_helper/queue_io.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import AsyncIterable, Iterable
 from typing import TYPE_CHECKING
 
 from ..util.sentinel import StopSentinel
@@ -61,10 +61,11 @@ async def __call__(self, *args, **kwargs):
 
 class _BranchingAsyncEnqueue(_AsyncEnqueue):
     async def __call__(self, *args, **kwargs):
-        if self.task.is_gen:
-            async for output in self.task.func(*args, **kwargs):
+        result = self.task.func(*args, **kwargs)
+        if isinstance(result, AsyncIterable):
+            async for output in result:
                 await self.q_out.put(output)
-        elif isinstance(result := await self.task.func(*args, **kwargs), Iterable):
+        elif isinstance(result := await result, Iterable):
             for output in result:
                 await self.q_out.put(output)
         else:

src/pyper/_core/async_helper/stage.py

@@ -19,9 +19,9 @@
 class AsyncProducer:
     def __init__(self, task: Task, next_task: Task):
         if task.workers > 1:
-            raise RuntimeError(f"The first task in a pipeline ({task.func.__qualname__}) cannot have more than 1 worker")
+            raise RuntimeError(f"The first task in a pipeline ({task.func}) 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")
+            raise RuntimeError(f"The first task in a pipeline ({task.func}) cannot join previous results")
         self.task = task
         self.q_out = asyncio.Queue(maxsize=task.throttle)
 

src/pyper/_core/sync_helper/output.py

@@ -1,30 +1,31 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-import queue
+from typing import TYPE_CHECKING, Union
 
 from .stage import Producer, ProducerConsumer
 from ..util.sentinel import StopSentinel
 from ..util.worker_pool import ProcessPool, ThreadPool
 
 if TYPE_CHECKING:
+    import multiprocessing as mp
+    import queue
     from ..pipeline import Pipeline
 
 
 class PipelineOutput:
     def __init__(self, pipeline: Pipeline):
         self.pipeline = pipeline
 
-    def _get_q_out(self, tp: ThreadPool, pp: ProcessPool, *args, **kwargs) -> queue.Queue:
+    def _get_q_out(self, tp: ThreadPool, pp: ProcessPool, *args, **kwargs) -> Union[mp.Queue, 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]):
             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
 
@@ -34,18 +35,10 @@ def __call__(self, *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:
-                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
+            try:
+                while (data := q_out.get()) is not StopSentinel:
                     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
+            except (KeyboardInterrupt, SystemExit):  # pragma: no cover
+                tp.shutdown_event.set()
+                pp.shutdown_event.set()
+                raise

src/pyper/_core/sync_helper/queue_io.py

@@ -30,13 +30,13 @@ def __call__(self):
 
 
 class _SingleDequeue(_Dequeue):
-     def __call__(self):
-         for data in self._input_stream():
+    def __call__(self):
+        for data in self._input_stream():
             yield data
 
 
 class _JoiningDequeue(_Dequeue):
-     def __call__(self):
+    def __call__(self):
         yield self._input_stream()
 
 
@@ -56,12 +56,12 @@ def __call__(self, *args, **kwargs):
 
 
 class _SingleEnqueue(_Enqueue):        
-     def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs):
         self.q_out.put(self.task.func(*args, **kwargs))
 
 
 class _BranchingEnqueue(_Enqueue):
-     def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs):
         if isinstance(result := self.task.func(*args, **kwargs), Iterable):
             for output in result:
                 self.q_out.put(output)