From 688788b83047d5b4f744a6ed1ea3e6df24d9f345 Mon Sep 17 00:00:00 2001 From: Moto Hira Date: Fri, 10 Jul 2026 13:50:30 -0400 Subject: [PATCH] [pipeline] follow-up: Remove fuse_subprocess_stages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the `.to()` region API (1/4–4/4; see https://github.com/facebookresearch/spdl/pull/1584 for the overall design and rationale). Removes the executor-identity fusion path now that `.to()` regions provide the same capability with an explicit, statically-configurable surface. Both removed features were added in the still-unreleased 0.6.0 cycle, so no deprecation shim is warranted. Removed: - The `fuse_subprocess_stages` keyword from `PipelineBuilder.build`, `build_pipeline`, and `run_pipeline_in_subprocess`. - The executor-identity fusion machinery in `_fuse.py` (`_find_fusable_runs`, `_scan_run`, `_FusableRun`, `_fusable_*`, `_fuse_subprocess_stages`, `_pool_params`, the identity `_build_fused_stage`) and the async-op-as-fusion-tag pass (`_strip_async_executor_tags`). The marker path (`_fuse_marked_regions` and friends) is kept. Reverted (the async-op executor relaxation from https://github.com/facebookresearch/spdl/pull/1582): an async op may no longer be given an `executor`. `PipeConfig.__post_init__` again rejects any executor on an async op, and `convert_to_async` asserts it is `None`. In a `.to()` region an async op is placed by the marker and needs no per-stage executor tag, so the relaxation is obsolete. `run_pipeline_in_subprocess` now fuses `.to()` regions in the main process (via `_fuse_marked_regions`), preserving the main-ownership of region worker pools that the old flag provided. Docs: the parallelism guide's "Multi-processing (fused)" section is rewritten as "Multi-processing (region)" using `.to()`. --- docs/source/getting_started/parallelism.rst | 117 +- src/spdl/pipeline/_build.py | 120 +- src/spdl/pipeline/_builder.py | 35 +- src/spdl/pipeline/_common/_convert.py | 7 +- src/spdl/pipeline/_fuse.py | 372 +---- src/spdl/pipeline/defs/_defs.py | 28 +- tests/pipeline/builder_to_test.py | 36 +- tests/pipeline/fuse_test.py | 162 -- tests/pipeline/marked_region_fuse_test.py | 474 +++++- tests/pipeline/region_hooks_test.py | 235 +++ tests/pipeline/subprocess_pipe_bridge_test.py | 133 ++ .../pipeline/subprocess_pipeline_fuse_test.py | 1313 ----------------- tools/skills/authoring/building_pipelines.md | 20 +- 13 files changed, 1011 insertions(+), 2041 deletions(-) delete mode 100644 tests/pipeline/fuse_test.py create mode 100644 tests/pipeline/region_hooks_test.py create mode 100644 tests/pipeline/subprocess_pipe_bridge_test.py delete mode 100644 tests/pipeline/subprocess_pipeline_fuse_test.py diff --git a/docs/source/getting_started/parallelism.rst b/docs/source/getting_started/parallelism.rst index 2c5e0aa1c..d6b51f688 100644 --- a/docs/source/getting_started/parallelism.rst +++ b/docs/source/getting_started/parallelism.rst @@ -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) `). 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 ----------------------------- diff --git a/src/spdl/pipeline/_build.py b/src/spdl/pipeline/_build.py index 4d71e3530..186857d32 100644 --- a/src/spdl/pipeline/_build.py +++ b/src/spdl/pipeline/_build.py @@ -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 @@ -142,7 +138,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]: if _DEFAULT_BUILD_CALLBACK is not None: try: @@ -150,31 +145,13 @@ def _build_pipeline( 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) @@ -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. @@ -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 @@ -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, ) @@ -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. @@ -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. @@ -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: @@ -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 @@ -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 diff --git a/src/spdl/pipeline/_builder.py b/src/spdl/pipeline/_builder.py index a9c2b425c..c12ee8dcf 100644 --- a/src/spdl/pipeline/_builder.py +++ b/src/spdl/pipeline/_builder.py @@ -253,21 +253,7 @@ 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. - 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. + It is invalid to provide this argument when the given op is already async. 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. @@ -455,7 +441,6 @@ def build( task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, use_thread_output_queue: bool = False, - fuse_subprocess_stages: bool = False, ) -> Pipeline[U]: """Build the pipeline. @@ -498,19 +483,10 @@ def build( ``queue.Queue``-backed queue for lower-latency batch handoff. 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 runs 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 (router and branches move into the worker), and fuses on its own. An - ``aggregate``/``disaggregate`` between pool stages is not fused and runs in the - main process with its usual batching. Default: ``False``. - - .. versionadded:: 0.6.0 - The ``fuse_subprocess_stages`` argument. + .. seealso:: + + :py:meth:`~spdl.pipeline.PipelineBuilder.to` + Run a region of stages in a subprocess (or subinterpreter) worker pool. """ return build_pipeline( self.get_config(), @@ -521,5 +497,4 @@ def build( task_hook_factory=task_hook_factory, stage_id=stage_id, use_thread_output_queue=use_thread_output_queue, - fuse_subprocess_stages=fuse_subprocess_stages, ) diff --git a/src/spdl/pipeline/_common/_convert.py b/src/spdl/pipeline/_common/_convert.py index 8b2723da4..5cac569c4 100644 --- a/src/spdl/pipeline/_common/_convert.py +++ b/src/spdl/pipeline/_common/_convert.py @@ -202,11 +202,8 @@ def convert_to_async( op = op.__call__ if inspect.iscoroutinefunction(op) or inspect.isasyncgenfunction(op): - # 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. + # op is async function. No need to convert. + assert executor is None # This has been checked in `PipelineBuilder.pipe()` return op # pyre-ignore: [7] if inspect.isgeneratorfunction(op): diff --git a/src/spdl/pipeline/_fuse.py b/src/spdl/pipeline/_fuse.py index 881707ac7..51dbc390c 100644 --- a/src/spdl/pipeline/_fuse.py +++ b/src/spdl/pipeline/_fuse.py @@ -6,28 +6,25 @@ # pyre-strict -"""Detect runs of consecutive pipe stages that can be fused into one subprocess sub-pipeline. - -When several consecutive :py:func:`~spdl.pipeline.defs.Pipe` stages share the **same** -process-pool (or interpreter-pool) executor instance, SPDL round-trips data back to the main -process between every stage (each stage compiles to ``loop.run_in_executor(executor, op, item)`` -in :py:mod:`spdl.pipeline._common._convert`). That inter-stage IPC is wasteful and fails -outright when an intermediate value is unpicklable. - -This module finds the maximal *fusable runs* — spans of the config's ``pipes`` that can be -combined into a single nested :py:class:`~spdl.pipeline.Pipeline` executed inside the worker -pool, so the op→op handoff stays in-process. The actual config rewrite and worker-pool spawning -build on the runs found here. - -A :py:class:`~spdl.pipeline.defs.PathVariantsConfig` whose branches all use the same pool -executor is fusable as a unit: the whole routing construct moves into the nested pipeline, where -the router and any async/thread branch stages run on the worker's loop/threads. Because the -nested pipeline is built by the normal :py:func:`~spdl.pipeline._build.build_pipeline`, which -already supports path-variants, no execution machinery is special-cased here — only detection and -the executor-stripping rewrite know about it. - -The detection is a pure function over the stage configs (no process spawning), which keeps the -fusion policy easy to reason about and to test in isolation. +"""Fuse each ``.to()`` execution region into one nested subprocess (or subinterpreter) pipeline. + +A :py:class:`~spdl.pipeline.defs.PlacementConfig` marker (appended by +:py:meth:`~spdl.pipeline.PipelineBuilder.to`) designates where the stages that follow it run: a +subprocess/subinterpreter worker pool, or back in the main process. This module walks a config's +``pipes``, and replaces each maximal span of stages under a worker-pool target with a single +stage that executes the span as one nested :py:class:`~spdl.pipeline.Pipeline` inside the pool — +so the op→op handoff stays in the worker and intermediate values are never round-tripped (or +pickled) back to the main process. + +Every stage in a region is carried into the worker, including +:py:class:`~spdl.pipeline.defs.AggregateConfig`, +:py:class:`~spdl.pipeline.defs.DisaggregateConfig`, and +:py:class:`~spdl.pipeline.defs.PathVariantsConfig`. Because the nested pipeline is built by the +normal :py:func:`~spdl.pipeline._build.build_pipeline`, no execution machinery is special-cased +here — only the segmentation on markers and the executor-stripping rewrite. + +The rewrite is a pure function over the stage configs (aside from spawning the worker pools), +which keeps the fusion policy easy to reason about and to test in isolation. """ from __future__ import annotations @@ -39,14 +36,11 @@ import threading import warnings from collections.abc import Sequence -from concurrent.futures import Executor -from dataclasses import dataclass, replace +from dataclasses import replace from functools import partial from typing import Any -from spdl.pipeline._common._convert import _is_isolating_pool from spdl.pipeline._components import _get_global_id, _set_global_id -from spdl.pipeline._executor_proxy import _ensure_executor_unused from spdl.pipeline._subprocess_pipeline_pool import ( _InterpreterBackend, _PoolBackend, @@ -55,7 +49,6 @@ ) from spdl.pipeline.defs._defs import ( _MainProcess, - _PipeType, _SubprocessPipelineConfig, InterpreterPoolExecutorConfig, MergeConfig, @@ -69,11 +62,7 @@ ) __all__ = [ - "_FusableRun", - "_find_fusable_runs", "_fuse_marked_regions", - "_fuse_subprocess_stages", - "_strip_async_executor_tags", ] # Buffer size for the fused sub-pipeline's sink. Small: the worker drains it straight onto the @@ -81,166 +70,6 @@ _FUSED_SINK_BUFFER: int = 3 -@dataclass(frozen=True) -class _FusableRun: - """A maximal span of ``pipes`` that can be fused into one subprocess sub-pipeline. - - The fused stages are ``pipes[start:stop]`` — *adjacent* stages sharing :py:attr:`executor`: - pool-pipes, and path-variants stages whose branches all use it. A span of two or more always - fuses; a single path-variants stage fuses on its own (its branches otherwise round-trip per - stage). Aggregate/disaggregate micro-stages are never absorbed (so their batching semantics - are unchanged) and instead bound a run. - """ - - start: int - """Index of the first fused stage in ``pipes`` (inclusive).""" - - stop: int - """Index just past the last fused stage in ``pipes`` (exclusive).""" - - executor: Executor - """The isolating-pool executor instance shared by every pool-pipe in the run.""" - - -def _fusable_pool_executor(cfg: object) -> Executor | None: - """Return the executor if ``cfg`` is a completion-ordered isolating-pool pipe. - - Returns ``None`` otherwise. Input-ordered (``output_order="input"``) pipes are not fusable: - their global input order cannot be preserved across a pool of independent workers. - """ - if not isinstance(cfg, PipeConfig): - return None - if cfg._type is not _PipeType.Pipe: - return None - executor = cfg._args.executor - if executor is not None and _is_isolating_pool(executor): - return executor - return None - - -def _scan_variant_pool_executors( - cfg: PathVariantsConfig[Any], acc: list[Executor] -) -> bool: - """Collect the isolating-pool executors of every pool-pipe in ``cfg``'s branches. - - 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"``) *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: - if isinstance(stage, PathVariantsConfig): - if not _scan_variant_pool_executors(stage, acc): - return False - continue - executor = _fusable_pool_executor(stage) - if executor is not None: - acc.append(executor) - elif ( - 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 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 - - -def _path_variants_executor(cfg: object) -> Executor | None: - """Return the shared isolating-pool executor if ``cfg`` is a fusable path-variants stage. - - A :py:class:`~spdl.pipeline.defs.PathVariantsConfig` is fusable when every pool-pipe across - all of its branches (recursively) shares the **same** isolating-pool executor instance — the - whole routing construct then moves into one worker's nested pipeline, where the router and any - async/thread branch stages run on the worker's loop/threads and the pool stages run on its - thread pool. Returns ``None`` for a non-path-variants stage, a stage with no pool-pipe (no - isolation to gain), or branches that mix more than one pool executor / hold an input-ordered - pool-pipe. - """ - if not isinstance(cfg, PathVariantsConfig): - return None - executors: list[Executor] = [] - if not _scan_variant_pool_executors(cfg, executors) or not executors: - return None - first = executors[0] - return first if all(e is first for e in executors) else None - - -def _fusable_executor(cfg: object) -> Executor | None: - """The isolating-pool executor a stage fuses onto: a pool-pipe's, or a path-variants' - shared one. ``None`` if the stage is not fusable.""" - return _fusable_pool_executor(cfg) or _path_variants_executor(cfg) - - -def _scan_run( - pipes: Sequence[object], start: int, executor: Executor -) -> tuple[_FusableRun | None, int]: - """Scan one fusable run of adjacent same-executor fusable stages from ``pipes[start]``. - - Greedily consumes the immediately-following stages that share ``executor`` — pool-pipes and - path-variants stages whose branches all use it — stopping at the first stage that is anything - else: a different executor, a thread/None/async pipe, an aggregate/disaggregate micro-stage, - or the end. A micro-stage between two fusable stages splits them into separate runs and keeps - its main-process batching semantics. - - A run of two or more stages always fuses. A lone stage fuses only if it is a path-variants - stage — its branches otherwise round-trip per stage, so moving the whole construct into one - worker is the win — whereas a lone pool-pipe gains nothing from fusion. - - Returns: - A tuple ``(run, next_index)``. ``run`` is the emitted :py:class:`_FusableRun`, or - ``None`` for a lone pool-pipe. ``next_index`` is where the outer scan should resume. - """ - j = start + 1 - n = len(pipes) - while j < n and _fusable_executor(pipes[j]) is executor: - j += 1 - if j - start >= 2 or isinstance(pipes[start], PathVariantsConfig): - return _FusableRun(start, j, executor), j - return None, j - - -def _find_fusable_runs(pipes: Sequence[object]) -> list[_FusableRun]: - """Find all maximal fusable runs in a pipeline config's ``pipes``. - - A run is a span of adjacent stages sharing the same isolating-pool executor instance: either - ≥2 such stages (completion-ordered :py:func:`~spdl.pipeline.defs.Pipe` stages and/or - path-variants stages), or a single :py:class:`~spdl.pipeline.defs.PathVariantsConfig` (whose - branches round-trip per stage when unfused). A lone pool-pipe, and any aggregate/disaggregate - micro-stage, break a run and are left unchanged in the main process. Returned runs are ordered - by position and never overlap. - - Args: - pipes: The ordered stage configs from a :py:class:`~spdl.pipeline.defs.PipelineConfig`. - - Returns: - The list of fusable runs; empty when no fusion applies. - """ - runs: list[_FusableRun] = [] - i = 0 - n = len(pipes) - while i < n: - executor = _fusable_executor(pipes[i]) - if executor is None: - i += 1 - continue - run, i = _scan_run(pipes, i, executor) - if run is not None: - runs.append(run) - return runs - - -################################################################################ -# Config rewrite: replace each fusable run with one subprocess-pipeline stage. -################################################################################ - - def _has_continuous_source(config: PipelineConfig[Any]) -> bool: """Whether any source in the (possibly merged) config re-iterates continuously.""" src = config.src @@ -307,15 +136,6 @@ def _worker_initializer( user_initializer(*user_initargs) -def _pool_params(executor: Executor) -> tuple[int, Any, tuple[Any, ...]]: - """Read worker count and initializer off a fresh pool executor without using it.""" - _ensure_executor_unused(executor) - max_workers = getattr(executor, "_max_workers", None) or os.cpu_count() or 1 - user_initializer = getattr(executor, "_initializer", None) - user_initargs = getattr(executor, "_initargs", ()) or () - return max_workers, user_initializer, user_initargs - - def _warn_fork_with_threads(ctx: Any, stacklevel: int) -> None: if ctx.get_start_method() == "fork" and threading.active_count() > 1: warnings.warn( @@ -339,9 +159,8 @@ def _build_fused_stage_core( continuous: bool, ) -> tuple[_SubprocessPipelineConfig, _SubprocessPipelinePool]: """Spawn a worker pool that runs ``stages`` as a nested pipeline; 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`); they differ only in - where the pool params and backend come from (a live executor vs. a spec).""" + replacement stage. Called by :py:func:`_build_fused_stage_from_spec` once the pool + parameters and backend have been resolved from a region's spec.""" stripped = [_strip_executor(s) for s in stages] sub_config: PipelineConfig[Any] = PipelineConfig( src=SourceConfig( @@ -370,27 +189,6 @@ def _build_fused_stage_core( 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, - backend=_ProcessBackend(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: ProcessPoolExecutorConfig | InterpreterPoolExecutorConfig, @@ -418,78 +216,6 @@ def _build_fused_stage_from_spec( ) -def _fuse_subprocess_stages( - config: PipelineConfig[Any], - *, - mp_context: str | None = None, - report_stats_interval: float = -1, - stacklevel: int = 3, -) -> tuple[PipelineConfig[Any], list[_SubprocessPipelinePool]]: - """Replace fusable runs of pool-pipe stages with single subprocess-pipeline stages. - - For each run found by :py:func:`_find_fusable_runs`, spawns a worker pool that runs the run - as a nested pipeline and replaces the run with a :py:class:`_SubprocessPipelineConfig`. The - spawned pools are returned so the caller can reap them at teardown. The input ``config`` is - not mutated. - - For a continuous-source pipeline, the fused workers run continuous sub-pipelines and the - bridge stage propagates epoch boundaries across the pool (see - :py:mod:`spdl.pipeline._components._subprocess_pipe`). - - Args: - config: The pipeline configuration to rewrite. - mp_context: Multiprocessing start-method name, as accepted by - :py:func:`multiprocessing.get_context`. - report_stats_interval: Forwarded to the nested ``build_pipeline`` so per-stage stats are - reported from inside the workers. - stacklevel: ``warnings.warn`` stack level measured at this function, set by the - caller so warnings point at the user's call site. The fork-with-threads warning - is one frame deeper and uses ``stacklevel + 1``. - - Returns: - A tuple ``(new_config, pools)``. ``pools`` is empty when no fusion applies. - """ - runs = _find_fusable_runs(config.pipes) - if not runs: - return config, [] - - continuous = _has_continuous_source(config) - ctx = mp.get_context(mp_context) - _warn_fork_with_threads(ctx, stacklevel + 1) - - pools: list[_SubprocessPipelinePool] = [] - by_start = {r.start: r for r in runs} - pipes = list(config.pipes) - new_pipes: list[object] = [] - i = 0 - try: - while i < len(pipes): - run = by_start.get(i) - if run is None: - new_pipes.append(pipes[i]) - i += 1 - continue - fused, pool = _build_fused_stage( - pipes[run.start : run.stop], - run.executor, - ctx, - report_stats_interval, - continuous, - ) - pools.append(pool) - new_pipes.append(fused) - i = run.stop - except BaseException: - # Each pool spawns its workers eagerly, so a failure partway through this loop (e.g. a - # later pool fails to spawn) would otherwise leak the workers already started by the - # pools built so far. Reap them before propagating, since the caller never receives the - # ``pools`` list to reap them itself. - for pool in pools: - pool.shutdown() - raise - 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, PlacementConfig) for p in pipes) @@ -505,16 +231,14 @@ def _fuse_marked_regions( Walks ``config.pipes`` tracking the current execution target set by :py:class:`~spdl.pipeline.defs.PlacementConfig` 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. + process). Every maximal span of stages under a subprocess/subinterpreter target — pipes *and* + the aggregate/disaggregate/path-variants stages between them — 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. + run unconditionally. Args: config: The pipeline configuration to rewrite. @@ -581,43 +305,3 @@ def _flush() -> None: 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] diff --git a/src/spdl/pipeline/defs/_defs.py b/src/spdl/pipeline/defs/_defs.py index 954223a0a..46e707032 100644 --- a/src/spdl/pipeline/defs/_defs.py +++ b/src/spdl/pipeline/defs/_defs.py @@ -21,7 +21,6 @@ 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 @@ -219,14 +218,8 @@ class PipeConfig(Generic[T, U]): def __post_init__(self) -> None: op = self._args.op if inspect.iscoroutinefunction(op) or inspect.isasyncgenfunction(op): - 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 self._args.executor is not None: + raise ValueError("`executor` cannot be specified when op is async.") if inspect.isasyncgenfunction(op): if self._type == _PipeType.OrderedPipe: raise ValueError( @@ -427,12 +420,13 @@ def __repr__(self) -> str: @dataclass(frozen=True) class _SubprocessPipelineConfig: - """Internal config for a fused run of pipe stages run in a subprocess worker pool. + """Internal config for a fused region of pipe stages run in a worker pool. - Not user-facing. Produced by the fusion pass (``_fuse_subprocess_stages``), which replaces a - span of consecutive pool-pipe stages with a single stage that executes the run as one nested - pipeline inside the worker pool — eliminating the inter-stage IPC of running each stage on a - process/interpreter pool separately. + Not user-facing. Produced by the fusion pass (``_fuse_marked_regions``), which replaces a + ``.to()`` region — the span of stages between two + :py:class:`~spdl.pipeline.defs.PlacementConfig` markers — with a single stage that executes + the span as one nested pipeline inside a subprocess (or subinterpreter) worker pool, + eliminating the inter-stage IPC of round-tripping data back to the main process. """ name: str @@ -908,11 +902,7 @@ 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. - 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. + It is invalid to provide this argument when the given op is already async. 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. diff --git a/tests/pipeline/builder_to_test.py b/tests/pipeline/builder_to_test.py index 7f543f86b..61e742d7e 100644 --- a/tests/pipeline/builder_to_test.py +++ b/tests/pipeline/builder_to_test.py @@ -8,10 +8,10 @@ import sys import unittest -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from typing import Any -from spdl.pipeline import build_pipeline, Pipeline, PipelineBuilder +from spdl.pipeline import Pipeline, PipelineBuilder from spdl.pipeline.defs import ( _MainProcess, InterpreterPoolExecutorConfig, @@ -30,6 +30,10 @@ def times_two(x: int) -> int: return x * 2 +async def aadd_one(x: int) -> int: + return x + 1 + + def _run(pipeline: Pipeline[Any], timeout: float = 60.0) -> list[Any]: with pipeline.auto_stop(): return list(pipeline.get_iterator(timeout=timeout)) @@ -194,3 +198,31 @@ def test_subprocess_region_runs(self) -> None: .build(num_threads=4) ) self.assertEqual(sorted(_run(pipeline)), sorted((x + 1) * 2 for x in range(n))) + + +class AsyncOpExecutorRejectedTest(unittest.TestCase): + """An async op may not be given an executor. + + Regions place stages by marker, so an async op inside a region carries no executor; the + per-stage ``executor=`` on an async op (previously accepted as a fusion-group tag) is + rejected again. + """ + + def test_pipe_async_with_executor_raises(self) -> None: + """pipe(async_op, executor=...) is rejected regardless of executor type. + + The validation only checks that an executor was supplied (no type-specific branch), so an + isolating (process) and a non-isolating (thread) pool are rejected identically. + """ + for factory in (ProcessPoolExecutor, ThreadPoolExecutor): + with self.subTest(executor=factory.__name__): + b = PipelineBuilder().add_source(range(4)) + with factory(max_workers=1) as ex: + with self.assertRaises(ValueError): + b.pipe(aadd_one, executor=ex) + + def test_pipe_config_async_with_executor_raises(self) -> None: + """The rejection is enforced at the config layer (Pipe factory).""" + with ProcessPoolExecutor(max_workers=1) as ex: + with self.assertRaises(ValueError): + Pipe(aadd_one, executor=ex) diff --git a/tests/pipeline/fuse_test.py b/tests/pipeline/fuse_test.py deleted file mode 100644 index 1a329020f..000000000 --- a/tests/pipeline/fuse_test.py +++ /dev/null @@ -1,162 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import unittest -from collections.abc import Callable -from concurrent.futures import Executor, ProcessPoolExecutor -from typing import Any - -from spdl.pipeline._fuse import _find_fusable_runs, _FusableRun -from spdl.pipeline.defs import Aggregate, Disaggregate, Pipe - - -def op(x: int) -> int: - return x - - -async def aop(x: int) -> int: - return x - - -class _FakeProcessPool(Executor): - """Stand-in that ``_is_process_pool`` recognizes, without spawning processes.""" - - _pool_executor_class: type[Executor] = ProcessPoolExecutor - - def submit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> Any: - raise NotImplementedError - - -class _FakeThreadPool(Executor): - """Stand-in for a non-isolating (thread) executor — not a fusion target.""" - - def submit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> Any: - raise NotImplementedError - - -def _pool_pipe(executor: Executor) -> Any: - return Pipe(op, executor=executor, concurrency=2) - - -class FindFusableRunsTest(unittest.TestCase): - def test_empty_pipes(self) -> None: - """No pipes yields no fusable runs.""" - self.assertEqual(_find_fusable_runs([]), []) - - def test_two_same_executor_pipes_fuse(self) -> None: - """Two adjacent pipes on the same pool executor fuse into one run.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), _pool_pipe(ex)]) - self.assertEqual(runs, [_FusableRun(0, 2, ex)]) - - def test_aggregate_between_pool_pipes_breaks_run(self) -> None: - """An aggregate between two pool pipes is not absorbed; the lone pipes do not fuse.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), Aggregate(4), _pool_pipe(ex)]) - self.assertEqual(runs, []) - - def test_disaggregate_between_pool_pipes_breaks_run(self) -> None: - """A disaggregate between two pool pipes is not absorbed; the lone pipes do not fuse.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), Disaggregate(), _pool_pipe(ex)]) - self.assertEqual(runs, []) - - def test_aggregate_between_two_runs_fuses_each_side(self) -> None: - """An aggregate splits the pipes; each adjacent same-executor pair fuses on its own.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs( - [ - _pool_pipe(ex), - _pool_pipe(ex), - Aggregate(4), - _pool_pipe(ex), - _pool_pipe(ex), - ] - ) - self.assertEqual(runs, [_FusableRun(0, 2, ex), _FusableRun(3, 5, ex)]) - - def test_trailing_aggregate_not_fused(self) -> None: - """A trailing aggregate is left in the main process; only the pool pipes fuse.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), _pool_pipe(ex), Aggregate(4)]) - self.assertEqual(runs, [_FusableRun(0, 2, ex)]) - - def test_single_pool_pipe_with_trailing_aggregate_not_fused(self) -> None: - """A lone pool pipe plus a trailing aggregate has no same-executor neighbour: no run.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), Aggregate(4)]) - self.assertEqual(runs, []) - - def test_trailing_disaggregate_not_fused(self) -> None: - """A trailing disaggregate is left in the main process; only the pool pipes fuse.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), _pool_pipe(ex), Disaggregate()]) - self.assertEqual(runs, [_FusableRun(0, 2, ex)]) - - def test_different_executors_do_not_fuse(self) -> None: - """Adjacent pool pipes on different executor instances do not fuse.""" - runs = _find_fusable_runs( - [_pool_pipe(_FakeProcessPool()), _pool_pipe(_FakeProcessPool())] - ) - self.assertEqual(runs, []) - - def test_thread_pool_breaks_run(self) -> None: - """A thread-pool stage between two pool pipes breaks the run (no fusion).""" - ex = _FakeProcessPool() - runs = _find_fusable_runs( - [_pool_pipe(ex), _pool_pipe(_FakeThreadPool()), _pool_pipe(ex)] - ) - self.assertEqual(runs, []) - - def test_default_executor_breaks_run(self) -> None: - """A stage with no executor (default thread pool) breaks the run.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), Pipe(op), _pool_pipe(ex)]) - self.assertEqual(runs, []) - - def test_async_op_breaks_run(self) -> None: - """An async stage breaks the run.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), Pipe(aop), _pool_pipe(ex)]) - self.assertEqual(runs, []) - - def test_input_ordered_pool_pipe_not_fused(self) -> None: - """An input-ordered pool pipe is not fusable and breaks the run.""" - ex = _FakeProcessPool() - ordered = Pipe(op, executor=ex, output_order="input") - runs = _find_fusable_runs([_pool_pipe(ex), ordered, _pool_pipe(ex)]) - self.assertEqual(runs, []) - - def test_lone_pool_pipe_untouched(self) -> None: - """A single pool pipe with nothing to combine is left unchanged.""" - self.assertEqual(_find_fusable_runs([_pool_pipe(_FakeProcessPool())]), []) - - def test_leading_micro_stage_stays_in_main(self) -> None: - """A micro-stage before the first pool pipe stays in the main process.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([Aggregate(4), _pool_pipe(ex), _pool_pipe(ex)]) - self.assertEqual(runs, [_FusableRun(1, 3, ex)]) - - def test_multiple_runs(self) -> None: - """Several disjoint fusable runs are all detected.""" - ex1, ex2 = _FakeProcessPool(), _FakeProcessPool() - pipes = [ - _pool_pipe(ex1), - _pool_pipe(ex1), - Pipe(op), # breaker - _pool_pipe(ex2), - _pool_pipe(ex2), - ] - runs = _find_fusable_runs(pipes) - self.assertEqual(runs, [_FusableRun(0, 2, ex1), _FusableRun(3, 5, ex2)]) - - def test_same_executor_separated_by_breaker_not_fused(self) -> None: - """The same executor reused across a breaker yields no paying run.""" - ex = _FakeProcessPool() - runs = _find_fusable_runs([_pool_pipe(ex), Pipe(op), _pool_pipe(ex)]) - self.assertEqual(runs, []) diff --git a/tests/pipeline/marked_region_fuse_test.py b/tests/pipeline/marked_region_fuse_test.py index b9294e0e0..c9fa76603 100644 --- a/tests/pipeline/marked_region_fuse_test.py +++ b/tests/pipeline/marked_region_fuse_test.py @@ -9,11 +9,22 @@ import multiprocessing as mp import os import sys +import threading +import time import unittest -from collections.abc import Sequence +from collections.abc import AsyncIterator, Iterator, Sequence from typing import Any -from spdl.pipeline import build_pipeline, Pipeline, run_pipeline_in_subprocess +from spdl.pipeline import ( + AsyncQueue, + build_pipeline, + Pipeline, + PipelineBuilder, + PipelineFailure, + run_pipeline_in_subprocess, +) +from spdl.pipeline._fuse import _fuse_marked_regions +from spdl.pipeline.config import set_default_queue_class from spdl.pipeline.defs import ( Aggregate, InterpreterPoolExecutorConfig, @@ -25,6 +36,7 @@ SinkConfig, SourceConfig, ) +from spdl.pipeline.defs._defs import _SubprocessPipelineConfig def add_one(x: int) -> int: @@ -223,6 +235,464 @@ def test_region_open_at_end_of_pipes(self) -> None: self.assertEqual(sorted(_run(pipeline)), sorted((x + 1) * 2 for x in range(n))) +# --------------------------------------------------------------------------- +# Behavioral coverage ported from the removed identity-fusion tests +# (subprocess_pipeline_fuse_test.py), now expressed via the public ``.to()`` API. +# These assert the guarantees the old ``fuse_subprocess_stages`` feature provided. +# --------------------------------------------------------------------------- + + +def boom(x: int) -> int: + raise ValueError("boom") + + +def _raise_initializer() -> None: + raise RuntimeError("initializer boom") + + +def dup(x: int) -> Iterator[int]: + """A sync-generator op: yields two values per input (1->2 fan-out).""" + yield x * 2 + yield x * 2 + 1 + + +async def adup(x: int) -> AsyncIterator[int]: + """An async-generator op: yields two values per input (1->2 fan-out).""" + yield x + yield x + 1 + + +async def _astamp(item: _PidStamp) -> _PidStamp: + """An async op that records its pid; runs on the worker's event loop inside the region.""" + item.pids["async"] = os.getpid() + return item + + +def _stamp_branch(item: _PidStamp) -> _PidStamp: + """Records its pid; used as a plain region stage and as a path-variants branch.""" + item.pids["branch"] = os.getpid() + return item + + +def _stamp_merge(item: _PidStamp) -> _PidStamp: + """A post-merge / fan-in region stage that records its pid.""" + item.pids["merge"] = os.getpid() + return item + + +def _stamp_router(item: _PidStamp) -> int: + """Path-variants router: records its pid, then routes by parity.""" + item.pids["router"] = os.getpid() + return item.value % 2 + + +def route_by_parity(x: int) -> int: + """Router: even items to path 0, odd items to path 1.""" + return x % 2 + + +def stall_on_first(x: int) -> int: + """Stall only on the earliest item, so one worker holds it while its peer races ahead.""" + if x == 0: + time.sleep(3.0) + return x + 1 + + +class _Buffer1Queue(AsyncQueue): + """An ``AsyncQueue`` pinned to ``buffer_size=1`` (no prefetch slack).""" + + def __init__( + self, info: Any, *, buffer_size: int = 1, interval: float = -1 + ) -> None: + super().__init__(info, buffer_size=1) + + +def _install_small_buffers() -> None: + """Region-worker initializer: pin the worker sub-pipeline's queues to ``buffer_size=1``. + + A single-slot buffer stops the worker holding the earliest item from prefetching past it, so + that worker stays parked on the slow item while its peer drains the rest -- surfacing the + startup race where an early-finishing worker could otherwise steal a slow peer's end marker. + """ + set_default_queue_class(_Buffer1Queue) + + +class RegionBehaviorTest(unittest.TestCase): + """End-to-end guarantees the fusing feature provided, now via ``.to()`` regions.""" + + def test_generator_op_in_region(self) -> None: + """A sync-generator op inside a region fans out 1->N in the worker sub-pipeline.""" + n = 8 + # dup yields {2x, 2x+1}; add_one shifts each by one, covering 1..2n exactly. + ref = list(range(1, 2 * n + 1)) + pipeline = ( + PipelineBuilder() + .add_source(range(n)) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(dup, concurrency=2) + .pipe(add_one, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + self.assertEqual(sorted(_run(pipeline)), ref) + + def test_async_op_in_region_runs_in_worker(self) -> None: + """An async op inside a region runs on the worker's loop, in the worker process. + + The sync stages on either side and the async op in the middle all stamp ``os.getpid()``; + the assertion proves the async stage shares the worker process with its neighbours rather + than hopping back to the main process. + """ + n = 16 + pipeline = ( + PipelineBuilder() + .add_source([_PidStamp(x) for x in range(n)]) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(_stamp_branch, concurrency=2) + .pipe(_astamp, concurrency=2) + .pipe(_stamp_merge, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + main_pid = os.getpid() + results = _run(pipeline) + self.assertEqual(len(results), n) + for item in results: + self.assertEqual(set(item.pids), {"branch", "async", "merge"}) + self.assertEqual(item.pids["branch"], item.pids["async"]) + self.assertEqual(item.pids["async"], item.pids["merge"]) + self.assertNotEqual(item.pids["async"], main_pid) + + def test_async_generator_op_in_region(self) -> None: + """An async-generator op inside a region fans out inside the worker.""" + n = 8 + # add_one -> v=x+1; adup(v) yields {v, v+1}. + ref = sorted(v for x in range(n) for v in (x + 1, x + 2)) + pipeline = ( + PipelineBuilder() + .add_source(range(n)) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(add_one, concurrency=2) + .pipe(adup, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + self.assertEqual(sorted(_run(pipeline)), ref) + + def test_initializer_failure_surfaces_not_hangs(self) -> None: + """A region-worker initializer that raises surfaces as a failure instead of hanging. + + A real failure surfaces as :py:class:`PipelineFailure`; a hung pipeline would instead + raise :py:class:`TimeoutError` from the finite ``get_iterator`` timeout, so asserting on + ``PipelineFailure`` proves the failed initializer is reported rather than wedging forever. + """ + pipeline = ( + PipelineBuilder() + .add_source(range(2)) + .to( + ProcessPoolExecutorConfig(max_workers=2, initializer=_raise_initializer) + ) + .pipe(add_one, concurrency=2) + .pipe(times_two, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(4) + .build(num_threads=4) + ) + with self.assertRaises(PipelineFailure): + _run(pipeline, timeout=30.0) + + def test_slow_worker_does_not_drop_earliest_items(self) -> None: + """A worker stalled on the earliest item still delivers it instead of dropping it.""" + n = 16 + ref = sorted((x + 1) * 2 for x in range(n)) + pipeline = ( + PipelineBuilder() + .add_source(range(n)) + .to( + ProcessPoolExecutorConfig( + max_workers=2, initializer=_install_small_buffers + ) + ) + .pipe(stall_on_first, concurrency=1) + .pipe(times_two, concurrency=1) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + self.assertEqual(sorted(_run(pipeline)), ref) + + def test_lone_path_variants_in_region(self) -> None: + """A path-variants stage inside a region routes and runs entirely in the worker. + + Even items take the ``add_one`` branch, odd items the ``times_two`` branch; the whole + routing construct moves into the worker, so the result must match an inline run. + """ + n = 16 + ref = sorted(x + 1 if x % 2 == 0 else x * 2 for x in range(n)) + pipeline = ( + PipelineBuilder() + .add_source(range(n)) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .path_variants(route_by_parity, [[Pipe(add_one)], [Pipe(times_two)]]) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + self.assertEqual(sorted(_run(pipeline)), ref) + + def test_path_variants_process_locality_in_region(self) -> None: + """Router, branches, and the post-merge stage all run in one worker process. + + Each stage stamps ``os.getpid()``; the assertion proves the whole + router -> branch -> merge chain runs inside a single worker process per item (distinct + from the main process), i.e. fusion leaves no per-stage process hop. + """ + n = 16 + pipeline = ( + PipelineBuilder() + .add_source([_PidStamp(x) for x in range(n)]) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .path_variants( + _stamp_router, [[Pipe(_stamp_branch)], [Pipe(_stamp_branch)]] + ) + .pipe(_stamp_merge) # fan-in / post-merge, same region + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + results = _run(pipeline) + self.assertEqual(len(results), n) + main_pid = os.getpid() + worker_pids = set() + for item in results: + self.assertEqual(set(item.pids), {"router", "branch", "merge"}) + self.assertEqual(item.pids["router"], item.pids["branch"]) + self.assertEqual(item.pids["branch"], item.pids["merge"]) + self.assertNotEqual(item.pids["router"], main_pid) + worker_pids.add(item.pids["router"]) + self.assertNotIn(main_pid, worker_pids) + + def test_run_pipeline_in_subprocess_with_region(self) -> None: + """A ``.to()`` region composes with ``run_pipeline_in_subprocess``.""" + n = 16 + ref = sorted((x + 1) * 2 for x in range(n)) + config = ( + PipelineBuilder() + .add_source(range(n)) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(add_one, concurrency=2) + .pipe(times_two, concurrency=3) + .to(MAIN_PROCESS) + .add_sink(n) + .get_config() + ) + src = run_pipeline_in_subprocess(config, num_threads=4) + self.assertEqual(sorted(src), ref) + + def test_run_pipeline_in_subprocess_unpicklable_intermediate(self) -> None: + """The region handle survives into the pipeline subprocess and the unpicklable + op->op handoff stays inside a worker.""" + n = 12 + ref = sorted((x + 1) * 2 for x in range(n)) + config = ( + PipelineBuilder() + .add_source(range(n)) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(wrap, concurrency=2) + .pipe(unwrap, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(n) + .get_config() + ) + src = run_pipeline_in_subprocess(config, num_threads=4) + self.assertEqual(sorted(src), ref) + + +class ContinuousRegionTest(unittest.TestCase): + """A ``.to()`` region under a continuous source stays warm and correct across epochs.""" + + def test_multi_epoch_correct(self) -> None: + """A continuous region yields the correct set each epoch from the same warm pool.""" + n = 12 + ref = sorted((x + 1) * 2 for x in range(n)) + pipeline = ( + PipelineBuilder() + .add_source(range(n), continuous=True) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(add_one, concurrency=2) + .pipe(times_two, concurrency=3) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + with pipeline.auto_stop(): + for _ in range(3): # three epochs from the same warm worker pool + self.assertEqual(sorted(pipeline.get_iterator(timeout=60)), ref) + + def test_fewer_items_than_workers(self) -> None: + """An epoch with fewer items than workers still completes (some workers run empty).""" + n = 2 + ref = sorted((x + 1) * 2 for x in range(n)) + pipeline = ( + PipelineBuilder() + .add_source(range(n), continuous=True) + .to(ProcessPoolExecutorConfig(max_workers=4)) + .pipe(add_one, concurrency=2) + .pipe(times_two, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + with pipeline.auto_stop(): + for _ in range(2): + self.assertEqual(sorted(pipeline.get_iterator(timeout=60)), ref) + + def test_op_failure_does_not_deadlock(self) -> None: + """Op failures (dropped per SPDL default) still let each epoch's barrier complete. + + ``boom`` raises on every item, so the region produces no results; the continuous epoch + barrier must still complete each epoch (workers report the boundary with zero results) + instead of deadlocking. + """ + n = 8 + pipeline = ( + PipelineBuilder() + .add_source(range(n), continuous=True) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(add_one, concurrency=2) + .pipe(boom, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + with pipeline.auto_stop(): + for _ in range(2): + self.assertEqual(list(pipeline.get_iterator(timeout=60)), []) + + def test_teardown_mid_stream_does_not_hang(self) -> None: + """Tearing a continuous region subprocess pipeline down mid-stream must not hang. + + A teardown before the stream is drained leaves the per-worker input queue full; pool + shutdown must cancel the queue feeder-thread join instead of blocking forever flushing + buffered items into a pipe the terminated workers no longer drain. + """ + n = 100_000 # far more than any buffer, so the stream is still full at teardown + config = ( + PipelineBuilder() + .add_source(range(n), continuous=True) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(add_one, concurrency=2) + .pipe(times_two, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(2) + .get_config() + ) + src = run_pipeline_in_subprocess(config, num_threads=4) + it = iter(src) + next(it) # consume a couple of items so the queues stay backed up + next(it) + torn_down = threading.Event() + + def _teardown() -> None: + src._finalizer() # pyre-ignore[16]: documented teardown handle + torn_down.set() + + t = threading.Thread(target=_teardown, daemon=True) + t.start() + self.assertTrue( + torn_down.wait(timeout=60), + "region-pool teardown hung on a full input queue after a mid-stream stop", + ) + t.join(timeout=10) + + def test_continuous_in_subprocess(self) -> None: + """A continuous region composes with ``run_pipeline_in_subprocess`` across epochs.""" + n = 12 + ref = sorted((x + 1) * 2 for x in range(n)) + config = ( + PipelineBuilder() + .add_source(range(n), continuous=True) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(add_one, concurrency=2) + .pipe(times_two, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(n) + .get_config() + ) + src = run_pipeline_in_subprocess(config, num_threads=4) + for _ in range(3): # one epoch per iteration + self.assertEqual(sorted(src), ref) + + +class MarkedRegionSegmentationTest(unittest.TestCase): + """Unit tests for the ``_fuse_marked_regions`` config rewrite. + + Replaces the pure-detection coverage of the removed ``_find_fusable_runs`` tests. Because the + rewrite eagerly spawns worker pools, each fused case reaps its pools in a ``finally``. + """ + + def test_no_markers_returns_config_unchanged(self) -> None: + """With no region markers the rewrite is a no-op: same config, no pools.""" + config = _cfg(range(4), [Pipe(add_one), Pipe(times_two)]) + new_config, pools = _fuse_marked_regions(config) + self.assertEqual(pools, []) + self.assertIs(new_config, config) + + def test_region_becomes_single_fused_stage_absorbing_aggregate(self) -> None: + """A region collapses to one fused stage (absorbing aggregate); main stages are kept.""" + config = _cfg( + range(4), + [ + Pipe(add_one), # main process, before the region + PlacementConfig(target=ProcessPoolExecutorConfig(max_workers=1)), + Pipe(times_two), + Aggregate(2), + PlacementConfig(target=MAIN_PROCESS), + Pipe(add_one), # main process, after the region + ], + ) + new_config, pools = _fuse_marked_regions(config) + try: + kinds = [type(p).__name__ for p in new_config.pipes] + self.assertEqual( + kinds, ["PipeConfig", "_SubprocessPipelineConfig", "PipeConfig"] + ) + fused = [ + p for p in new_config.pipes if isinstance(p, _SubprocessPipelineConfig) + ] + self.assertEqual(len(fused), 1) + finally: + for pool in pools: + pool.shutdown() + + def test_two_regions_produce_two_fused_stages(self) -> None: + """Two separate regions produce two independent fused stages.""" + config = _cfg( + range(4), + [ + PlacementConfig(target=ProcessPoolExecutorConfig(max_workers=1)), + Pipe(add_one), + PlacementConfig(target=MAIN_PROCESS), + Pipe(times_two), + PlacementConfig(target=ProcessPoolExecutorConfig(max_workers=1)), + Pipe(add_one), + PlacementConfig(target=MAIN_PROCESS), + ], + ) + new_config, pools = _fuse_marked_regions(config) + try: + fused = [ + p for p in new_config.pipes if isinstance(p, _SubprocessPipelineConfig) + ] + self.assertEqual(len(fused), 2) + finally: + for pool in pools: + pool.shutdown() + + if sys.version_info >= (3, 14): class SubinterpreterRegionFuseTest(unittest.TestCase): diff --git a/tests/pipeline/region_hooks_test.py b/tests/pipeline/region_hooks_test.py new file mode 100644 index 000000000..da75a6b9c --- /dev/null +++ b/tests/pipeline/region_hooks_test.py @@ -0,0 +1,235 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Per-stage hooks and queues fire inside the workers of a ``.to()`` region. + +Ported from the removed ``subprocess_pipeline_fuse_test.py``: the region worker runs a nested +pipeline, so its default task hooks and queues must fire *in the worker process*. The worker is a +separate process and cannot share in-memory counters with the test, so recording hooks write +small JSON files that the test reads back as evidence. +""" + +import itertools +import json +import os +import tempfile +import threading +import time +import unittest +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any + +from spdl.pipeline import AsyncQueue, Pipeline, PipelineBuilder, TaskHook +from spdl.pipeline.config import set_default_hook_class, set_default_queue_class +from spdl.pipeline.defs import MAIN_PROCESS, ProcessPoolExecutorConfig + + +def add_one(x: int) -> int: + return x + 1 + + +def times_two(x: int) -> int: + return x * 2 + + +def _run(pipeline: Pipeline[Any], timeout: float = 60.0) -> list[Any]: + with pipeline.auto_stop(): + return list(pipeline.get_iterator(timeout=timeout)) + + +# Set inside each region worker process by ``_install_recording_hooks`` (the region initializer). +# The worker is a separate process, so the recording hooks below cannot share in-memory counters +# with the test; each instead writes a small JSON file into this directory, which the test reads +# back to prove the hooks fired inside the worker. +_EVIDENCE_DIR: str | None = None + +# A per-process monotonic counter for evidence filenames. ``id(self)`` is only unique among +# simultaneously-live objects -- CPython recycles addresses, so two recorders created and +# destroyed in sequence could collide on one path and clobber each other's count. This counter +# (combined with ``pid``) is guaranteed distinct for every write within a process. +_EVIDENCE_SEQ = itertools.count() +_EVIDENCE_SEQ_LOCK = threading.Lock() + + +def _write_evidence(record: dict[str, Any]) -> None: + if (d := _EVIDENCE_DIR) is None: + return + with _EVIDENCE_SEQ_LOCK: + seq = next(_EVIDENCE_SEQ) + # ``pid`` proves the write came from a worker process (never the main/test process); ``seq`` + # keeps each writer's file distinct within that process, so no two writers ever race on one + # path. Write to a temp file then atomically rename, so a reader polling mid-write never sees + # a truncated/partial file. + path = os.path.join(d, f"{record['kind']}-{record['pid']}-{seq}.json") + tmp = f"{path}.tmp" + try: + with open(tmp, "w") as f: + f.write(json.dumps(record)) + os.replace(tmp, path) + except OSError: + # This runs inside the worker's stage_hook ``finally`` (stage teardown). It must never + # raise into the stage: an I/O failure here -- e.g. the evidence ``TemporaryDirectory`` + # already torn down in a teardown race -- would otherwise crash the stage and drop items + # from the data path the test is validating. + pass + + +class _RecordingTaskHook(TaskHook): + """A ``TaskHook`` that records how many tasks ran in its (region-worker) stage.""" + + def __init__(self, info: Any, interval: float = -1) -> None: + self.info = info + self.n_tasks = 0 + + @asynccontextmanager + async def stage_hook(self) -> AsyncIterator[None]: + try: + yield + finally: + _write_evidence( + { + "kind": "task", + "name": str(self.info), + "pid": os.getpid(), + "iid": id(self), + "n_tasks": self.n_tasks, + } + ) + + @asynccontextmanager + async def task_hook(self, input_item: Any = None) -> AsyncIterator[None]: + self.n_tasks += 1 + yield + + +class _RecordingQueue(AsyncQueue): + """An ``AsyncQueue`` whose ``stage_hook`` records that it ran in the (region-worker) stage.""" + + def __init__( + self, info: Any, *, buffer_size: int = 1, interval: float = -1 + ) -> None: + super().__init__(info, buffer_size=buffer_size) + self.n_get = 0 + + async def get(self) -> object: + item = await super().get() + self.n_get += 1 + return item + + @asynccontextmanager + async def stage_hook(self) -> AsyncIterator[None]: + try: + yield + finally: + _write_evidence( + { + "kind": "queue", + "name": str(self.info), + "pid": os.getpid(), + "iid": id(self), + "n_get": self.n_get, + } + ) + + +def _install_recording_hooks(evidence_dir: str) -> None: + """Region-worker initializer: runs inside each worker before its sub-pipeline is built. + + The region reads this off the :py:class:`ProcessPoolExecutorConfig` and runs it in every worker + process, so the worker's nested ``build_pipeline`` -- given no explicit + ``task_hook_factory``/``queue_class`` -- picks up these recording classes as its defaults. + """ + global _EVIDENCE_DIR + _EVIDENCE_DIR = evidence_dir + set_default_hook_class(_RecordingTaskHook) + set_default_queue_class(_RecordingQueue) + + +class RegionHookTest(unittest.TestCase): + """Per-stage hooks/queues fire inside the workers of a ``.to()`` region.""" + + def test_hooks_fire_in_region_workers(self) -> None: + """Hooks fire in the workers spawned for a + ``.to(ProcessPoolExecutorConfig(...))`` region.""" + n = 16 + ref = sorted((x + 1) * 2 for x in range(n)) + with tempfile.TemporaryDirectory() as evidence_dir: + pipeline = ( + PipelineBuilder() + .add_source(range(n)) + .to( + ProcessPoolExecutorConfig( + max_workers=2, + initializer=_install_recording_hooks, + initargs=(evidence_dir,), + ) + ) + .pipe(add_one, concurrency=2) + .pipe(times_two, concurrency=2) + .to(MAIN_PROCESS) + .add_sink(n) + .build(num_threads=4) + ) + self.assertEqual(sorted(_run(pipeline)), ref) + self._assert_hooks_fired(self._await_evidence(evidence_dir, n), n) + + def _read_evidence(self, evidence_dir: str) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for name in os.listdir(evidence_dir): + if not name.endswith(".json"): + continue # skip ``.json.tmp`` files still being written + with open(os.path.join(evidence_dir, name)) as f: + records.append(json.loads(f.read())) + return records + + def _await_evidence( + self, evidence_dir: str, n: int, timeout: float = 60.0 + ) -> list[dict[str, Any]]: + """Re-read evidence until the workers' task total lands, then return the records. + + Each worker writes its evidence files during nested-pipeline teardown, a side channel + with no happens-before relationship to the main iterator completing. Polling until the + recorded task total reaches the expected ``2 * n`` closes that read-too-early window + without masking regressions: on a genuine failure (hooks never fire) the poll times out + and the caller's assertions report the shortfall. + """ + deadline = time.monotonic() + timeout + records = self._read_evidence(evidence_dir) + while sum(r["n_tasks"] for r in records if r["kind"] == "task") < 2 * n: + if time.monotonic() > deadline: + break # fall through; the assertions below report the shortfall + time.sleep(0.02) + records = self._read_evidence(evidence_dir) + return records + + def _assert_hooks_fired(self, records: list[dict[str, Any]], n: int) -> None: + task_records = [r for r in records if r["kind"] == "task"] + queue_records = [r for r in records if r["kind"] == "queue"] + + self.assertTrue(task_records, "TaskHook never fired in any region worker") + self.assertTrue( + queue_records, "queue stage_hook never fired in any region worker" + ) + + # Every record came from a worker process, never the main/test process -- this is what + # proves the hooks fired "in the worker". + main_pid = os.getpid() + for r in records: + self.assertNotEqual(r["pid"], main_pid) + + # The two region pipe stages each see every item once: ``add_one`` over n items, then + # ``times_two`` over n items => 2n ``task_hook`` invocations, summed across workers. + self.assertEqual(sum(r["n_tasks"] for r in task_records), 2 * n) + + # The region sub-pipeline's queues carried the data inside the worker(s). + self.assertGreaterEqual(sum(r["n_get"] for r in queue_records), n) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pipeline/subprocess_pipe_bridge_test.py b/tests/pipeline/subprocess_pipe_bridge_test.py new file mode 100644 index 000000000..280aebe29 --- /dev/null +++ b/tests/pipeline/subprocess_pipe_bridge_test.py @@ -0,0 +1,133 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Tests for the subprocess-pipeline bridge internals (``_subprocess_pipe``). + +These exercise the feeder/collector/stall-guard machinery that streams items to and from the +worker pool backing a ``.to()`` region. They are independent of how a region is expressed -- +they drive ``_subprocess_pipe`` helpers directly -- so they live apart from the region tests. +Moved (unchanged) from the removed ``subprocess_pipeline_fuse_test.py``. +""" + +import asyncio +import queue as _queue +import threading +import time +import unittest +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +from spdl.pipeline import AsyncQueue +from spdl.pipeline._components import _subprocess_pipe +from spdl.pipeline._components._common import StageInfo + + +class FeedAbortTest(unittest.TestCase): + """The bridge feeder must wind down promptly when the collector signals abort.""" + + def test_feed_ends_session_when_aborted_while_idle(self) -> None: + """A feeder parked on an empty input queue still emits the per-worker _SESSION_END. + + On a worker error the collector sets ``abort`` while the feeder is typically + blocked waiting on a slow/idle upstream. The feeder must wake and send exactly one + ``_SESSION_END`` onto each worker's own queue so the collector can drain every + ``_DONE`` instead of hanging until its stall timeout. Driving ``_feed`` directly + keeps the abort-while-idle race deterministic. + """ + num_workers = 3 + + async def _scenario() -> list[list[Any]]: + in_qs: list[_queue.Queue[Any]] = [ + _queue.Queue() for _ in range(num_workers) + ] + input_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="input") + ) # stays empty -> get() blocks + abort = asyncio.Event() + feeder_idle = asyncio.Event() + put_stop = threading.Event() + with ThreadPoolExecutor(max_workers=num_workers + 1) as ex: + task = asyncio.ensure_future( + _subprocess_pipe._feed( + input_queue, in_qs, ex, abort, feeder_idle, put_stop + ) + ) + await asyncio.sleep(0.1) # let the feeder park on input_queue.get() + self.assertFalse(task.done(), "feeder should be parked on empty queue") + abort.set() + await asyncio.wait_for(task, timeout=5.0) + return [[q.get_nowait() for _ in range(q.qsize())] for q in in_qs] + + msgs = asyncio.run(_scenario()) + # Every worker's own queue receives exactly one _SESSION_END. + self.assertEqual(msgs, [[(_subprocess_pipe._SESSION_END, None)]] * num_workers) + + +class StallGuardTest(unittest.TestCase): + """The collector's stall guard against an abruptly-dead worker.""" + + def test_check_stall_raises_past_timeout(self) -> None: + """``_check_stall`` raises once no message has arrived for longer than the bound.""" + orig = _subprocess_pipe._WORKER_STALL_TIMEOUT + _subprocess_pipe._WORKER_STALL_TIMEOUT = 0.0 + try: + with self.assertRaises(TimeoutError): + _subprocess_pipe._check_stall(time.monotonic() - 1.0) + finally: + _subprocess_pipe._WORKER_STALL_TIMEOUT = orig + + def test_check_stall_quiet_within_timeout(self) -> None: + """``_check_stall`` does not raise while progress is within the bound.""" + orig = _subprocess_pipe._WORKER_STALL_TIMEOUT + _subprocess_pipe._WORKER_STALL_TIMEOUT = 60.0 + try: + _subprocess_pipe._check_stall(time.monotonic()) # should not raise + finally: + _subprocess_pipe._WORKER_STALL_TIMEOUT = orig + + def test_collect_suppresses_stall_while_feeder_idle(self) -> None: + """An idle feeder suppresses the collector's stall guard during input starvation. + + With the timeout pinned to zero, any stall check on an empty queue would trip instantly; + the collector must instead keep draining while ``feeder_idle`` is set (nothing dispatched, + no worker message due) and still finish once the worker reports ``_DONE``. + """ + orig = _subprocess_pipe._WORKER_STALL_TIMEOUT + _subprocess_pipe._WORKER_STALL_TIMEOUT = 0.0 + + async def _scenario() -> None: + out_q: _queue.Queue[Any] = _queue.Queue() + output_queue = AsyncQueue( + StageInfo(pipeline_id=0, stage_id="0", stage_name="output") + ) + abort = asyncio.Event() + feeder_idle = asyncio.Event() + feeder_idle.set() # feeder parked on an idle upstream -> no message expected + with ThreadPoolExecutor(max_workers=2) as ex: + task = asyncio.ensure_future( + _subprocess_pipe._collect( + out_q, 1, output_queue, ex, abort, feeder_idle + ) + ) + await asyncio.sleep( + 0.6 + ) # several empty poll cycles; must not trip the guard + self.assertFalse( + task.done(), "idle feeder must suppress the stall guard" + ) + out_q.put((_subprocess_pipe._DONE, None)) + await asyncio.wait_for(task, timeout=5.0) + + try: + asyncio.run(_scenario()) + finally: + _subprocess_pipe._WORKER_STALL_TIMEOUT = orig + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pipeline/subprocess_pipeline_fuse_test.py b/tests/pipeline/subprocess_pipeline_fuse_test.py deleted file mode 100644 index 99c6bd957..000000000 --- a/tests/pipeline/subprocess_pipeline_fuse_test.py +++ /dev/null @@ -1,1313 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import asyncio -import itertools -import json -import os -import queue as _queue -import sys -import tempfile -import threading -import time -import unittest -from collections.abc import AsyncIterator, Iterator -from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor -from contextlib import asynccontextmanager -from typing import Any - -from spdl.pipeline import ( - AsyncQueue, - build_pipeline, - Pipeline, - PipelineBuilder, - PipelineFailure, - run_pipeline_in_subprocess, - TaskHook, -) -from spdl.pipeline._components import _subprocess_pipe -from spdl.pipeline._components._common import StageInfo -from spdl.pipeline._fuse import _find_fusable_runs, _strip_async_executor_tags -from spdl.pipeline.config import set_default_hook_class, set_default_queue_class -from spdl.pipeline.defs import Pipe, PipeConfig - - -def add_one(x: int) -> int: - return x + 1 - - -def _raise_initializer() -> None: - raise RuntimeError("initializer boom") - - -def times_two(x: int) -> int: - return x * 2 - - -def boom(x: int) -> int: - raise ValueError("boom") - - -def route_by_parity(x: int) -> int: - """Router: even items to path 0, odd items to path 1.""" - return x % 2 - - -async def aident(x: int) -> int: - """An async (no-executor) branch stage; runs on the worker's loop once fused.""" - return x - - -class _PidStamp: - """A mutable item each stage stamps its ``os.getpid()`` into, to prove process-locality. - - Picklable, so it survives the main -> worker -> main round-trip; the stamps written inside a - worker travel back on the returned copy. - """ - - def __init__(self, value: int) -> None: - self.value = value - self.pids: dict[str, int] = {} - - -def _stamp_router(item: _PidStamp) -> int: - """Path-variants router: records its pid, then routes by parity.""" - item.pids["router"] = os.getpid() - return item.value % 2 - - -def _stamp_branch(item: _PidStamp) -> _PidStamp: - """A branch stage (used on both paths): records its pid.""" - item.pids["branch"] = os.getpid() - return item - - -def _stamp_merge(item: _PidStamp) -> _PidStamp: - """Post-merge (fan-in) stage: records its pid after the branches merge back.""" - item.pids["merge"] = os.getpid() - return item - - -class _Unpicklable: - """An object that refuses to be pickled, to prove it never crosses a process boundary.""" - - def __init__(self, value: int) -> None: - self.value = value - - def __reduce__(self) -> Any: - raise TypeError("_Unpicklable must not be pickled") - - -def wrap(x: int) -> _Unpicklable: - return _Unpicklable(x + 1) - - -def unwrap(o: _Unpicklable) -> int: - return o.value * 2 - - -def dup(x: int) -> Iterator[int]: - """A sync-generator op: yields two values per input (1->2 fan-out).""" - yield x * 2 - yield x * 2 + 1 - - -async def adup(x: int) -> AsyncIterator[int]: - """An async-generator op: yields two values per input (1->2 fan-out).""" - yield x - yield x + 1 - - -async def aadd_one(x: int) -> int: - """An async op: adds one. Picklable (module-level) so it can be fused into a worker.""" - return x + 1 - - -async def aunwrap(o: _Unpicklable) -> int: - """An async op consuming an unpicklable intermediate; runs inside the worker when fused.""" - return o.value * 2 - - -async def _astamp(item: _PidStamp) -> _PidStamp: - """An async op that records its pid; runs on the worker's event loop once fused.""" - item.pids["async"] = os.getpid() - return item - - -def _run(pipeline: Pipeline[Any], timeout: float = 60.0) -> list[Any]: - with pipeline.auto_stop(): - return list(pipeline.get_iterator(timeout=timeout)) - - -def stall_on_first(x: int) -> int: - """Stall only on the earliest item, so one worker holds it while its peer races ahead.""" - if x == 0: - time.sleep(3.0) - return x + 1 - - -class _Buffer1Queue(AsyncQueue): - """An ``AsyncQueue`` pinned to ``buffer_size=1`` (no prefetch slack).""" - - def __init__( - self, info: Any, *, buffer_size: int = 1, interval: float = -1 - ) -> None: - super().__init__(info, buffer_size=1) - - -def _install_small_buffers() -> None: - """Executor initializer: pin the worker sub-pipeline's queues to ``buffer_size=1``. - - A single-slot buffer stops the worker holding the earliest item from prefetching past it, so - that worker stays parked on the slow item while its peer drains the rest — the timing that - surfaces the startup race below (a cheap, deterministic stand-in for the recording overhead - that first exposed it under stress). - """ - set_default_queue_class(_Buffer1Queue) - - -class SubprocessPipelineFuseTest(unittest.TestCase): - def test_two_pool_stages_fused_match_unfused(self) -> None: - """A fused two-stage process-pool pipeline matches the unfused result.""" - n = 16 - ref = sorted((x + 1) * 2 for x in range(n)) - - ex = ProcessPoolExecutor(max_workers=2) - fused = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=3) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(fused)), ref) - - def test_unpicklable_intermediate_passes_through_fused(self) -> None: - """A fused run keeps the op->op handoff in-process, so an unpicklable value works.""" - n = 12 - ref = sorted((x + 1) * 2 for x in range(n)) - - ex = ProcessPoolExecutor(max_workers=2) - fused = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(wrap, executor=ex, concurrency=2) - .pipe(unwrap, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(fused)), ref) - - def test_generator_op_fused(self) -> None: - """A generator op is fusable, fanning out 1->N inside the worker sub-pipeline.""" - n = 8 - # dup yields {2x, 2x+1}; add_one shifts each by one, covering 1..2n exactly. - ref = list(range(1, 2 * n + 1)) - ex = ProcessPoolExecutor(max_workers=2) - fused = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(dup, executor=ex, concurrency=2) - .pipe(add_one, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(fused)), ref) - - def test_async_generator_op_composes_with_fused(self) -> None: - """An async-generator op (main-process) fans out downstream of a fused run.""" - n = 8 - # add_one+times_two fuse to (x+1)*2; adup then yields {v, v+1}, covering 2..2n+1. - ref = list(range(2, 2 * n + 2)) - ex = ProcessPoolExecutor(max_workers=2) - fused = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .pipe(adup) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(fused)), ref) - - def test_aggregate_between_pool_stages_not_absorbed(self) -> None: - """An aggregate between two pool stages stays in the main process (not absorbed). - - Its main-process batching semantics are unchanged: each batch holds exactly - ``aggregate``'s size. - """ - n = 12 - ex = ProcessPoolExecutor(max_workers=2) - fused = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .aggregate(3) - .pipe(len, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - # Each aggregated batch has exactly 3 items, so every produced value is 3. - out = _run(fused) - self.assertTrue(all(v == 3 for v in out)) - self.assertEqual(sum(out), n) - - def test_initializer_failure_surfaces_not_hangs(self) -> None: - """A worker initializer that raises surfaces as a failure instead of hanging. - - A real failure surfaces as :py:class:`PipelineFailure`; a hung pipeline would instead - raise :py:class:`TimeoutError` from the finite ``get_iterator`` timeout, so asserting on - ``PipelineFailure`` proves the failed initializer is reported rather than wedging the - collector forever. - """ - ex = ProcessPoolExecutor(max_workers=2, initializer=_raise_initializer) - fused = ( - PipelineBuilder() - .add_source(range(2)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(4) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with self.assertRaises(PipelineFailure): - _run(fused, timeout=30.0) - - def test_flag_off_is_unaffected(self) -> None: - """Without the flag, the same pipeline runs the stages normally and matches.""" - n = 10 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=False) - ) - self.assertEqual(sorted(_run(pipeline)), ref) - - -class AsyncFuseTest(unittest.TestCase): - """An async op tagged with a pool executor joins the surrounding fused run.""" - - def test_async_op_between_pool_stages_fuses_and_matches(self) -> None: - """An async op tagged with the neighbours' pool joins their run (one run, span 0..3).""" - n = 16 - ref = sorted((x + 1 + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - builder = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(aadd_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - ) - config = builder.get_config() - runs = _find_fusable_runs(config.pipes) - self.assertEqual(len(runs), 1) - self.assertEqual((runs[0].start, runs[0].stop), (0, 3)) - pipeline = build_pipeline(config, num_threads=4, fuse_subprocess_stages=True) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_async_op_runs_in_worker_process(self) -> None: - """A fused async op runs on the worker's event loop, in the worker process (not main). - - The sync pool stages on either side and the async op in the middle all stamp - ``os.getpid()``; the assertion proves the async stage shares the worker process with its - neighbours rather than hopping back to main. - """ - n = 16 - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source([_PidStamp(x) for x in range(n)]) - .pipe(_stamp_branch, executor=ex, concurrency=2) - .pipe(_astamp, executor=ex, concurrency=2) - .pipe(_stamp_merge, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - main_pid = os.getpid() - results = _run(pipeline) - self.assertEqual(len(results), n) - for item in results: - self.assertEqual(set(item.pids), {"branch", "async", "merge"}) - self.assertEqual(item.pids["branch"], item.pids["async"]) - self.assertEqual(item.pids["async"], item.pids["merge"]) - self.assertNotEqual(item.pids["async"], main_pid) - - def test_async_generator_op_fused(self) -> None: - """An async-generator op tagged with the pool fuses, fanning out inside the worker.""" - n = 8 - ex = ProcessPoolExecutor(max_workers=2) - # add_one -> v=x+1; adup(v) yields {v, v+1}. - ref = sorted(v for x in range(n) for v in (x + 1, x + 2)) - pipeline = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(adup, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_unpicklable_intermediate_through_async_fused(self) -> None: - """An unpicklable value handed from a sync pool op to a fused async op stays in-worker.""" - n = 12 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(wrap, executor=ex, concurrency=2) - .pipe(aunwrap, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_async_tag_ignored_when_fusion_off(self) -> None: - """With fusion off, the async op's pool tag is ignored and it runs in the main process.""" - n = 10 - ref = sorted((x + 1 + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(aadd_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=False) - ) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_thread_pool_executor_on_async_op_raises(self) -> None: - """A non-isolating (thread-pool) executor on an async op is rejected at config time.""" - with self.assertRaises(ValueError): - Pipe(aadd_one, executor=ThreadPoolExecutor(max_workers=1)) - # An isolating-pool executor is accepted (used only as a fusion-group tag). - Pipe(aadd_one, executor=ProcessPoolExecutor(max_workers=1)) - - def test_async_fused_multi_epoch(self) -> None: - """A fused async stage stays correct across epochs of a continuous source.""" - n = 12 - ref = sorted((x + 1 + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(aadd_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(3): # three epochs from the same warm worker pool - self.assertEqual(sorted(pipeline.get_iterator(timeout=60)), ref) - - def test_strip_async_executor_tags_clears_only_async(self) -> None: - """The strip pass clears an async op's pool tag but leaves a sync op's tag intact.""" - ex1 = ProcessPoolExecutor(max_workers=2) - ex2 = ProcessPoolExecutor(max_workers=2) - config = ( - PipelineBuilder() - .add_source(range(4)) - .pipe(add_one, executor=ex1) # sync pool: tag preserved - .pipe(aadd_one, executor=ex2) # async on a different pool: tag stripped - .add_sink(4) - .get_config() - ) - # Different executors -> neither adjacent stage fuses, so both tags stay on config.pipes. - self.assertEqual(_find_fusable_runs(config.pipes), []) - stripped = _strip_async_executor_tags(config) - sync_pipe, async_pipe = stripped.pipes[0], stripped.pipes[1] - assert isinstance(sync_pipe, PipeConfig) - assert isinstance(async_pipe, PipeConfig) - self.assertIsNotNone(sync_pipe._args.executor) # sync tag kept - self.assertIsNone(async_pipe._args.executor) # async tag cleared - - def test_nonfused_async_tag_runs_in_subprocess(self) -> None: - """A non-fused async op carrying a pool tag runs correctly via run_pipeline_in_subprocess. - - The strip pass removes the (meaningless) tag before executor hoisting, so no spurious - worker pool is spawned for an executor the async op never submits to. - """ - n = 12 - ref = sorted(((x + 1) + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - config = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one) # main-process sync op - .pipe(aadd_one, executor=ex) # lone async pool tag -> not fused - .pipe(times_two) # main-process sync op - .add_sink(n) - .get_config() - ) - self.assertEqual(_find_fusable_runs(config.pipes), []) - src = run_pipeline_in_subprocess( - config, num_threads=4, fuse_subprocess_stages=True - ) - self.assertEqual(sorted(src), ref) - - -class StartupRaceFuseTest(unittest.TestCase): - """A slow worker holding the earliest items must not have them silently dropped. - - Regression test for the silent earliest-item drop in the non-continuous fused - pipeline. With a single shared work-stealing input queue, a worker that drains its - stream and reaches ``_SESSION_END`` early could loop back and consume a second end - marker meant for a slower peer still holding (un-flushed) items; that peer then - never ended, the collector reached its ``_DONE`` count from the wrong workers and - finished, and the peer's earliest items were dropped with no error. Per-worker input - queues (one ``_SESSION_END`` each) make that impossible. ``stall_on_first`` + - ``buffer_size=1`` deterministically pin the worker that grabs item ``0`` while its - peer races through the rest — exactly the interleaving the race needs. - """ - - def test_slow_worker_does_not_drop_earliest_items(self) -> None: - """A worker stalled on the earliest item still delivers it instead of dropping it.""" - n = 16 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2, initializer=_install_small_buffers) - fused = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(stall_on_first, executor=ex, concurrency=1) - .pipe(times_two, executor=ex, concurrency=1) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(fused)), ref) - - -class ContinuousFuseTest(unittest.TestCase): - """Fusion with a continuous (multi-epoch) source.""" - - def test_multi_epoch_correct(self) -> None: - """A continuous fused pipeline yields the correct set each epoch.""" - n = 12 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=3) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(3): # three epochs from the same warm worker pool - epoch = sorted(pipeline.get_iterator(timeout=60)) - self.assertEqual(epoch, ref) - - def test_unpicklable_intermediate_multi_epoch(self) -> None: - """The unpicklable op->op handoff keeps working across epochs.""" - n = 10 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(wrap, executor=ex, concurrency=2) - .pipe(unwrap, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(2): - self.assertEqual(sorted(pipeline.get_iterator(timeout=60)), ref) - - def test_aggregate_not_absorbed_multi_epoch(self) -> None: - """A non-absorbed aggregate keeps single-flush-per-epoch semantics across epochs. - - Because the aggregate runs in the main process (not per worker), each epoch produces - full-size batches plus one combined partial, and every item is accounted for. - """ - n = 12 - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(add_one, executor=ex, concurrency=2) - .aggregate(3) - .pipe(len, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(2): # each epoch's batches cover exactly n items - out = list(pipeline.get_iterator(timeout=60)) - self.assertEqual(sum(out), n) - - def test_fewer_items_than_workers(self) -> None: - """An epoch with fewer items than workers still completes (some workers run empty).""" - n = 2 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=4) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(2): - self.assertEqual(sorted(pipeline.get_iterator(timeout=60)), ref) - - def test_generator_op_multi_epoch(self) -> None: - """A fused generator op keeps its 1->N fan-out correct across epochs.""" - n = 8 - ref = list(range(1, 2 * n + 1)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(dup, executor=ex, concurrency=2) - .pipe(add_one, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(3): - self.assertEqual(sorted(pipeline.get_iterator(timeout=60)), ref) - - def test_async_generator_op_multi_epoch(self) -> None: - """A main-process async-generator op fans out downstream of a fused run each epoch.""" - n = 8 - ref = list(range(2, 2 * n + 2)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .pipe(adup) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(3): - self.assertEqual(sorted(pipeline.get_iterator(timeout=60)), ref) - - def test_continuous_op_failure_does_not_deadlock(self) -> None: - """Op failures (dropped per SPDL default) still let each epoch's barrier complete. - - ``boom`` raises on every item, so the fused workers produce no results; the test checks - the continuous epoch barrier still completes each epoch (workers report the boundary - with zero results) instead of deadlocking, matching unfused drop-on-failure behavior. - """ - n = 8 - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(boom, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(2): - self.assertEqual(list(pipeline.get_iterator(timeout=60)), []) - - def test_continuous_in_subprocess(self) -> None: - """Continuous fusion composes with run_pipeline_in_subprocess across epochs.""" - n = 12 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - config = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - .get_config() - ) - src = run_pipeline_in_subprocess( - config, num_threads=4, fuse_subprocess_stages=True - ) - for _ in range(3): # one epoch per iteration - self.assertEqual(sorted(src), ref) - - def test_teardown_mid_stream_does_not_hang(self) -> None: - """Tearing a continuous fused subprocess pipeline down mid-stream must not hang. - - A teardown before the stream is drained leaves the per-worker input queue full; the - pool shutdown must cancel the queue feeder-thread join instead of blocking forever - waiting to flush buffered items into a pipe the (terminated) workers no longer drain. - Uses the same path data loaders do — ``run_pipeline_in_subprocess`` + fusion — and - triggers teardown via the iterable's documented ``_finalizer`` handle. - """ - n = 100_000 # far more than any buffer, so the stream is still full at teardown - ex = ProcessPoolExecutor(max_workers=2) - self.addCleanup(ex.shutdown) - config = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(2) - .get_config() - ) - src = run_pipeline_in_subprocess( - config, num_threads=4, fuse_subprocess_stages=True - ) - it = iter(src) - next(it) # consume a couple of items so the queues stay backed up - next(it) - torn_down = threading.Event() - - def _teardown() -> None: - src._finalizer() # pyre-ignore[16]: documented teardown handle - torn_down.set() - - t = threading.Thread(target=_teardown, daemon=True) - t.start() - self.assertTrue( - torn_down.wait(timeout=60), - "fused-pool teardown hung on a full input queue after a mid-stream stop", - ) - t.join(timeout=10) - - -class PathVariantsFuseTest(unittest.TestCase): - """Fusion of a path-variants stage whose branches share one pool executor.""" - - def test_lone_path_variants_pool_branches_fused_match_unfused(self) -> None: - """A single path-variants stage with same-pool branches fuses and matches unfused. - - Even items take the ``add_one`` branch, odd items the ``times_two`` branch. The whole - routing construct moves into one worker, so the result must match the unfused run. - """ - n = 16 - ref = sorted(x + 1 if x % 2 == 0 else x * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - builder = ( - PipelineBuilder() - .add_source(range(n)) - .path_variants( - route_by_parity, - [[Pipe(add_one, executor=ex)], [Pipe(times_two, executor=ex)]], - ) - .add_sink(n) - ) - config = builder.get_config() - # A lone same-pool path-variants is a fusable run on its own. - self.assertEqual(len(_find_fusable_runs(config.pipes)), 1) - pipeline = build_pipeline(config, num_threads=4, fuse_subprocess_stages=True) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_path_variants_async_branch_stage_fuses(self) -> None: - """A branch mixing an async (no-executor) stage with a pool stage still fuses. - - Mirrors a cache miss-path shape (async I/O then pool CPU): the async stage runs on the - worker's loop and the pool stage on its threads, all inside one worker. - """ - n = 16 - ref = sorted(x + 1 if x % 2 == 0 else x * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - builder = ( - PipelineBuilder() - .add_source(range(n)) - .path_variants( - route_by_parity, - [ - [Pipe(add_one, executor=ex)], - [Pipe(aident), Pipe(times_two, executor=ex)], - ], - ) - .add_sink(n) - ) - config = builder.get_config() - self.assertEqual(len(_find_fusable_runs(config.pipes)), 1) - pipeline = build_pipeline(config, num_threads=4, fuse_subprocess_stages=True) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_path_variants_ordered_async_tag_does_not_block_fusion(self) -> None: - """An input-ordered async op carrying a pool tag doesn't poison path-variants fusion. - - The async op runs on the worker's loop, not the pool, so its ``output_order="input"`` - can't be broken by pool parallelism and must not mark the same-pool stage non-fusable — - unlike an input-ordered *sync* pool-pipe, which genuinely blocks fusion. - """ - n = 16 - ref = sorted(x + 1 if x % 2 == 0 else (x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - builder = ( - PipelineBuilder() - .add_source(range(n)) - .path_variants( - route_by_parity, - [ - [Pipe(add_one, executor=ex)], - [ - Pipe(aadd_one, executor=ex, output_order="input"), - Pipe(times_two, executor=ex), - ], - ], - ) - .add_sink(n) - ) - config = builder.get_config() - # The ordered async op is only a fusion-group tag, so the same-pool stage still fuses. - self.assertEqual(len(_find_fusable_runs(config.pipes)), 1) - pipeline = build_pipeline(config, num_threads=4, fuse_subprocess_stages=True) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_path_variants_mixed_executors_not_fused(self) -> None: - """Branches on two different pool executors are not fused, and still run correctly.""" - n = 16 - ref = sorted(x + 1 if x % 2 == 0 else x * 2 for x in range(n)) - ex1 = ProcessPoolExecutor(max_workers=2) - ex2 = ProcessPoolExecutor(max_workers=2) - builder = ( - PipelineBuilder() - .add_source(range(n)) - .path_variants( - route_by_parity, - [[Pipe(add_one, executor=ex1)], [Pipe(times_two, executor=ex2)]], - ) - .add_sink(n) - ) - config = builder.get_config() - self.assertEqual(_find_fusable_runs(config.pipes), []) - pipeline = build_pipeline(config, num_threads=4, fuse_subprocess_stages=True) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_path_variants_fused_with_adjacent_pool_pipe(self) -> None: - """A pool-pipe adjacent to a same-pool path-variants stage fuses into one run.""" - n = 16 - # add_one shifts each item to v=x+1; then even v -> add_one (v+1), odd v -> times_two (v*2). - ref = sorted( - (v + 1) if v % 2 == 0 else (v * 2) for v in (x + 1 for x in range(n)) - ) - ex = ProcessPoolExecutor(max_workers=2) - builder = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex) - .path_variants( - route_by_parity, - [[Pipe(add_one, executor=ex)], [Pipe(times_two, executor=ex)]], - ) - .add_sink(n) - ) - config = builder.get_config() - runs = _find_fusable_runs(config.pipes) - self.assertEqual(len(runs), 1) - self.assertEqual( - (runs[0].start, runs[0].stop), (0, 2) - ) # both stages in one run - pipeline = build_pipeline(config, num_threads=4, fuse_subprocess_stages=True) - self.assertEqual(sorted(_run(pipeline)), ref) - - def test_path_variants_fused_multi_epoch(self) -> None: - """A fused path-variants stage stays correct across epochs of a continuous source.""" - n = 12 - ref = sorted(x + 1 if x % 2 == 0 else x * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source(range(n), continuous=True) - .path_variants( - route_by_parity, - [[Pipe(add_one, executor=ex)], [Pipe(times_two, executor=ex)]], - ) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - with pipeline.auto_stop(): - for _ in range(3): # three epochs from the same warm worker pool - self.assertEqual(sorted(pipeline.get_iterator(timeout=60)), ref) - - def _assert_one_worker_process(self, item: _PidStamp, main_pid: int) -> None: - """The router, branch, and post-merge stages all ran in one worker process (not main).""" - self.assertEqual(set(item.pids), {"router", "branch", "merge"}) - # router -> branch -> merge stayed in the SAME process for this item ... - self.assertEqual(item.pids["router"], item.pids["branch"]) - self.assertEqual(item.pids["branch"], item.pids["merge"]) - # ... and that process is a worker, not the main process. - self.assertNotEqual(item.pids["router"], main_pid) - - def test_path_variants_router_branches_merge_share_one_process(self) -> None: - """Router, branches, and the fan-in (post-merge) stage all run in one worker process. - - Each stage stamps ``os.getpid()`` into its item, so the assertion proves the whole - router -> branch -> merge -> post-merge chain runs inside a single worker process per - item (distinct from the main process): fusion leaves no per-stage process hop. A - regression that stopped fusing path-variants would run them in the main process and - fail the ``assertNotEqual(main_pid)`` check. - """ - n = 16 - ex = ProcessPoolExecutor(max_workers=2) - builder = ( - PipelineBuilder() - .add_source([_PidStamp(x) for x in range(n)]) - .path_variants( - _stamp_router, - [ - [Pipe(_stamp_branch, executor=ex)], - [Pipe(_stamp_branch, executor=ex)], - ], - ) - .pipe(_stamp_merge, executor=ex) # fan-in / post-merge, same pool - .add_sink(n) - ) - config = builder.get_config() - # path-variants + the post-merge pool-pipe fuse into a single run. - self.assertEqual(len(_find_fusable_runs(config.pipes)), 1) - pipeline = build_pipeline(config, num_threads=4, fuse_subprocess_stages=True) - results = _run(pipeline) - - self.assertEqual(len(results), n) - main_pid = os.getpid() - worker_pids = set() - for item in results: - self._assert_one_worker_process(item, main_pid) - worker_pids.add(item.pids["router"]) - # work really ran in worker process(es), never the main process. - self.assertTrue(worker_pids) - self.assertNotIn(main_pid, worker_pids) - - def test_path_variants_process_locality_multi_epoch(self) -> None: - """Process-locality holds for a continuous (multi-epoch) fused path-variants too. - - This is the shape ``op_mode=mp`` uses. Pickling isolates each epoch's stamps (the - source objects are copied into the worker), so every epoch's output items carry a fresh - same-process router/branch/merge stamp. - """ - n = 12 - ex = ProcessPoolExecutor(max_workers=2) - pipeline = ( - PipelineBuilder() - .add_source([_PidStamp(x) for x in range(n)], continuous=True) - .path_variants( - _stamp_router, - [ - [Pipe(_stamp_branch, executor=ex)], - [Pipe(_stamp_branch, executor=ex)], - ], - ) - .pipe(_stamp_merge, executor=ex) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - main_pid = os.getpid() - with pipeline.auto_stop(): - for _ in range(2): # two epochs from the same warm worker pool - items = list(pipeline.get_iterator(timeout=60)) - self.assertEqual(len(items), n) - for item in items: - self._assert_one_worker_process(item, main_pid) - - -class FuseInSubprocessTest(unittest.TestCase): - """Fusion composed with whole-pipeline subprocess execution.""" - - def test_fused_run_in_subprocess_matches(self) -> None: - """A fused config run via run_pipeline_in_subprocess yields the same result.""" - n = 16 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - config = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=3) - .add_sink(n) - .get_config() - ) - src = run_pipeline_in_subprocess( - config, num_threads=4, fuse_subprocess_stages=True - ) - self.assertEqual(sorted(src), ref) - - def test_fused_unpicklable_intermediate_in_subprocess(self) -> None: - """The fused handle survives the trip into the pipeline subprocess and the - unpicklable op->op handoff stays inside a worker.""" - n = 12 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = ProcessPoolExecutor(max_workers=2) - config = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(wrap, executor=ex, concurrency=2) - .pipe(unwrap, executor=ex, concurrency=2) - .add_sink(n) - .get_config() - ) - src = run_pipeline_in_subprocess( - config, num_threads=4, fuse_subprocess_stages=True - ) - self.assertEqual(sorted(src), ref) - - -class FeedAbortTest(unittest.TestCase): - """The bridge feeder must wind down promptly when the collector signals abort.""" - - def test_feed_ends_session_when_aborted_while_idle(self) -> None: - """A feeder parked on an empty input queue still emits the per-worker _SESSION_END. - - On a worker error the collector sets ``abort`` while the feeder is typically - blocked waiting on a slow/idle upstream. The feeder must wake and send exactly one - ``_SESSION_END`` onto each worker's own queue so the collector can drain every - ``_DONE`` instead of hanging until its stall timeout. Driving ``_feed`` directly - keeps the abort-while-idle race deterministic. - """ - num_workers = 3 - - async def _scenario() -> list[list[Any]]: - in_qs: list[_queue.Queue[Any]] = [ - _queue.Queue() for _ in range(num_workers) - ] - input_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="input") - ) # stays empty -> get() blocks - abort = asyncio.Event() - feeder_idle = asyncio.Event() - put_stop = threading.Event() - with ThreadPoolExecutor(max_workers=num_workers + 1) as ex: - task = asyncio.ensure_future( - _subprocess_pipe._feed( - input_queue, in_qs, ex, abort, feeder_idle, put_stop - ) - ) - await asyncio.sleep(0.1) # let the feeder park on input_queue.get() - self.assertFalse(task.done(), "feeder should be parked on empty queue") - abort.set() - await asyncio.wait_for(task, timeout=5.0) - return [[q.get_nowait() for _ in range(q.qsize())] for q in in_qs] - - msgs = asyncio.run(_scenario()) - # Every worker's own queue receives exactly one _SESSION_END. - self.assertEqual(msgs, [[(_subprocess_pipe._SESSION_END, None)]] * num_workers) - - -class StallGuardTest(unittest.TestCase): - """The collector's stall guard against an abruptly-dead worker.""" - - def test_check_stall_raises_past_timeout(self) -> None: - """``_check_stall`` raises once no message has arrived for longer than the bound.""" - orig = _subprocess_pipe._WORKER_STALL_TIMEOUT - _subprocess_pipe._WORKER_STALL_TIMEOUT = 0.0 - try: - with self.assertRaises(TimeoutError): - _subprocess_pipe._check_stall(time.monotonic() - 1.0) - finally: - _subprocess_pipe._WORKER_STALL_TIMEOUT = orig - - def test_check_stall_quiet_within_timeout(self) -> None: - """``_check_stall`` does not raise while progress is within the bound.""" - orig = _subprocess_pipe._WORKER_STALL_TIMEOUT - _subprocess_pipe._WORKER_STALL_TIMEOUT = 60.0 - try: - _subprocess_pipe._check_stall(time.monotonic()) # should not raise - finally: - _subprocess_pipe._WORKER_STALL_TIMEOUT = orig - - def test_collect_suppresses_stall_while_feeder_idle(self) -> None: - """An idle feeder suppresses the collector's stall guard during input starvation. - - With the timeout pinned to zero, any stall check on an empty queue would trip instantly; - the collector must instead keep draining while ``feeder_idle`` is set (nothing dispatched, - no worker message due) and still finish once the worker reports ``_DONE``. - """ - orig = _subprocess_pipe._WORKER_STALL_TIMEOUT - _subprocess_pipe._WORKER_STALL_TIMEOUT = 0.0 - - async def _scenario() -> None: - out_q: _queue.Queue[Any] = _queue.Queue() - output_queue = AsyncQueue( - StageInfo(pipeline_id=0, stage_id="0", stage_name="output") - ) - abort = asyncio.Event() - feeder_idle = asyncio.Event() - feeder_idle.set() # feeder parked on an idle upstream -> no message expected - with ThreadPoolExecutor(max_workers=2) as ex: - task = asyncio.ensure_future( - _subprocess_pipe._collect( - out_q, 1, output_queue, ex, abort, feeder_idle - ) - ) - await asyncio.sleep( - 0.6 - ) # several empty poll cycles; must not trip the guard - self.assertFalse( - task.done(), "idle feeder must suppress the stall guard" - ) - out_q.put((_subprocess_pipe._DONE, None)) - await asyncio.wait_for(task, timeout=5.0) - - try: - asyncio.run(_scenario()) - finally: - _subprocess_pipe._WORKER_STALL_TIMEOUT = orig - - -# Set inside each fused worker process by ``_install_recording_hooks`` (the executor -# initializer). The worker is a separate process, so the recording hooks below cannot share -# in-memory counters with the test; each instead writes a small JSON file into this directory, -# which the test reads back to prove the hooks fired inside the worker. -_EVIDENCE_DIR: str | None = None - -# A per-process monotonic counter for evidence filenames. ``id(self)`` is only unique among -# simultaneously-live objects -- CPython recycles addresses, so two recorders created and -# destroyed in sequence could collide on one path and clobber each other's count. This counter -# (combined with ``pid``) is guaranteed distinct for every write within a process. -_EVIDENCE_SEQ = itertools.count() -_EVIDENCE_SEQ_LOCK = threading.Lock() - - -def _write_evidence(record: dict[str, Any]) -> None: - if (d := _EVIDENCE_DIR) is None: - return - with _EVIDENCE_SEQ_LOCK: - seq = next(_EVIDENCE_SEQ) - # ``pid`` proves the write came from a worker process (never the main/test process); ``seq`` - # keeps each writer's file distinct within that process, so no two writers ever race on one - # path. Write to a temp file then atomically rename, so a reader polling mid-write never sees - # a truncated/partial file. - path = os.path.join(d, f"{record['kind']}-{record['pid']}-{seq}.json") - tmp = f"{path}.tmp" - try: - with open(tmp, "w") as f: - f.write(json.dumps(record)) - os.replace(tmp, path) - except OSError: - # This runs inside the worker's stage_hook ``finally`` (stage teardown). It must never - # raise into the stage: an I/O failure here -- e.g. the evidence ``TemporaryDirectory`` - # already torn down in a teardown race -- would otherwise crash the stage and drop items - # from the data path the test is validating. A lost evidence file at most makes an - # assertion under-count, which surfaces as a clear assertion failure (not a baffling - # dropped-item one). - pass - - -class _RecordingTaskHook(TaskHook): - """A ``TaskHook`` that records how many tasks ran in its (subprocess) stage.""" - - def __init__(self, info: Any, interval: float = -1) -> None: - self.info = info - self.n_tasks = 0 - - @asynccontextmanager - async def stage_hook(self) -> AsyncIterator[None]: - try: - yield - finally: - _write_evidence( - { - "kind": "task", - "name": str(self.info), - "pid": os.getpid(), - "iid": id(self), - "n_tasks": self.n_tasks, - } - ) - - @asynccontextmanager - async def task_hook(self, input_item: Any = None) -> AsyncIterator[None]: - self.n_tasks += 1 - yield - - -class _RecordingQueue(AsyncQueue): - """An ``AsyncQueue`` whose ``stage_hook`` records that it ran in the (subprocess) stage.""" - - def __init__( - self, info: Any, *, buffer_size: int = 1, interval: float = -1 - ) -> None: - super().__init__(info, buffer_size=buffer_size) - self.n_get = 0 - - async def get(self) -> object: - item = await super().get() - self.n_get += 1 - return item - - @asynccontextmanager - async def stage_hook(self) -> AsyncIterator[None]: - try: - yield - finally: - _write_evidence( - { - "kind": "queue", - "name": str(self.info), - "pid": os.getpid(), - "iid": id(self), - "n_get": self.n_get, - } - ) - - -def _install_recording_hooks(evidence_dir: str) -> None: - """Executor initializer: runs inside each fused worker before its sub-pipeline is built. - - Fusion reads this off the ``ProcessPoolExecutor`` (``_pool_params``) and runs it in every - worker process, so the worker's nested ``build_pipeline`` — which is given no explicit - ``task_hook_factory``/``queue_class`` — picks up these recording classes as its defaults. - """ - global _EVIDENCE_DIR - _EVIDENCE_DIR = evidence_dir - set_default_hook_class(_RecordingTaskHook) - set_default_queue_class(_RecordingQueue) - - -class FuseHookTest(unittest.TestCase): - """``TaskHook`` and the queue ``stage_hook`` fire inside the fused worker subprocess. - - The fused run executes as a nested pipeline inside main-process-owned worker processes; the - per-stage hooks/stats fire there, not in the bridge stage. These tests install recording - hook/queue classes as the worker defaults (via the pool's ``initializer``) and assert, from - the evidence files those recorders leave behind, that both fired in a non-main process. - """ - - def _read_evidence(self, evidence_dir: str) -> list[dict[str, Any]]: - records: list[dict[str, Any]] = [] - for name in os.listdir(evidence_dir): - if not name.endswith(".json"): - continue # skip ``.json.tmp`` files still being written - with open(os.path.join(evidence_dir, name)) as f: - records.append(json.loads(f.read())) - return records - - def _await_evidence( - self, evidence_dir: str, n: int, timeout: float = 60.0 - ) -> list[dict[str, Any]]: - """Re-read evidence until the workers' task total lands, then return the records. - - Each fused worker writes its evidence files to disk during nested-pipeline teardown, a - side channel with no happens-before relationship to the main iterator completing. Polling - until the recorded task total reaches the expected ``2 * n`` closes that read-too-early - window without masking regressions: on a genuine failure (hooks never fire) the poll times - out and the caller's assertions report the shortfall. - """ - deadline = time.monotonic() + timeout - records = self._read_evidence(evidence_dir) - while sum(r["n_tasks"] for r in records if r["kind"] == "task") < 2 * n: - if time.monotonic() > deadline: - break # fall through; the assertions below report the shortfall - time.sleep(0.02) - records = self._read_evidence(evidence_dir) - return records - - def _assert_hooks_fired(self, records: list[dict[str, Any]], n: int) -> None: - task_records = [r for r in records if r["kind"] == "task"] - queue_records = [r for r in records if r["kind"] == "queue"] - - self.assertTrue(task_records, "TaskHook never fired in any fused worker") - self.assertTrue( - queue_records, "queue stage_hook never fired in any fused worker" - ) - - # Every record came from a worker process, never the main/test process — this is what - # proves the hooks fired "in the subprocess". - main_pid = os.getpid() - for r in records: - self.assertNotEqual(r["pid"], main_pid) - - # The two fused pipe stages each see every item once: ``add_one`` over n items, then - # ``times_two`` over n items => 2n ``task_hook`` invocations, summed across workers. - self.assertEqual(sum(r["n_tasks"] for r in task_records), 2 * n) - - # The fused sub-pipeline's queues carried the data inside the worker(s). - self.assertGreaterEqual(sum(r["n_get"] for r in queue_records), n) - - def test_hooks_fire_in_fused_workers(self) -> None: - """Hooks fire in the fused workers spawned by ``PipelineBuilder.build``.""" - n = 16 - ref = sorted((x + 1) * 2 for x in range(n)) - with tempfile.TemporaryDirectory() as evidence_dir: - ex = ProcessPoolExecutor( - max_workers=2, - initializer=_install_recording_hooks, - initargs=(evidence_dir,), - ) - fused = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(fused)), ref) - self._assert_hooks_fired(self._await_evidence(evidence_dir, n), n) - - def test_hooks_fire_in_fused_workers_via_subprocess(self) -> None: - """Hooks fire in the fused workers when the run is driven from a pipeline subprocess.""" - n = 16 - ref = sorted((x + 1) * 2 for x in range(n)) - with tempfile.TemporaryDirectory() as evidence_dir: - ex = ProcessPoolExecutor( - max_workers=2, - initializer=_install_recording_hooks, - initargs=(evidence_dir,), - ) - config = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - .get_config() - ) - src = run_pipeline_in_subprocess( - config, num_threads=4, fuse_subprocess_stages=True - ) - # Each worker writes its evidence during nested-pipeline teardown, on a side channel - # (the filesystem) with no happens-before relationship to iteration ending — so - # ``_await_evidence`` polls for the files rather than assuming they are already there. - self.assertEqual(sorted(src), ref) - self._assert_hooks_fired(self._await_evidence(evidence_dir, n), n) - - -def _has_interpreter_pool_executor() -> bool: - if sys.version_info < (3, 14): - return False - try: - from concurrent.futures.interpreter import InterpreterPoolExecutor # noqa: F401 - except ImportError: - return False - return True - - -_HAS_INTERPRETER: bool = _has_interpreter_pool_executor() - - -if _HAS_INTERPRETER: - - class SubprocessPipelineInterpreterPoolFuseTest(unittest.TestCase): - def test_interpreter_pool_stages_fused(self) -> None: - """Stages sharing an InterpreterPoolExecutor are recognized and fused.""" - from concurrent.futures import InterpreterPoolExecutor # pyre-ignore[21] - - n = 12 - ref = sorted((x + 1) * 2 for x in range(n)) - ex = InterpreterPoolExecutor(max_workers=2) - fused = ( - PipelineBuilder() - .add_source(range(n)) - .pipe(add_one, executor=ex, concurrency=2) - .pipe(times_two, executor=ex, concurrency=2) - .add_sink(n) - .build(num_threads=4, fuse_subprocess_stages=True) - ) - self.assertEqual(sorted(_run(fused)), ref) diff --git a/tools/skills/authoring/building_pipelines.md b/tools/skills/authoring/building_pipelines.md index b2f17c4d5..ed1e41f1f 100644 --- a/tools/skills/authoring/building_pipelines.md +++ b/tools/skills/authoring/building_pipelines.md @@ -142,24 +142,24 @@ def _build_mt_pipeline(...): ) ``` -**Extract a shared builder only when the variants differ purely by argument.** When two modes really do produce the identical chain except for one value (say an executor or a buffer size), factor out *the entire chain* into a helper parameterized by that value — never a helper that builds only part of it: +**Extract a shared builder only when the variants differ purely by argument.** When two modes really do produce the identical chain except for one value (say a region target or a buffer size), factor out *the entire chain* into a helper parameterized by that value — never a helper that builds only part of it: ```python -# Two modes that are the SAME stage graph, differing only by the executor: -def _build_thread_pipeline(args): - return _build_local_pipeline(args, ThreadPoolExecutor(max_workers=args.n)) +def _build_subprocess_pipeline(args): + return _build_local_pipeline(args, ProcessPoolExecutorConfig(max_workers=args.n)) +def _build_subinterp_pipeline(args): + return _build_local_pipeline(args, InterpreterPoolExecutorConfig(max_workers=args.n)) -def _build_interp_pipeline(args): - return _build_local_pipeline(args, InterpreterPoolExecutor(max_workers=args.n)) - -def _build_local_pipeline(args, executor: Executor): +def _build_local_pipeline(args, target: ProcessPoolExecutorConfig | InterpreterPoolExecutorConfig): return ( PipelineBuilder() .add_source(...) .aggregate(...) - .pipe(decode, executor=executor) + .to(target) + .pipe(decode) + .to(MAIN_PROCESS) .add_sink(...) - .build(num_threads=args.n, fuse_subprocess_stages=True) + .build(num_threads=args.n) ) ```