Skip to content

[pipeline] Remove fuse_subprocess_stages #1588

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mthrok wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[pipeline] Remove fuse_subprocess_stages #1588

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 54 additions & 63 deletions docs/source/getting_started/parallelism.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,85 +309,76 @@ It is also cumbersome: you must hand-combine the stages into one function,
which collapses the per-stage performance stats into a single number and
discards each stage's individual ``concurrency``.

Multi-processing (fused)
------------------------
Multi-processing (region)
-------------------------

Instead of hand-combining stages, you can ask the builder to fuse them for
you by passing ``fuse_subprocess_stages=True`` to
:py:meth:`~spdl.pipeline.PipelineBuilder.build`. Runs of consecutive
:py:meth:`~spdl.pipeline.PipelineBuilder.pipe` stages that share the **same**
process-pool (or interpreter-pool) executor instance are fused into a single
stage that runs the run as one nested :py:class:`Pipeline` inside the worker
pool:
Instead of hand-combining stages, you can mark a **region** of the pipeline to
run together in a worker pool with :py:meth:`~spdl.pipeline.PipelineBuilder.to`.
Every stage between ``.to(ProcessPoolExecutorConfig(...))`` and ``.to(MAIN_PROCESS)`` runs
as one nested :py:class:`Pipeline` inside a pool of worker processes:

.. code-block::

executor = ProcessPoolExecutor(max_workers=4)
from spdl.pipeline.defs import MAIN_PROCESS, ProcessPoolExecutorConfig

pipeline = (
PipelineBuilder()
.add_source(...)
.pipe(op1, executor=executor, concurrency=2)
.pipe(op2, executor=executor, concurrency=3)
.to(ProcessPoolExecutorConfig(max_workers=4))
.pipe(op1, concurrency=2) # runs in a worker process
.aggregate(batch_size) # runs in a worker process
.pipe(op2, concurrency=3) # runs in a worker process
.to(MAIN_PROCESS) # data returns to the main process
.add_sink(...)
.build(num_threads=..., fuse_subprocess_stages=True)
.build(num_threads=...)
)

Because ``op1`` and ``op2`` now run back-to-back inside one worker, the
intermediate value is **not** copied back to the main process between them.
This removes the inter-stage IPC entirely, and — unlike the per-stage
multi-processing above — the value handed from ``op1`` to ``op2`` does **not**
need to be picklable. Each fused stage keeps its own ``concurrency`` and its
own per-stage performance stats (the nested pipeline is built with the usual
hooks, so the stats are reported from inside the worker).

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 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
:py:meth:`~spdl.pipeline.PipelineBuilder.disaggregate` between two pool stages
is **not** fused — it runs in the main process and keeps its usual batching
semantics, and it splits the surrounding pool stages into separate runs (so
each side fuses on its own only if it has two or more adjacent pool stages).

The same option is accepted by
:py:func:`spdl.pipeline.run_pipeline_in_subprocess`, where the fused worker
pool is owned by the main process and the run executes in those workers — so
the per-stage round-trip between the pipeline subprocess and the pool is
removed as well.

Fusion also works with a **continuous source**
Because the region's stages run back-to-back inside one worker, the value handed
from one stage to the next is **not** copied back to the main process between
them. This removes the inter-stage IPC entirely, and — unlike the per-stage
multi-processing above — those intermediate values do **not** need to be
picklable; only the region's inputs and outputs cross the process boundary. Each
stage keeps its own ``concurrency`` and its own per-stage performance stats (the
nested pipeline is built with the usual hooks, so the stats are reported from
inside the worker).

Unlike passing ``executor=`` to individual
:py:meth:`~spdl.pipeline.PipelineBuilder.pipe` calls, a region also carries
:py:meth:`~spdl.pipeline.PipelineBuilder.aggregate`,
:py:meth:`~spdl.pipeline.PipelineBuilder.disaggregate`, and
:py:meth:`~spdl.pipeline.PipelineBuilder.path_variants` stages into the worker,
and gives the worker-pool configuration (worker count, ``mp_context``,
``initializer``) a single home. A pipeline starts on the main process, so a
region is opened by a ``.to(ProcessPoolExecutorConfig(...))`` and closed by
``.to(MAIN_PROCESS)``; the region must be closed before
:py:meth:`~spdl.pipeline.PipelineBuilder.add_sink`.

