[pipeline] 2/4: Fuse .to() subprocess regions in the engine

Part 2/4 of the `.to()` region API; see D110604018 for the overall design and rationale. Behind-the-scenes engine change — backward compatible and dormant until region markers exist, so it lands before the public `.to()` surface (4/4) without altering any current behavior.

Adds `_fuse_marked_regions`: it walks `PipelineConfig.pipes`, tracks the current execution target set by `ExecutorConfig` markers (a pipeline starts on the main process), and replaces each maximal span of stages under a `SubprocessConfig` target with one stage that runs the span as a nested pipeline in a worker pool — eliminating inter-stage IPC within the region. Crucially, `aggregate`/`disaggregate`/`path_variants` stages inside a region are absorbed into it, which the existing executor-identity fusion cannot do (it bounds a run at those stages). Pool parameters come from the serializable spec rather than a live executor.

`_build_fused_stage` is refactored to share a core with the new spec-driven builder. `build_pipeline` now runs `_fuse_marked_regions` unconditionally; it is a no-op when the config has no markers, so the existing `fuse_subprocess_stages` path is untouched. Subinterpreter regions raise `NotImplementedError` pending the subinterpreter worker-pool backend (3/4).

Author: mthrok

src/spdl/pipeline/_build.py

@@ -35,7 +35,11 @@
     TaskHook,
 )
 from spdl.pipeline._executor_proxy import _make_config_executors_picklable
