[pipeline] support async ops in fused subprocess stages (#1582)

`fuse_subprocess_stages=True` fuses runs of adjacent pipe stages that
share the same isolating-pool (process/interpreter) executor instance
into one nested `Pipeline` that runs inside a worker process,
eliminating the per-stage IPC round-trip. Previously an async op could
never be part of a fused run: `PipeConfig` rejected any `executor` on an
async op, and fusion groups stages purely by executor-instance identity.
An async op between two pool stages therefore split the fusable run in
two, forcing the intermediate value to round-trip (and be pickled)
through the main process.

This lets an async op join a fused run by tagging it with the same
isolating-pool executor as its neighbours. The executor is never used to
run the coroutine (an async op always runs on the event loop) — it is
only the fusion-group key. When fused, the tag is stripped and the op
runs on the worker's own event loop, exactly as fused `path_variants`
async branches already do. When not fused (fusion off, or a lone async
op), the tag is ignored and the op runs on the main loop as before. The
execution engine is unchanged: each worker already rebuilds the
sub-config with the normal `build_pipeline`, which runs a full event
loop.

Builds on the parent diff, which moves the
`_is_isolating_pool`/`_is_interpreter_pool` helpers into
`_common/_convert`; this diff adds their new call sites (notably in
`defs`).

Details:
- `PipeConfig.__post_init__` now allows an isolating-pool executor on an
async op and rejects only non-isolating executors (e.g. a thread pool),
which have no effect on an async op.
- `convert_to_async` ignores the executor for async ops instead of
asserting it is `None`.
- Fusion detection: an async op that now carries an executor flows
through `_fusable_pool_executor` automatically. Inside a `path_variants`
stage, `_scan_variant_pool_executors` ignores an async op's pool tag
instead of treating it as an input-ordered pool-pipe — an async op's
executor is only a fusion-group tag (it runs on the loop, not the pool),
so its `output_order="input"` cannot be broken by pool parallelism and
must not block fusing the same-pool stage. Only a sync input-ordered
pool-pipe still blocks fusion. `_stage_concurrency` counts async ops as
zero worker threads (they run on the loop, not the thread pool).
- `run_pipeline_in_subprocess` strips any executor tag left on an
unfused async op before the op-agnostic executor-hoisting pass, so a tag
never spawns an idle worker pool the op will not use.

A fused async op must be picklable, like any fused stage. This is
documented on `PipelineBuilder.pipe` and in the parallelism guide.

Author: mthrok

docs/source/getting_started/parallelism.rst

@@ -345,9 +345,24 @@ A **generator op** (a function that ``yield``\ s) is supported as a fused
 process-pool stage: each input item fans out into the values the generator
 yields, exactly as in an unfused pipeline. As with any sync generator on a
 process-pool executor, the yielded items are materialized once the generator is
-exhausted rather than streamed out incrementally. (An *async* generator cannot
-take an ``executor`` and so is never itself a fused stage, but it composes with
-fusion when placed before or after a fused run, running in the main process.)
+exhausted rather than streamed out incrementally.
+
+An **async op** (an ``async def`` function or an async generator) can be fused
+too. Because an async op always runs on the event loop, it takes no executor to
+*run* it; instead, tag it with the **same** pool executor as its neighbours and
+it joins their fused run, executing on the worker's own event loop:
+
+.. code-block::
+
+   .pipe(sync_op,  executor=executor)
+   .pipe(async_op, executor=executor)   # runs on the worker's event loop
+   .pipe(sync_op,  executor=executor)
+
+All three fuse into one subprocess run, so an async op between two pool stages no
+longer splits the run in two. The executor is used only to group the stage, not
+to run the coroutine; a fused async op must be picklable, like any fused stage.
+Passing a non-isolating executor (e.g. a thread pool) to an async op is an error.
+Unfused, the tag is ignored and the async op runs in the main process.
 
 Only *adjacent* pool stages on the same executor are fused. An
 :py:meth:`~spdl.pipeline.PipelineBuilder.aggregate` or

src/spdl/pipeline/_build.py

@@ -35,7 +35,7 @@
     TaskHook,
 )
 from spdl.pipeline._executor_proxy import _make_config_executors_picklable