A **generator op** (a function that ``yield``\ s) works inside a region: each
input item fans out into the values it yields, exactly as in an unfused pipeline.
An **async op** works too — it runs on the worker's own event loop. Every stage
inside a region must be picklable, since the region config is shipped to the
worker.

Regions also compose with :py:func:`spdl.pipeline.run_pipeline_in_subprocess`:
the region's worker pool is owned by the main process and the run executes in
those workers, so the per-stage round-trip between the pipeline subprocess and
the pool is removed as well.

Regions also work with a **continuous source**
(:py:meth:`~spdl.pipeline.PipelineBuilder.add_source(..., continuous=True)
<spdl.pipeline.PipelineBuilder.add_source>`). The worker sub-pipelines run in
continuous mode and stay warm across epochs, and epoch boundaries are
propagated across the pool: each fused stage emits one epoch boundary per epoch
just like an unfused pipeline.
continuous mode and stay warm across epochs, and epoch boundaries are propagated
across the pool: each region emits one epoch boundary per epoch just like an
unfused pipeline.

To run a region in **subinterpreters** (Python 3.14+) instead of subprocesses,
pass a :py:class:`~spdl.pipeline.defs.InterpreterPoolExecutorConfig`; the region's ops
must avoid NumPy/PyTorch, which cannot be imported in a subinterpreter.

.. note::

Fusion preserves results but produces them in completion order across the
pool workers. Stages built with ``output_order="input"`` are not fused.
A region produces results in completion order across its pool workers, so a
stage built with ``output_order="input"`` cannot appear inside a region.

Multi-threading in subprocess
-----------------------------
Expand Down
120 changes: 29 additions & 91 deletions src/spdl/pipeline/_build.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,11 +35,7 @@
TaskHook,
)
from spdl.pipeline._executor_proxy import _make_config_executors_picklable
from spdl.pipeline._fuse import (
_fuse_marked_regions,
_fuse_subprocess_stages,
_strip_async_executor_tags,
)
from spdl.pipeline._fuse import _fuse_marked_regions
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
Expand Down Expand Up @@ -142,39 +138,20 @@ def _build_pipeline(
stage_id: int = 0,
background_tasks: list[BackgroundTaskFactory] | None = None,
use_thread_output_queue: bool = False,
fuse_subprocess_stages: bool = False,
) -> Pipeline[U]:
if _DEFAULT_BUILD_CALLBACK is not None:
try:
_DEFAULT_BUILD_CALLBACK(pipeline_cfg)
except Exception:
_LG.exception("Build callback failed.")

pools: list[Any] = []
# 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
# Fuse each `.to()` region into one nested-pipeline stage that runs in a worker pool,
# eliminating the inter-stage IPC within the region. A no-op when the config has no markers.
# The pools are owned by the returned Pipeline and reaped when it stops.
# stacklevel=4: _fuse_marked_regions -> _build_pipeline -> build_pipeline -> user.
pipeline_cfg, pools = _fuse_marked_regions(
pipeline_cfg, report_stats_interval=report_stats_interval, stacklevel=4
)

desc = repr(pipeline_cfg)

Expand Down Expand Up @@ -218,7 +195,6 @@ def build_pipeline(
stage_id: int = 0,
background_tasks: list[BackgroundTaskFactory] | None = None,
use_thread_output_queue: bool = False,
fuse_subprocess_stages: bool = False,
) -> Pipeline[U]:
"""Build a pipeline from the config.

Expand Down Expand Up @@ -291,24 +267,11 @@ def build_pipeline(
``asyncio.run_coroutine_threadsafe``, reducing per-batch latency from
~200-400us to ~10us. Default: ``False``.

fuse_subprocess_stages: If ``True``, fuse runs of two or more adjacent pipe stages that
share the same process-pool (or interpreter-pool) executor instance into a single
stage that executes the run as one nested pipeline inside a worker pool. This
eliminates the inter-stage IPC that otherwise round-trips data back to this process
between each stage (so intermediate values need not be picklable), while each fused
stage keeps its own ``concurrency`` and per-stage stats. A ``path_variants`` stage
whose branches all use the same pool executor is fused too — the whole routing
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. 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.
.. seealso::

:py:meth:`~spdl.pipeline.PipelineBuilder.to`
Designate a region of stages to run together in a subprocess (or subinterpreter)
worker pool, eliminating the inter-stage IPC within the region.
"""
from . import _profile

