[Pipeline] Support async ops in fused subprocess stages (#1579)

`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 commit 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.

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`.
- `_is_isolating_pool`/`_is_interpreter_pool` moved from `_fuse` to
`_common/_convert` so `defs` can reuse them without an import cycle.
- 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/autoresearch/pipeline_optimization/prompts/apply_changes.md

@@ -76,7 +76,7 @@ __PIPELINE_CODE__
 13. **TorchTNT scripts**: If the code uses TorchTNT (`torchtnt.framework.fit`, `train`, `AutoUnit`), the SPDL `Pipeline` is passed directly to TorchTNT as the `train_dataloader` (Pipeline is iterable — no wrapper class is required unless instrumentation was added). When applying changes:
    - **Pipeline construction changes** (concurrency, MTP, batch size): Modify the function that builds the `PipelineBuilder`, same as non-TorchTNT code. The `Pipeline` abstracts MTP vs pure multithreading, so the code passing Pipeline to TorchTNT does not change.
    - **Do NOT modify TorchTNT internals** (`fit()`, `train()`, `AutoUnit.train_step`). Only modify the pipeline construction.
-   - The pipeline is built once and iterated many times. `auto_stop()` is obsolete — do not call it, and do not rebuild per epoch.
+   - The pipeline is built once and iterated many times. There is no need for `auto_stop()` or rebuilding per epoch.
 
 Output the modified file:
 

src/spdl/autoresearch/pipeline_optimization/prompts/knowledge/knowledge.md

@@ -268,7 +268,7 @@ This wrapper is passed to TorchTNT's `fit()` or `train()` as the `train_dataload
 **Key facts about Pipeline:**
 - `Pipeline` is both iterable and iterator. It can technically be directly iterated, but the wrapper pattern above is recommended because it provides `get_iterator(timeout=...)` for timeout handling and `__len__` for epoch length.
 - **Preferred iteration**: Always use `pipeline.get_iterator(timeout=<seconds>)` to obtain an iterator with a timeout, so that jobs do not get stuck. Directly iterating with `for batch in pipeline:` or `iter(pipeline)` is discouraged because it lacks timeout handling.
-- The pipeline is built once and iterated many times. `auto_stop()` is obsolete — do not call it, and do not rebuild per epoch.
+- The pipeline is built once and iterated many times. There is no need for `auto_stop()` or rebuilding per epoch.
 - The `Pipeline` abstracts away whether it uses MTP (subprocess) or pure multithreading internally, so switching between them does not affect how the Pipeline is consumed.
 - TorchTNT calls `iter(dataloader)` at the start of each epoch, then `next(data_iter)` per step until `StopIteration`.
 - The pipeline's source controls how many items are produced per epoch. When the source is exhausted, the iterator raises `StopIteration`, ending the epoch in TorchTNT.

src/spdl/autoresearch/pipeline_optimization/prompts/knowledge/optimization_strategies.md

@@ -33,37 +33,6 @@ pipeline = (
 | `buffer_size` | `.add_sink(N)` | 2-10; affects memory, NOT throughput |
 | `output_order` | `.pipe(fn, output_order=...)` | "completion" for I/O (avoids head-of-line blocking), "input" for determinism |
 
-### Async Stage Functions: Pass Them As-Is
-
-`.pipe()` accepts both regular (`def`) and asynchronous (`async def`) callables. **Pass an async function directly — do not wrap it in `asyncio.run()`:**
-
-```python
-async def fetch(key):
-    async with session.get(url(key)) as resp:
-        return await resp.read()
-
-pipeline = (
-    PipelineBuilder()
-    .add_source(source)
-    .pipe(fetch, concurrency=32)   # CORRECT: coroutine function passed as-is
-    ...
-)
-```
-
-SPDL runs stages on its own event loop, so it awaits async stage functions natively and drives up to `concurrency` of them at once as cheap cooperative coroutines — ideal for I/O-bound work.
-
-**Anti-pattern:** wrapping the coroutine so the stage becomes synchronous.
-
-```python
-# WRONG — asyncio.run() spins up and tears down a brand-new event loop on
-# every single item. That per-item loop init/finalize is pure overhead, and
-# it forces the work onto pipeline worker threads instead of running as cheap
-# cooperative coroutines on SPDL's shared event loop.
-.pipe(lambda key: asyncio.run(fetch(key)), concurrency=32)
-```
-
-If a stage is already synchronous, just pass it as a normal function — no event loop is involved either way.
-
 ## Recommended Architecture: Multi-Threading in Subprocess (MTP)
 
 The production pattern runs the CPU-heavy pipeline in a **subprocess** and only does GPU transfer in the main process.

src/spdl/autoresearch/pipeline_optimization/prompts/plan_next.md

@@ -77,7 +77,7 @@ Before you can consider stopping, ALL of the following must have been attempted
    - **Critical — understand the existing code first**: Inspect the pipeline builder function's return type and structure. If it calls `.build()` or `.get_iterator()` (returning a `Pipeline` or iterator), the MTP refactor must change it to return a `PipelineConfig` (via `.get_config()`) or an unbuilt `PipelineBuilder` instead. `run_pipeline_in_subprocess()` accepts ONLY `PipelineConfig` — NOT `Pipeline` objects or iterators. Check type annotations and return statements.
    - **Separate CPU and GPU stages**: The backend pipeline (subprocess) must contain ONLY CPU-bound stages (fetch, decode, aggregate, collate). GPU stages like `transfer_tensor` require a CUDA context and MUST be in the frontend pipeline (main process). If the existing code includes `transfer_tensor` or similar GPU ops in the pipeline, exclude them from the backend config.
    - **Write the `description` field to explicitly instruct the code modifier** about these structural requirements, e.g.: "Refactor `build_pipeline()` to NOT call `.build()`. Instead, build a `PipelineBuilder` with only CPU stages (no `transfer_tensor`), call `.get_config()`, pass the config to `run_pipeline_in_subprocess(config, num_threads=N)`, then create a frontend `PipelineBuilder` that takes the subprocess iterable as source and applies `transfer_tensor`."
-   - **TorchTNT scripts**: If the code uses TorchTNT (`fit()`, `train()`, `AutoUnit`), the MTP refactor targets only the pipeline builder function — NOT TorchTNT internals. The `Pipeline` object abstracts away MTP vs pure multithreading, so the wrapper class that calls `get_iterator(timeout=...)` does not need to change. Do NOT use `auto_stop()` (obsolete) — the pipeline is built once and iterated many times.
+   - **TorchTNT scripts**: If the code uses TorchTNT (`fit()`, `train()`, `AutoUnit`), the MTP refactor targets only the pipeline builder function — NOT TorchTNT internals. The `Pipeline` object abstracts away MTP vs pure multithreading, so the wrapper class that calls `get_iterator(timeout=...)` does not need to change. Do NOT use `auto_stop()` — the pipeline is built once and iterated many times.
    - Use the **two-tier approach** for pickling: first try module-level functions with `functools.partial` (Tier 1). If the subprocess crashes silently (0 batches), retry with picklable callable classes that pickle objects directly from the main process (Tier 2). See "Pickling Constraints" in the knowledge base.
    - **HuggingFace tokenizers are NOT thread-safe** — use thread-local storage (TLS) when tokenizing with concurrency > 1.
    - Wrap the sampler with **`spdl.source.utils.embed_shuffle()`** for correct sampling behavior.
@@ -204,6 +204,6 @@ Rules:
 - torchx entrypoint args use **underscores**, not dashes (e.g. `--num_threads`)
 - The orchestrator will call a separate Claude session to apply the code changes described in `description`, commit them, rebuild the image, and then launch. **Write the `description` field as precise instructions for what code to modify** — specify function names, SPDL API calls to add/change, and the exact transformation. Do not write vague descriptions like "enable MTP" — instead write "Refactor `build_pipeline()` to return a `PipelineBuilder` config (without `.build()`), wrap it with `spdl.pipeline.run_pipeline_in_subprocess(config, num_threads=16, mp_context='forkserver')`, and create an outer pipeline that takes the subprocess source and applies GPU transfer via `pipe(transfer_tensor, executor=ThreadPoolExecutor(1))`."
 - Each experiment should differ from the baseline in exactly one dimension (or a small, justified set of changes)
-- **TorchTNT scripts**: When writing `description` for rebuild experiments in TorchTNT code, specify changes to the pipeline builder function only. The `Pipeline` abstracts away MTP vs pure multithreading, so the wrapper class (which calls `get_iterator(timeout=...)`) and TorchTNT internals (`fit()`, `train()`, `AutoUnit`) do not change. Do NOT use `auto_stop()` (obsolete) — Pipeline supports multiple iterations natively.
+- **TorchTNT scripts**: When writing `description` for rebuild experiments in TorchTNT code, specify changes to the pipeline builder function only. The `Pipeline` abstracts away MTP vs pure multithreading, so the wrapper class (which calls `get_iterator(timeout=...)`) and TorchTNT internals (`fit()`, `train()`, `AutoUnit`) do not change. Do NOT use `auto_stop()` — Pipeline supports multiple iterations without it.
 - **`best_practices_tags`**: Tag each experiment with which best practices it covers from the valid tags list. This is how the orchestrator tracks progress. If an experiment covers multiple practices, include all relevant tags.
 - **`goto`** (per-experiment): The orchestrator checks out the instrumentation (anchor) commit before applying each experiment's code changes by default. This ensures every experiment starts from a clean slate. Set `goto` to `null` in most cases. Only set it to a specific commit hash if you want to stack changes on top of a previous successful experiment (e.g., adding batch size tuning on top of an MTP experiment that already improved metrics). Never stack incompatible changes (e.g., GPU decode on top of MTP — they are mutually exclusive). **The `goto` field also determines the experiment's parent in the hypothesis tree**: `null` means the experiment branches from baseline; a specific commit means it branches from the experiment that produced that commit. This allows a single planning round to propose experiments with different parents (e.g., NVDEC from baseline + batch_size from a successful MTP experiment).

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):