-from spdl.pipeline._fuse import _fuse_subprocess_stages
+from spdl.pipeline._fuse import _fuse_subprocess_stages, _strip_async_executor_tags
 from spdl.pipeline._iter_utils import iterate_in_subinterpreter, iterate_in_subprocess
 from spdl.pipeline._random_seed import _capture_rng_initializers
 from spdl.pipeline._subprocess_pipeline_pool import _shutdown_pipeline_pools
@@ -281,9 +281,11 @@ def build_pipeline(
             construct (router and branches) moves into the worker — and fuses on its own even
             when it is the only such stage. An ``aggregate``/``disaggregate`` between two pool
             stages is not fused (it keeps its main-process batching) and splits them into
-            separate runs. Continuous sources are supported (the fused worker sub-pipelines stay
-            warm across epochs and epoch boundaries are propagated across the pool). Default:
-            ``False``.
+            separate runs. An async op joins a fused run when tagged with the same executor as
+            its neighbours (see :py:meth:`~spdl.pipeline.PipelineBuilder.pipe`), running on the
+            worker's own event loop. Continuous sources are supported (the fused worker
+            sub-pipelines stay warm across epochs and epoch boundaries are propagated across the
+            pool). Default: ``False``.
 
             .. versionadded:: 0.6.0
                The ``fuse_subprocess_stages`` argument.
@@ -565,7 +567,9 @@ def run_pipeline_in_subprocess(
             This removes the per-stage round-trip between the pipeline subprocess and the pool
             workers (so intermediate values need not be picklable). A ``path_variants`` stage
             whose branches all use the same pool executor is fused too (router and branches move
-            into the worker). Continuous sources are supported. Default: ``False``.
+            into the worker). An async op joins a fused run when tagged with the same executor as
+            its neighbours (see :py:meth:`~spdl.pipeline.PipelineBuilder.pipe`), running on the
+            worker's own event loop. Continuous sources are supported. Default: ``False``.
 
             .. versionadded:: 0.6.0
                The ``fuse_subprocess_stages`` argument.
@@ -616,6 +620,11 @@ def run_pipeline_in_subprocess(
             stacklevel=3,
         )
 
+    # Clear executor tags left on any unfused async op: they are subprocess fusion-group hints,
+    # not real pools, and the executor-hoisting/pickling passes below are op-agnostic -- an async
+    # op's process-pool tag would otherwise spawn an idle worker pool it never submits to.
+    config = _strip_async_executor_tags(config)
+
     # Spawn workers for any stdlib ``ProcessPoolExecutor`` in the main process (as children of
     # main, not grandchildren via the pipeline subprocess), then replace the executor with a
     # queue-backed proxy that the subprocess submits to. The remaining stdlib executors

src/spdl/pipeline/_builder.py