Expand All @@ -325,7 +288,6 @@ def build_pipeline(
stage_id=stage_id,
background_tasks=background_tasks,
use_thread_output_queue=use_thread_output_queue,
fuse_subprocess_stages=fuse_subprocess_stages,
)


Expand Down Expand Up @@ -427,7 +389,6 @@ def run_pipeline_in_subprocess(
task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None,
background_tasks: list[BackgroundTaskFactory] | None = None,
use_thread_output_queue: bool = False,
fuse_subprocess_stages: bool = False,
**kwargs: Any,
) -> Iterable[T]:
"""Run the given Pipeline in a subprocess, and iterate on the result.
Expand Down Expand Up @@ -579,22 +540,16 @@ def run_pipeline_in_subprocess(

num_threads,max_failures,report_stats_interval,queue_class,task_hook_factory,background_tasks:
Passed to :py:func:`build_pipeline`.
fuse_subprocess_stages: If ``True``, fuse runs of two or more adjacent pipe stages that
share the same process-pool (or interpreter-pool) executor instance into a single
stage that runs the run as one nested pipeline inside a worker pool. The worker
processes are spawned in (and owned by) the main process, exactly like a hoisted
``ProcessPoolExecutor``; the pipeline subprocess drives them through a queue handle.
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). 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.
kwargs: Passed to :py:func:`iterate_in_subprocess`.

.. seealso::

:py:meth:`~spdl.pipeline.PipelineBuilder.to`
Designate a region of stages to run together in a subprocess (or subinterpreter)
worker pool. When the config has such a region, its worker pool is spawned in (and
owned by) the main process, so it is not orphaned if the pipeline subprocess is
force-killed.

Yields:
The results yielded from the pipeline.

Expand Down Expand Up @@ -625,11 +580,11 @@ def run_pipeline_in_subprocess(
else config_or_builder.get_config() # pyre-ignore[16]
)

# Every pass below eagerly spawns worker pools, so they all run inside one try/except:
# ``_fuse_marked_regions`` spawns ``fuse_pools`` up front, and if any *later* pass
# (``_fuse_subprocess_stages``, ``_hoist_process_pools``) or the iterable creation
# raises, both ``fuse_pools`` and the hoisted ``pools`` must be reaped -- this half-built
# iterable is never returned to the caller to be stopped. Mirrors ``_build_pipeline``.
# Both passes below eagerly spawn worker pools, so they run inside one try/except:
# ``_fuse_marked_regions`` spawns ``fuse_pools`` up front, so if a later pass
# (``_hoist_process_pools``) or the iterable creation raises, both ``fuse_pools`` and the
# hoisted ``pools`` must be reaped -- this half-built iterable is never returned to the
# caller to be stopped. Mirrors the guard in ``_build_pipeline``.
fuse_pools: list[Any] = []
pools: list[Any] = []
try:
Expand All @@ -648,23 +603,6 @@ def run_pipeline_in_subprocess(
stacklevel=3,
)
fuse_pools.extend(region_pools)
# Also fuse runs of same-pool stages tagged with an identical executor into one nested
# pipeline inside a worker pool, eliminating the inter-stage IPC (before hoisting, so
# only unfused ProcessPoolExecutor stages remain).
if fuse_subprocess_stages:
# stacklevel=3: _fuse_subprocess_stages -> run_pipeline_in_subprocess -> user.
config, id_pools = _fuse_subprocess_stages(
config,
mp_context=kwargs.get("mp_context"),
report_stats_interval=report_stats_interval,
stacklevel=3,
)
fuse_pools.extend(id_pools)

# 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 pool it never uses.
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
Expand Down Expand Up @@ -693,9 +631,9 @@ def run_pipeline_in_subprocess(
**kwargs,
)
except BaseException:
# Any eager-spawn pass above (region fusion, identity fusion, hoisting) or the iterable
# creation failed; the iterable is never returned to the caller, so reap the region and
# hoisted pools here to avoid leaking their worker processes and pipe fds.
# Any eager-spawn pass above (region fusion, hoisting) or the iterable creation failed;
# the iterable is never returned to the caller, so reap the region and hoisted pools
# here to avoid leaking their worker processes and pipe fds.
_shutdown_pools(pools)
_shutdown_pipeline_pools(fuse_pools)
raise
Expand Down
Loading
Loading