diff --git a/src/spdl/pipeline/_builder.py b/src/spdl/pipeline/_builder.py index e500cf32e..a9c2b425c 100644 --- a/src/spdl/pipeline/_builder.py +++ b/src/spdl/pipeline/_builder.py @@ -7,6 +7,7 @@ # pyre-strict import logging +import sys from collections.abc import AsyncIterable, Callable, Iterable, Sequence from concurrent.futures import Executor from fractions import Fraction @@ -15,17 +16,22 @@ from spdl._internal import log_api_usage_once from spdl.pipeline._components import AsyncQueue, StageInfo, TaskHook from spdl.pipeline.defs import ( + _MainProcess, + _PipeType, _TPipeInputs, Aggregate, AggregateConfig, Aggregator, Disaggregate, DisaggregateConfig, + InterpreterPoolExecutorConfig, PathVariants, PathVariantsConfig, Pipe, PipeConfig, PipelineConfig, + PlacementConfig, + ProcessPoolExecutorConfig, SinkConfig, SourceConfig, ) @@ -45,6 +51,79 @@ U_ = TypeVar("U_") +def _has_ordered_pipe(cfg: object) -> bool: + """Whether ``cfg`` is (or, for path-variants, contains) an ``output_order="input"`` pipe. + + Recurses into :py:class:`~spdl.pipeline.defs.PathVariantsConfig` branches: those stages run + inside the region worker too, so an input-ordered pipe nested in a branch is just as invalid + as a top-level one (global input order cannot be preserved across the pool's workers). + """ + if isinstance(cfg, PipeConfig): + return cfg._type is _PipeType.OrderedPipe + if isinstance(cfg, PathVariantsConfig): + return any(_has_ordered_pipe(s) for path in cfg.paths for s in path) + return False + + +def _validate_executor_regions( + pipes: Sequence[ + PipeConfig + | AggregateConfig + | DisaggregateConfig + | PathVariantsConfig + | PlacementConfig + ], +) -> None: + """Validate the :py:meth:`PipelineBuilder.to` region markers in ``pipes``. + + Stateless scan (a pipeline starts on the main process). Rejects: a subinterpreter region on + Python < 3.14; a stage using ``output_order="input"`` inside a region (including one nested + in a ``path_variants`` branch, since order cannot be preserved across independent workers); an + empty region (a ``.to(...)`` marker with no stages before the next marker/sink); and a region + left open at the sink. + """ + in_region = False + region_has_stage = False + for p in pipes: + if isinstance(p, PlacementConfig): + # A non-main region that opened but never received a stage does nothing; reject it + # rather than silently dropping it (usually a stray or duplicated ``.to(...)``). + # Adjacent *non-empty* regions with different targets remain valid. + if in_region and not region_has_stage: + raise ValueError( + "An empty `to(...)` region has no stages. Add stages before the next " + "`.to(...)`, or remove the redundant marker." + ) + target = p.target + in_region = not isinstance(target, _MainProcess) + region_has_stage = False + if isinstance( + target, InterpreterPoolExecutorConfig + ) and sys.version_info < (3, 14): + raise RuntimeError( + "A subinterpreter region (`to(InterpreterPoolExecutorConfig(...))`) requires " + "Python 3.14 or later. Current version: " + f"{sys.version_info.major}.{sys.version_info.minor}" + ) + elif in_region: + region_has_stage = True + if _has_ordered_pipe(p): + raise ValueError( + "A stage with `output_order='input'` cannot run inside a `to()` region: " + "input order cannot be preserved across the region's independent workers." + ) + if in_region and not region_has_stage: + raise ValueError( + "An empty `to(...)` region has no stages. Add stages before closing it, or remove " + "the redundant marker." + ) + if in_region: + raise ValueError( + "A `to()` execution region must be closed with `to(MAIN_PROCESS)` before " + "`add_sink()`. The sink always runs in the main process." + ) + + ################################################################################ # Builder ################################################################################ @@ -81,7 +160,11 @@ def __init__(self) -> None: self._src: SourceConfig[T] | None = None self._process_args: list[ - PipeConfig | AggregateConfig | DisaggregateConfig | PathVariantsConfig + PipeConfig + | AggregateConfig + | DisaggregateConfig + | PathVariantsConfig + | PlacementConfig ] = [] self._sink: SinkConfig[U] | None = None @@ -246,6 +329,66 @@ def disaggregate(self) -> "PipelineBuilder[T, U]": self._process_args.append(Disaggregate()) return self + def to( + self, + target: "ProcessPoolExecutorConfig | InterpreterPoolExecutorConfig | _MainProcess", + /, + ) -> "PipelineBuilder[T, U]": + """**[Experimental]** Designate where the subsequent stages execute. + + Opens (or closes) an *execution region*: every stage added after this call runs on + ``target`` until the next :py:meth:`to`. A pipeline starts on the main process, so a + region is opened by ``to(ProcessPoolExecutorConfig(...))`` or ``to(InterpreterPoolExecutorConfig(...))`` + and closed by ``to(MAIN_PROCESS)``. The stages inside a region are fused into one nested + pipeline that runs together in a worker process (or subinterpreter), so the value handed + from one stage to the next stays in the worker — it is **not** copied back to the main + process between stages and need not be picklable. Only the region's inputs and outputs + cross the boundary. + + Unlike passing ``executor=`` to individual :py:meth:`pipe` calls, a region also carries + :py:meth:`aggregate`, :py:meth:`disaggregate`, and :py:meth:`path_variants` stages into + the worker, and gives the worker-pool configuration a single home. + + Args: + target: Where the following stages run. + + - :py:class:`~spdl.pipeline.defs.ProcessPoolExecutorConfig` — a pool of worker processes. + - :py:class:`~spdl.pipeline.defs.InterpreterPoolExecutorConfig` — a pool of subinterpreters + (Python 3.14+; the region's ops must avoid NumPy/PyTorch, which cannot be + imported in a subinterpreter). + - :py:data:`~spdl.pipeline.defs.MAIN_PROCESS` — close the current region; the + following stages run in the main process. + + A live :py:class:`~concurrent.futures.Executor` is **not** accepted — pass a spec + so the pipeline stays expressible as static config. To run a single stage on a + custom executor, use ``pipe(executor=...)`` instead. + + .. note:: + + The region must be closed with ``to(MAIN_PROCESS)`` before :py:meth:`add_sink`, a + stage inside a region may not use ``output_order="input"`` (order cannot be preserved + across independent workers), and a subinterpreter region requires Python 3.14+. These + are checked when the pipeline is built. + + .. versionadded:: 0.6.0 + """ + if isinstance(target, Executor): + raise TypeError( + "`to()` takes a serializable execution target (ProcessPoolExecutorConfig, " + "InterpreterPoolExecutorConfig, or MAIN_PROCESS), not a live Executor. To run a single " + "stage on a custom executor, pass it to `pipe(executor=...)` instead." + ) + if not isinstance( + target, + (ProcessPoolExecutorConfig, InterpreterPoolExecutorConfig, _MainProcess), + ): + raise TypeError( + "`to()` target must be a ProcessPoolExecutorConfig, InterpreterPoolExecutorConfig, or " + f"MAIN_PROCESS. Got: {type(target).__name__}." + ) + self._process_args.append(PlacementConfig(target=target)) + return self + def path_variants( self, router: Callable, @@ -286,7 +429,11 @@ def get_config(self) -> PipelineConfig[U]: A PipelineConfig object representing the current pipeline configuration. Raises: - RuntimeError: If source or sink is not set. + RuntimeError: If source or sink is not set, or a subinterpreter region is used on + Python < 3.14. + ValueError: If an execution region opened by :py:meth:`to` is not closed with + ``to(MAIN_PROCESS)`` before the sink, or a stage inside a region uses + ``output_order="input"``. """ if (src := self._src) is None: raise RuntimeError("Source is not set. Did you call `add_source`?") @@ -294,6 +441,8 @@ def get_config(self) -> PipelineConfig[U]: if (sink := self._sink) is None: raise RuntimeError("Sink is not set. Did you call `add_sink`?") + _validate_executor_regions(self._process_args) + return PipelineConfig(src, self._process_args, sink) def build( diff --git a/tests/pipeline/builder_to_test.py b/tests/pipeline/builder_to_test.py new file mode 100644 index 000000000..7f543f86b --- /dev/null +++ b/tests/pipeline/builder_to_test.py @@ -0,0 +1,196 @@ +# 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 sys +import unittest +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +from spdl.pipeline import build_pipeline, Pipeline, PipelineBuilder +from spdl.pipeline.defs import ( + _MainProcess, + InterpreterPoolExecutorConfig, + MAIN_PROCESS, + Pipe, + PlacementConfig, + 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)) + + +class ToMethodTest(unittest.TestCase): + """The PipelineBuilder.to() region-marker method.""" + + def test_appends_executor_config(self) -> None: + """to() appends an PlacementConfig marker carrying the given target.""" + config = ( + PipelineBuilder() + .add_source(range(4)) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(add_one) + .to(MAIN_PROCESS) + .add_sink() + .get_config() + ) + markers = [p for p in config.pipes if isinstance(p, PlacementConfig)] + self.assertEqual(len(markers), 2) + self.assertEqual(markers[0].target, ProcessPoolExecutorConfig(max_workers=2)) + self.assertIs(markers[1].target, MAIN_PROCESS) + + def test_returns_self(self) -> None: + """to() returns the builder for chaining.""" + b = PipelineBuilder().add_source(range(4)) + self.assertIs(b.to(MAIN_PROCESS), b) + + def test_rejects_live_executor(self) -> None: + """Passing a live Executor is a TypeError pointing at pipe(executor=...).""" + b = PipelineBuilder().add_source(range(4)) + with ThreadPoolExecutor(max_workers=1) as ex: + with self.assertRaisesRegex(TypeError, "not a live Executor"): + b.to(ex) # type: ignore[arg-type] + + def test_rejects_wrong_type(self) -> None: + """Passing a non-spec, non-sentinel target is a TypeError.""" + b = PipelineBuilder().add_source(range(4)) + with self.assertRaisesRegex( + TypeError, "target must be a ProcessPoolExecutorConfig" + ): + b.to("subprocess") # type: ignore[arg-type] + + +class ToValidationTest(unittest.TestCase): + """Validation performed on to() regions at get_config()/build().""" + + def test_unclosed_region_before_sink_raises(self) -> None: + """A region left open before the sink is rejected.""" + b = ( + PipelineBuilder() + .add_source(range(4)) + .to(ProcessPoolExecutorConfig()) + .pipe(add_one) + .add_sink() + ) + with self.assertRaisesRegex(ValueError, "must be closed with"): + b.get_config() + + def test_input_order_inside_region_raises(self) -> None: + """A stage with output_order='input' inside a region is rejected.""" + b = ( + PipelineBuilder() + .add_source(range(4)) + .to(ProcessPoolExecutorConfig()) + .pipe(add_one, output_order="input") + .to(MAIN_PROCESS) + .add_sink() + ) + with self.assertRaisesRegex(ValueError, "output_order='input'"): + b.get_config() + + def test_output_order_input_in_path_variants_region_raises(self) -> None: + """output_order='input' nested in a path_variants branch in a region is rejected.""" + b = ( + PipelineBuilder() + .add_source(range(4)) + .to(ProcessPoolExecutorConfig()) + .path_variants(lambda _x: 0, [[Pipe(add_one, output_order="input")]]) + .to(MAIN_PROCESS) + .add_sink() + ) + with self.assertRaisesRegex(ValueError, "output_order='input'"): + b.get_config() + + def test_empty_region_raises(self) -> None: + """Opening a region with no stages before the next marker is rejected.""" + b = ( + PipelineBuilder() + .add_source(range(4)) + .to(ProcessPoolExecutorConfig()) + .to(MAIN_PROCESS) + .add_sink() + ) + with self.assertRaisesRegex(ValueError, "has no stages"): + b.get_config() + + def test_adjacent_nonempty_regions_are_valid(self) -> None: + """Back-to-back non-empty regions (different targets, no MAIN between) are allowed.""" + config = ( + PipelineBuilder() + .add_source(range(4)) + .to(ProcessPoolExecutorConfig(max_workers=1)) + .pipe(add_one) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(times_two) + .to(MAIN_PROCESS) + .add_sink() + .get_config() # must not raise: both regions have stages + ) + markers = [p for p in config.pipes if isinstance(p, PlacementConfig)] + self.assertEqual(len(markers), 3) + + def test_closed_region_is_valid(self) -> None: + """A properly closed region passes validation and preserves the marker sequence.""" + config = ( + PipelineBuilder() + .add_source(range(4)) + .to(ProcessPoolExecutorConfig()) + .pipe(add_one) + .to(MAIN_PROCESS) + .add_sink() + .get_config() # must not raise: the region is closed before the sink + ) + markers = [p for p in config.pipes if isinstance(p, PlacementConfig)] + self.assertEqual( + [type(m.target) for m in markers], [ProcessPoolExecutorConfig, _MainProcess] + ) + + @unittest.skipIf( + sys.version_info >= (3, 14), "subinterpreters are supported on 3.14+" + ) + def test_subinterpreter_region_rejected_before_314(self) -> None: + """On Python < 3.14, a subinterpreter region is rejected at get_config().""" + b = ( + PipelineBuilder() + .add_source(range(4)) + .to(InterpreterPoolExecutorConfig()) + .pipe(add_one) + .to(MAIN_PROCESS) + .add_sink() + ) + with self.assertRaisesRegex(RuntimeError, "requires Python 3.14"): + b.get_config() + + +class ToEndToEndTest(unittest.TestCase): + """The full public path: build and run a pipeline with a subprocess region.""" + + def test_subprocess_region_runs(self) -> None: + """A .to(ProcessPoolExecutorConfig()) region fuses and runs, matching the inline result.""" + n = 16 + pipeline = ( + PipelineBuilder() + .add_source(range(n)) + .to(ProcessPoolExecutorConfig(max_workers=2)) + .pipe(add_one) + .pipe(times_two) + .to(MAIN_PROCESS) + .add_sink() + .build(num_threads=4) + ) + self.assertEqual(sorted(_run(pipeline)), sorted((x + 1) * 2 for x in range(n))) diff --git a/tests/pipeline/marked_region_fuse_test.py b/tests/pipeline/marked_region_fuse_test.py index add3fe40c..b9294e0e0 100644 --- a/tests/pipeline/marked_region_fuse_test.py +++ b/tests/pipeline/marked_region_fuse_test.py @@ -6,13 +6,14 @@ # pyre-strict +import multiprocessing as mp import os import sys import unittest from collections.abc import Sequence from typing import Any -from spdl.pipeline import build_pipeline, Pipeline +from spdl.pipeline import build_pipeline, Pipeline, run_pipeline_in_subprocess from spdl.pipeline.defs import ( Aggregate, InterpreterPoolExecutorConfig, @@ -72,6 +73,30 @@ def stamp_main(item: _PidStamp) -> _PidStamp: return item +class _ProcStamp: + """Item recording ``(pid, ppid, daemon)`` of the stages that touch it, so a test can check + which process a stage ran in and whether that process is a daemon.""" + + def __init__(self, value: int, main_pid: int) -> None: + self.value = value + self.main_pid = main_pid + self.intermediate: tuple[int, int, bool] | None = None + self.region: tuple[int, int, bool] | None = None + + +def stamp_intermediate(item: _ProcStamp) -> _ProcStamp: + """A stage OUTSIDE any region: under ``run_pipeline_in_subprocess`` it runs in the + intermediate pipeline subprocess.""" + item.intermediate = (os.getpid(), os.getppid(), mp.current_process().daemon) + return item + + +def stamp_region_proc(item: _ProcStamp) -> _ProcStamp: + """A stage INSIDE a subprocess region: it runs in a region worker process.""" + item.region = (os.getpid(), os.getppid(), mp.current_process().daemon) + return item + + def _cfg(src: Any, pipes: Sequence[Any], buffer: int = 16) -> PipelineConfig[Any]: return PipelineConfig( src=SourceConfig(src), pipes=list(pipes), sink=SinkConfig(buffer) @@ -277,3 +302,67 @@ def test_raises_runtime_error(self) -> None: with self.assertRaises(RuntimeError) as cm: build_pipeline(config, num_threads=2) self.assertIn("3.14", str(cm.exception)) + + +class RegionUnderSubprocessTest(unittest.TestCase): + """A ``.to()`` region driven by :py:func:`run_pipeline_in_subprocess`. + + ``run_pipeline_in_subprocess`` runs the source/sink and any non-region stages in an + *intermediate* subprocess, while each ``.to(ProcessPoolExecutorConfig(...))`` region is fused in + the **main** process (via ``_fuse_marked_regions``) so its worker pool is main-owned -- + spawned by main, not by the intermediate subprocess (which, as a daemon, cannot have + children). These tests pin that placement and the daemon flags that let teardown reap + everything at exit. + """ + + def test_region_pool_is_main_owned_and_daemon(self) -> None: + """The region worker pool is spawned from main (hoisted), and both the intermediate + subprocess and the region workers are daemons. + + A stage outside the region records the intermediate subprocess's + ``(pid, ppid, daemon)``; a stage inside the region records the region worker's. The + region worker's parent must be the main process (not the intermediate subprocess), + and every spawned process must be a daemon so it is terminated at interpreter exit + rather than hang-joined. + """ + main_pid = os.getpid() + n = 4 + config = _cfg( + [_ProcStamp(i, main_pid) for i in range(n)], + [ + Pipe( + stamp_intermediate + ), # outside the region -> intermediate subprocess + PlacementConfig( + target=ProcessPoolExecutorConfig(max_workers=1, mp_context="spawn") + ), + Pipe(stamp_region_proc), # inside the region -> region worker + PlacementConfig(target=MAIN_PROCESS), + ], + ) + src = run_pipeline_in_subprocess( + config, + num_threads=2, + use_thread_output_queue=True, + buffer_size=8, + timeout=60.0, + mp_context="spawn", + daemon=True, + ) + results = list(src) + + self.assertEqual(len(results), n) + for item in results: + self.assertIsNotNone(item.intermediate) + self.assertIsNotNone(item.region) + inter_pid, inter_ppid, inter_daemon = item.intermediate + region_pid, region_ppid, region_daemon = item.region + # The intermediate subprocess is a daemon child of main. + self.assertEqual(inter_ppid, main_pid) + self.assertTrue(inter_daemon) + # The region worker pool is spawned from MAIN (hoisted), not the intermediate + # subprocess -- its parent is main, and distinct from the intermediate. + self.assertEqual(region_ppid, main_pid) + self.assertNotEqual(region_ppid, inter_pid) + # The region workers are daemons. + self.assertTrue(region_daemon) diff --git a/tests/pipeline/region_exit_test.py b/tests/pipeline/region_exit_test.py new file mode 100644 index 000000000..e756e9753 --- /dev/null +++ b/tests/pipeline/region_exit_test.py @@ -0,0 +1,91 @@ +# 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 + +"""Interpreter-exit teardown of a held `.to(ProcessPoolExecutorConfig)` region pipeline. + +A held-but-unstopped pipeline is stopped at interpreter shutdown by the hook that +:py:meth:`Pipeline.start` registers. This exercises that for the worker-region topology: the +scenario holds a `.to()` region pipeline (mimicking a framework keeping the dataloader), iterates +a few items, and returns WITHOUT stopping. Run in a child process: the child must exit cleanly +rather than hang joining the region's still-running worker pool at shutdown. +""" + +import multiprocessing as mp +import unittest +from typing import Any, Callable + +from spdl.pipeline import PipelineBuilder +from spdl.pipeline.defs import MAIN_PROCESS, ProcessPoolExecutorConfig + +# Strong reference that outlives the scenario function inside the child process. +_HELD: dict[str, object] = {} + + +def _double(x: int) -> int: + return x * 2 + + +def _scenario_to_region(mp_ctx: str) -> None: + """Hold a `.to(ProcessPoolExecutorConfig)` region pipeline; never stop it.""" + p = ( + PipelineBuilder() + .add_source(range(1000), continuous=True) + .aggregate(8) + .to(ProcessPoolExecutorConfig(max_workers=2, mp_context=mp_ctx)) + .pipe(_double) + .to(MAIN_PROCESS) + .add_sink(4) + .build(num_threads=2) + ) + _HELD["dl"] = p + for i, _ in enumerate(p.get_iterator(timeout=60)): + if i >= 10: + break + + +def _child_exit_code( + scenario: Callable[[str], None], mp_ctx: str, timeout: float = 120.0 +) -> int | None: + """Run ``scenario`` in a child process; return its exit code, or None if it hung.""" + ctx: Any = mp.get_context(mp_ctx) + proc = ctx.Process(target=scenario, args=(mp_ctx,)) + proc.start() + proc.join(timeout) + if proc.is_alive(): + proc.terminate() + proc.join(10) + return None + return proc.exitcode + + +class RegionInterpreterExitTest(unittest.TestCase): + """A held, unstopped `.to()` region pipeline exits cleanly at interpreter exit.""" + + @unittest.skipUnless( + "forkserver" in mp.get_all_start_methods(), + "forkserver start method is unavailable on this platform (e.g. Windows)", + ) + def test_to_region_forkserver(self) -> None: + """`.to(ProcessPoolExecutorConfig)` region pipeline exits cleanly (forkserver).""" + exit_code = _child_exit_code(_scenario_to_region, "forkserver") + self.assertIsNotNone( + exit_code, + "Child process hung (did not exit within timeout); the region worker pool " + "likely blocked interpreter shutdown.", + ) + self.assertEqual(exit_code, 0) + + def test_to_region_spawn(self) -> None: + """`.to(ProcessPoolExecutorConfig)` region pipeline exits cleanly (spawn).""" + exit_code = _child_exit_code(_scenario_to_region, "spawn") + self.assertIsNotNone( + exit_code, + "Child process hung (did not exit within timeout); the region worker pool " + "likely blocked interpreter shutdown.", + ) + self.assertEqual(exit_code, 0)