-from spdl.pipeline._fuse import _fuse_subprocess_stages, _strip_async_executor_tags
+from spdl.pipeline._fuse import (
+    _fuse_marked_regions,
+    _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
@@ -147,14 +151,30 @@ def _build_pipeline(
             _LG.exception("Build callback failed.")
 
     pools: list[Any] = []
-    if fuse_subprocess_stages:
-        # Fuse consecutive same-pool stages so each run executes as one nested pipeline inside a
-        # worker pool, eliminating the inter-stage IPC. The pools are owned by the returned
-        # Pipeline and reaped when it stops.
-        # stacklevel=4: _fuse_subprocess_stages -> _build_pipeline -> build_pipeline -> user.
-        pipeline_cfg, pools = _fuse_subprocess_stages(
+    # Both fusion passes eagerly spawn worker pools. Reap them together on failure: each pass
+    # only reaps its own pools if it raises, so without this a failure in the second pass would
+    # leak the pools the first already spawned -- this half-built pipeline is never returned to
+    # the caller to be stopped.
+    try:
+        # Honor explicit `.to()` region markers first. This is a no-op when the config has no
+        # markers, so it is always safe to run and independent of `fuse_subprocess_stages`.
+        # stacklevel=4: _fuse_marked_regions -> _build_pipeline -> build_pipeline -> user.
+        pipeline_cfg, region_pools = _fuse_marked_regions(
             pipeline_cfg, report_stats_interval=report_stats_interval, stacklevel=4
         )
+        pools.extend(region_pools)
+        if fuse_subprocess_stages:
+            # Fuse consecutive same-pool stages so each run executes as one nested pipeline
+            # inside a worker pool, eliminating the inter-stage IPC. The pools are owned by the
+            # returned Pipeline and reaped when it stops.
+            # stacklevel=4: _fuse_subprocess_stages -> _build_pipeline -> build_pipeline -> user.
+            pipeline_cfg, id_pools = _fuse_subprocess_stages(
+                pipeline_cfg, report_stats_interval=report_stats_interval, stacklevel=4
+            )
+            pools.extend(id_pools)
+    except BaseException:
+        _shutdown_pipeline_pools(pools)
+        raise
 
     desc = repr(pipeline_cfg)
 

src/spdl/pipeline/_fuse.py

@@ -48,19 +48,24 @@
 from spdl.pipeline._executor_proxy import _ensure_executor_unused
 from spdl.pipeline._subprocess_pipeline_pool import _SubprocessPipelinePool
 from spdl.pipeline.defs._defs import (
+    _MainProcess,
     _PipeType,
     _SubprocessPipelineConfig,
+    ExecutorConfig,
     MergeConfig,
     PathVariantsConfig,
     PipeConfig,
     PipelineConfig,
     SinkConfig,
     SourceConfig,
+    SubinterpreterConfig,
+    SubprocessConfig,
 )
 
 __all__ = [
     "_FusableRun",
     "_find_fusable_runs",
+    "_fuse_marked_regions",
     "_fuse_subprocess_stages",
     "_strip_async_executor_tags",
 ]
@@ -316,24 +321,29 @@ def _warn_fork_with_threads(ctx: Any, stacklevel: int) -> None:
         )
 
 
-def _build_fused_stage(
+def _build_fused_stage_core(
     stages: Sequence[object],
-    executor: Executor,
+    *,
     ctx: Any,
+    num_threads: int,
+    max_workers: int,
+    user_initializer: Any,
+    user_initargs: tuple[Any, ...],
     report_stats_interval: float,
     continuous: bool,
 ) -> tuple[_SubprocessPipelineConfig, _SubprocessPipelinePool]:
-    """Build the worker pool and replacement stage for one fusable run."""
+    """Spawn a worker pool that runs ``stages`` as a nested pipeline, and return the pool and its
+    replacement stage. Shared by the executor-identity fusion (:py:func:`_build_fused_stage`) and
+    the ``.to()`` marker fusion (:py:func:`_build_fused_stage_from_spec`) — the two differ only in
+    where the pool parameters come from (a live executor vs. a serializable spec)."""
     stripped = [_strip_executor(s) for s in stages]
-    num_threads = max(1, sum(_stage_concurrency(s) for s in stages))
     sub_config: PipelineConfig[Any] = PipelineConfig(
         src=SourceConfig(
             []
         ),  # placeholder; the worker sets the real source per session
         pipes=stripped,  # pyre-ignore[6]
         sink=SinkConfig(_FUSED_SINK_BUFFER),
     )
-    max_workers, user_initializer, user_initargs = _pool_params(executor)
     build_kwargs = {
         "num_threads": num_threads,
         "report_stats_interval": report_stats_interval,
@@ -354,6 +364,62 @@ def _build_fused_stage(
     return _SubprocessPipelineConfig(name=name, handle=pool.make_handle()), pool
 
 
+def _build_fused_stage(
+    stages: Sequence[object],
+    executor: Executor,
+    ctx: Any,
+    report_stats_interval: float,
+    continuous: bool,
+) -> tuple[_SubprocessPipelineConfig, _SubprocessPipelinePool]:
+    """Build the worker pool and replacement stage for one executor-identity fusable run."""
+    max_workers, user_initializer, user_initargs = _pool_params(executor)
+    return _build_fused_stage_core(
+        stages,
+        ctx=ctx,
+        num_threads=max(1, sum(_stage_concurrency(s) for s in stages)),
+        max_workers=max_workers,
+        user_initializer=user_initializer,
+        user_initargs=user_initargs,
+        report_stats_interval=report_stats_interval,
+        continuous=continuous,
+    )
+
+
+def _build_fused_stage_from_spec(
+    stages: Sequence[object],
+    spec: SubprocessConfig,
+    ctx: Any,
+    report_stats_interval: float,
+    continuous: bool,
+) -> tuple[_SubprocessPipelineConfig, _SubprocessPipelinePool]:
+    """Build the worker pool and replacement stage for one ``.to()`` region, reading the pool
+    parameters from ``spec`` rather than a live executor. Unset spec fields fall back to the same
+    defaults the identity path uses: ``num_threads`` to the sum of the stages' concurrency,
+    ``max_workers`` to the CPU count, and ``report_stats_interval`` to the value passed to
+    :py:func:`~spdl.pipeline._build.build_pipeline`."""
+    num_threads = (
+        max(1, spec.num_threads)
+        if spec.num_threads is not None
+        else max(1, sum(_stage_concurrency(s) for s in stages))
+    )
+    max_workers = spec.max_workers or os.cpu_count() or 1
+    rsi = (
+        spec.report_stats_interval
+        if spec.report_stats_interval is not None
+        else report_stats_interval
+    )
+    return _build_fused_stage_core(
+        stages,
+        ctx=ctx,
+        num_threads=num_threads,
+        max_workers=max_workers,
+        user_initializer=spec.initializer,
+        user_initargs=spec.initargs,
+        report_stats_interval=rsi,
+        continuous=continuous,
+    )
+
+
 def _fuse_subprocess_stages(
     config: PipelineConfig[Any],
     *,
@@ -426,6 +492,92 @@ def _fuse_subprocess_stages(
     return replace(config, pipes=new_pipes), pools  # pyre-ignore[6]
 
 
+def _has_executor_markers(pipes: Sequence[object]) -> bool:
+    """Whether ``pipes`` contains any ``.to()`` region marker."""
+    return any(isinstance(p, ExecutorConfig) for p in pipes)
+
+
+def _fuse_marked_regions(
+    config: PipelineConfig[Any],
+    *,
+    report_stats_interval: float = -1,
+    stacklevel: int = 3,
+) -> tuple[PipelineConfig[Any], list[_SubprocessPipelinePool]]:
+    """Fuse each ``.to()`` region into one subprocess-pipeline stage.
+
+    Walks ``config.pipes`` tracking the current execution target set by
+    :py:class:`~spdl.pipeline.defs.ExecutorConfig` markers (a pipeline starts on the main
+    process). Every maximal span of stages under a subprocess target — pipes *and* the
+    aggregate/disaggregate/path-variants stages between them, which the executor-identity fusion
+    leaves in the main process — is replaced by a single stage that runs the span as one nested
+    pipeline inside a worker pool. Main-process spans are kept unchanged, and the marker nodes
+    themselves are dropped. The spawned pools are returned for the caller to reap at teardown; the
+    input ``config`` is not mutated.
+
+    This is a no-op (returns ``config`` and no pools) when there are no markers, so it is safe to
+    run unconditionally and leaves the executor-identity :py:func:`_fuse_subprocess_stages` path
+    unchanged.
+
+    Args:
+        config: The pipeline configuration to rewrite.
+        report_stats_interval: Fallback stats interval for a region whose spec does not set one.
+        stacklevel: ``warnings.warn`` stack level measured at this function. The fork-with-threads
+            warning is raised from the nested ``_flush`` closure (two frames deeper), so it uses
+            ``stacklevel + 2``.
+
+    Returns:
+        A tuple ``(new_config, pools)``. ``pools`` is empty when no region is fused.
+    """
+    pipes = list(config.pipes)
+    if not _has_executor_markers(pipes):
+        return config, []
+
+    continuous = _has_continuous_source(config)
+    pools: list[_SubprocessPipelinePool] = []
+    new_pipes: list[object] = []
+    target: SubprocessConfig | SubinterpreterConfig | _MainProcess = _MainProcess()
+    region: list[object] = []
+
+    def _flush() -> None:
+        if not region:
+            return
+        if isinstance(target, SubinterpreterConfig):
+            raise NotImplementedError(
+                "Subinterpreter regions (`.to(SubinterpreterConfig(...))`) are not yet "
+                "supported; use `SubprocessConfig` for now."
+            )
+        assert isinstance(
+            target, SubprocessConfig
+        )  # narrowed: not main, not subinterpreter
+        ctx = mp.get_context(target.mp_context)
+        # +2, not +1: this runs inside the nested ``_flush`` closure, one frame deeper than
+        # ``_fuse_marked_regions`` itself, so the warning still points at the user's call site.
+        _warn_fork_with_threads(ctx, stacklevel + 2)
+        fused, pool = _build_fused_stage_from_spec(
+            region, target, ctx, report_stats_interval, continuous
+        )
+        pools.append(pool)
+        new_pipes.append(fused)
+        region.clear()
+
+    try:
+        for p in pipes:
+            if isinstance(p, ExecutorConfig):
+                _flush()  # close the span running under the previous target
+                target = p.target
+            elif isinstance(target, _MainProcess):
+                new_pipes.append(p)
+            else:
+                region.append(p)
+        _flush()  # close a region left open at the end of the pipes
+    except BaseException:
+        # Reap any pools spawned before the failure; the caller never receives them to reap.
+        for pool in pools:
+            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.