@@ -170,7 +170,21 @@ def pipe(
             executor: A custom executor object to be used to convert the synchronous operation
                 into asynchronous one. If ``None``, the default executor is used.
 
-                It is invalid to provide this argument when the given op is already async.
+                When ``op`` is already async, ``executor`` must be an isolating-pool (process
+                or interpreter) executor and is **not** used to run the op -- an async op always
+                runs on the event loop. It serves only as a subprocess fusion-group tag: with
+                ``fuse_subprocess_stages=True`` (see
+                :py:meth:`~spdl.pipeline.PipelineBuilder.build`), adjacent stages sharing the
+                same executor instance are fused into one worker sub-pipeline, so tagging an
+                async op lets it join such a run (it then runs on the worker's own event loop
+                and, like any fused stage, must be picklable). Unfused, the tag is ignored and
+                the op runs in the main process. Passing a non-isolating executor (e.g. a thread
+                pool) with an async op is an error.
+
+                .. versionchanged:: 0.6.0
+                   An async ``op`` may now be given an isolating-pool ``executor`` as a
+                   subprocess fusion-group tag; previously any ``executor`` on an async op was
+                   rejected.
             name: The name (prefix) to give to the task.
             output_order: If ``"completion"`` (default), the items are put to output queue
                 in the order their process is completed.

src/spdl/pipeline/_common/_convert.py

@@ -202,8 +202,11 @@ def convert_to_async(
         op = op.__call__
 
     if inspect.iscoroutinefunction(op) or inspect.isasyncgenfunction(op):
-        # op is async function. No need to convert.
-        assert executor is None  # This has been checked in `PipelineBuilder.pipe()`
+        # op is async function. No need to convert. An async op always runs on the event loop,
+        # so any ``executor`` it carries is not used to run it here -- it is only a subprocess
+        # fusion-group tag. When the stage is fused the tag is stripped by ``_strip_executor``;
+        # when it is not fused the ``_strip_async_executor_tags`` pass clears it before executor
+        # hoisting, so no idle worker pool is spawned for a pool the async op never submits to.
         return op  # pyre-ignore: [7]
 
     if inspect.isgeneratorfunction(op):

src/spdl/pipeline/_fuse.py

@@ -32,6 +32,7 @@
 
 from __future__ import annotations
 
+import inspect
 import multiprocessing as mp
 import os
 import threading
@@ -61,6 +62,7 @@
     "_FusableRun",
     "_find_fusable_runs",
     "_fuse_subprocess_stages",
+    "_strip_async_executor_tags",
 ]
 
 # Buffer size for the fused sub-pipeline's sink. Small: the worker drains it straight onto the
@@ -113,8 +115,8 @@ def _scan_variant_pool_executors(
     Recurses into nested path-variants. Appends each completion-ordered pool-pipe's executor to
     ``acc``; async/thread/default pipes are ignored (they run on the worker's own loop/threads
     once the stage is fused). Returns ``False`` if any branch holds an *input-ordered*
-    (``output_order="input"``) pool-pipe, which is not fusable for the same reason a top-level
-    one is not — its global input order cannot be preserved across the pool workers.
+    (``output_order="input"``) *sync* pool-pipe, which is not fusable for the same reason a
+    top-level one is not — its global input order cannot be preserved across the pool workers.
     """
     for path in cfg.paths:
         for stage in path:
@@ -129,8 +131,12 @@ def _scan_variant_pool_executors(
                 isinstance(stage, PipeConfig)
                 and stage._args.executor is not None
                 and _is_isolating_pool(stage._args.executor)
+                and not _is_async_op(stage._args.op)
             ):
-                # A pool-pipe that _fusable_pool_executor rejected -> input-ordered: not fusable.
+                # A sync pool-pipe that _fusable_pool_executor rejected -> input-ordered:
+                # its global input order can't be preserved across pool workers, so it is not
+                # fusable. An async op's executor is only a fusion-group tag (it runs on the
+                # loop, not the pool), so its ordering is unaffected and it never blocks fusion.
                 return False
     return True
 
@@ -251,11 +257,20 @@ def _strip_executor(cfg: object) -> object:
     return cfg
 
 
+def _is_async_op(op: object) -> bool:
+    """Whether ``op`` runs on the event loop rather than the worker's thread pool."""
+    return inspect.iscoroutinefunction(op) or inspect.isasyncgenfunction(op)
+
+
 def _stage_concurrency(cfg: object) -> int:
     """Total worker-thread demand of a fused stage: a pipe's ``concurrency``, or the sum across
-    every branch pipe of a path-variants stage (recursively). Other stages contribute 0."""
+    every branch pipe of a path-variants stage (recursively). Other stages contribute 0.
+
+    An async pipe contributes 0: its ``concurrency`` bounds concurrent coroutines on the
+    worker's event loop, which do not consume worker threads.
+    """
     if isinstance(cfg, PipeConfig):
-        return cfg._args.concurrency
+        return 0 if _is_async_op(cfg._args.op) else cfg._args.concurrency
     if isinstance(cfg, PathVariantsConfig):
         return sum(_stage_concurrency(s) for path in cfg.paths for s in path)
     return 0
@@ -409,3 +424,43 @@ def _fuse_subprocess_stages(
             pool.shutdown()
         raise
     return replace(config, pipes=new_pipes), pools  # pyre-ignore[6]
+
+
+def _strip_async_executor_tag(cfg: object) -> object:
+    """Clear the executor tag on an async-op stage, recursing into path-variants branches.
+
+    An async op's executor is only a subprocess fusion-group tag; once fusion has run, any tag
+    left on an *unfused* async op is a no-op that must not reach the subprocess
+    executor-hoisting machinery (which is op-agnostic and would otherwise spawn an idle worker
+    pool for it). Sync stages are returned unchanged.
+    """
+    if isinstance(cfg, PipeConfig):
+        if cfg._args.executor is not None and _is_async_op(cfg._args.op):
+            return replace(cfg, _args=replace(cfg._args, executor=None))
+        return cfg
+    if isinstance(cfg, PathVariantsConfig):
+        new_paths = tuple(
+            tuple(_strip_async_executor_tag(stage) for stage in path)
+            for path in cfg.paths
+        )
+        return replace(cfg, paths=new_paths)
+    return cfg
+
+
+def _strip_async_executor_tags(config: PipelineConfig[Any]) -> PipelineConfig[Any]:
+    """Return ``config`` with executor tags cleared from every unfused async-op stage.
+
+    Walks the top-level pipes, path-variants branches, and merged sub-configs (mirroring the
+    executor rewrite in :py:mod:`spdl.pipeline._executor_proxy`). Fused async ops are already
+    inside a :py:class:`_SubprocessPipelineConfig` handle (their tags stripped when the run was
+    built), so this only touches tags left on stages that stayed in the enclosing pipeline.
+    """
+    src = config.src
+    if isinstance(src, MergeConfig):
+        new_configs = tuple(
+            _strip_async_executor_tags(plc) for plc in src.pipeline_configs
+        )
+        # pyre-ignore[6]: MergeConfig annotates a 1-tuple but holds N configs.
+        src = replace(src, pipeline_configs=new_configs)
+    new_pipes = [_strip_async_executor_tag(p) for p in config.pipes]
+    return replace(config, src=src, pipes=new_pipes)  # pyre-ignore[6]

src/spdl/pipeline/defs/_defs.py

@@ -21,6 +21,7 @@
 from functools import partial
 from typing import Any, Generic, Protocol, runtime_checkable, TypeAlias, TypeVar
 
+from spdl.pipeline._common._convert import _is_isolating_pool
 from spdl.pipeline._common._source_locator import locate_source
 from spdl.pipeline._common._types import _TCallables, _TMergeOp
 
@@ -212,8 +213,14 @@ class PipeConfig(Generic[T, U]):
     def __post_init__(self) -> None:
         op = self._args.op
         if inspect.iscoroutinefunction(op) or inspect.isasyncgenfunction(op):
-            if self._args.executor is not None:
-                raise ValueError("`executor` cannot be specified when op is async.")
+            executor = self._args.executor
+            if executor is not None and not _is_isolating_pool(executor):
+                raise ValueError(
+                    "An async op may only be given an isolating-pool (process or interpreter) "
+                    "executor, which is used solely to group the op into a subprocess fusion "
+                    "run -- not to execute it. A non-isolating executor (e.g. a thread pool) "
+                    "has no effect on an async op and is not allowed."
+                )
         if inspect.isasyncgenfunction(op):
             if self._type == _PipeType.OrderedPipe:
                 raise ValueError(
@@ -790,7 +797,11 @@ def Pipe(
         executor: A custom executor object to be used to convert the synchronous operation
             into asynchronous one. If ``None``, the default executor is used.
 
-            It is invalid to provide this argument when the given op is already async.
+            When ``op`` is already async, ``executor`` must be an isolating-pool (process or
+            interpreter) executor and is not used to run the op (an async op always runs on the
+            event loop) -- it serves only as a subprocess fusion-group tag (see
+            ``fuse_subprocess_stages`` on :py:meth:`~spdl.pipeline.PipelineBuilder.build`).
+            Passing a non-isolating executor (e.g. a thread pool) with an async op is an error.
         name: The name (prefix) to give to the task.
         output_order: If ``"completion"`` (default), the items are put to output queue
             in the order their process is completed.