Skip to content

Commit c37f897

Browse files
committed
[pipeline] 4/4: Add PipelineBuilder region-placement method and validation
Part 4/4 of the `.to()` region API; see #1584 for the overall design and rationale. The public surface — the engine (#1585) and backend (#1586) already support regions end-to-end, so this makes the feature usable. - Adds `PipelineBuilder.to(target)`: appends a `PlacementConfig` marker designating where the subsequent stages run. `target` is a `ProcessPoolExecutorConfig`, `InterpreterPoolExecutorConfig`, or `MAIN_PROCESS`. The builder stays a stateless append-only list (no region state). A live `Executor` is rejected with a `TypeError` pointing at `pipe(executor=...)`, keeping the pipeline expressible as static config. - Adds stateless validation in `get_config()` (so it also covers `build()` and static configs): rejects a region left open before the sink, a stage with `output_order="input"` inside a region (order cannot be preserved across independent workers), and a subinterpreter region on Python < 3.14. - `fuse_subprocess_stages` is left in place; the two paths coexist until the follow-up removal.
1 parent 4597cbd commit c37f897

4 files changed

Lines changed: 528 additions & 3 deletions

File tree

src/spdl/pipeline/_builder.py

Lines changed: 151 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
# pyre-strict
88

99
import logging
10+
import sys
1011
from collections.abc import AsyncIterable, Callable, Iterable, Sequence
1112
from concurrent.futures import Executor
1213
from fractions import Fraction
@@ -15,17 +16,22 @@
1516
from spdl._internal import log_api_usage_once
1617
from spdl.pipeline._components import AsyncQueue, StageInfo, TaskHook
1718
from spdl.pipeline.defs import (
19+
_MainProcess,
20+
_PipeType,
1821
_TPipeInputs,
1922
Aggregate,
2023
AggregateConfig,
2124
Aggregator,
2225
Disaggregate,
2326
DisaggregateConfig,
27+
InterpreterPoolExecutorConfig,
2428
PathVariants,
2529
PathVariantsConfig,
2630
Pipe,
2731
PipeConfig,
2832
PipelineConfig,
33+
PlacementConfig,
34+
ProcessPoolExecutorConfig,
2935
SinkConfig,
3036
SourceConfig,
3137
)
@@ -45,6 +51,79 @@
4551
U_ = TypeVar("U_")
4652

4753

54+
def _has_ordered_pipe(cfg: object) -> bool:
55+
"""Whether ``cfg`` is (or, for path-variants, contains) an ``output_order="input"`` pipe.
56+
57+
Recurses into :py:class:`~spdl.pipeline.defs.PathVariantsConfig` branches: those stages run
58+
inside the region worker too, so an input-ordered pipe nested in a branch is just as invalid
59+
as a top-level one (global input order cannot be preserved across the pool's workers).
60+
"""
61+
if isinstance(cfg, PipeConfig):
62+
return cfg._type is _PipeType.OrderedPipe
63+
if isinstance(cfg, PathVariantsConfig):
64+
return any(_has_ordered_pipe(s) for path in cfg.paths for s in path)
65+
return False
66+
67+
68+
def _validate_executor_regions(
69+
pipes: Sequence[
70+
PipeConfig
71+
| AggregateConfig
72+
| DisaggregateConfig
73+
| PathVariantsConfig
74+
| PlacementConfig
75+
],
76+
) -> None:
77+
"""Validate the :py:meth:`PipelineBuilder.to` region markers in ``pipes``.
78+
79+
Stateless scan (a pipeline starts on the main process). Rejects: a subinterpreter region on
80+
Python < 3.14; a stage using ``output_order="input"`` inside a region (including one nested
81+
in a ``path_variants`` branch, since order cannot be preserved across independent workers); an
82+
empty region (a ``.to(...)`` marker with no stages before the next marker/sink); and a region
83+
left open at the sink.
84+
"""
85+
in_region = False
86+
region_has_stage = False
87+
for p in pipes:
88+
if isinstance(p, PlacementConfig):
89+
# A non-main region that opened but never received a stage does nothing; reject it
90+
# rather than silently dropping it (usually a stray or duplicated ``.to(...)``).
91+
# Adjacent *non-empty* regions with different targets remain valid.
92+
if in_region and not region_has_stage:
93+
raise ValueError(
94+
"An empty `to(...)` region has no stages. Add stages before the next "
95+
"`.to(...)`, or remove the redundant marker."
96+
)
97+
target = p.target
98+
in_region = not isinstance(target, _MainProcess)
99+
region_has_stage = False
100+
if isinstance(
101+
target, InterpreterPoolExecutorConfig
102+
) and sys.version_info < (3, 14):
103+
raise RuntimeError(
104+
"A subinterpreter region (`to(InterpreterPoolExecutorConfig(...))`) requires "
105+
"Python 3.14 or later. Current version: "
106+
f"{sys.version_info.major}.{sys.version_info.minor}"
107+
)
108+
elif in_region:
109+
region_has_stage = True
110+
if _has_ordered_pipe(p):
111+
raise ValueError(
112+
"A stage with `output_order='input'` cannot run inside a `to()` region: "
113+
"input order cannot be preserved across the region's independent workers."
114+
)
115+
if in_region and not region_has_stage:
116+
raise ValueError(
117+
"An empty `to(...)` region has no stages. Add stages before closing it, or remove "
118+
"the redundant marker."
119+
)
120+
if in_region:
121+
raise ValueError(
122+
"A `to()` execution region must be closed with `to(MAIN_PROCESS)` before "
123+
"`add_sink()`. The sink always runs in the main process."
124+
)
125+
126+
48127
################################################################################
49128
# Builder
50129
################################################################################
@@ -81,7 +160,11 @@ def __init__(self) -> None:
81160

82161
self._src: SourceConfig[T] | None = None
83162
self._process_args: list[
84-
PipeConfig | AggregateConfig | DisaggregateConfig | PathVariantsConfig
163+
PipeConfig
164+
| AggregateConfig
165+
| DisaggregateConfig
166+
| PathVariantsConfig
167+
| PlacementConfig
85168
] = []
86169
self._sink: SinkConfig[U] | None = None
87170

@@ -246,6 +329,66 @@ def disaggregate(self) -> "PipelineBuilder[T, U]":
246329
self._process_args.append(Disaggregate())
247330
return self
248331

332+
def to(
333+
self,
334+
target: "ProcessPoolExecutorConfig | InterpreterPoolExecutorConfig | _MainProcess",
335+
/,
336+
) -> "PipelineBuilder[T, U]":
337+
"""**[Experimental]** Designate where the subsequent stages execute.
338+
339+
Opens (or closes) an *execution region*: every stage added after this call runs on
340+
``target`` until the next :py:meth:`to`. A pipeline starts on the main process, so a
341+
region is opened by ``to(ProcessPoolExecutorConfig(...))`` or ``to(InterpreterPoolExecutorConfig(...))``
342+
and closed by ``to(MAIN_PROCESS)``. The stages inside a region are fused into one nested
343+
pipeline that runs together in a worker process (or subinterpreter), so the value handed
344+
from one stage to the next stays in the worker — it is **not** copied back to the main
345+
process between stages and need not be picklable. Only the region's inputs and outputs
346+
cross the boundary.
347+
348+
Unlike passing ``executor=`` to individual :py:meth:`pipe` calls, a region also carries
349+
:py:meth:`aggregate`, :py:meth:`disaggregate`, and :py:meth:`path_variants` stages into
350+
the worker, and gives the worker-pool configuration a single home.
351+
352+
Args:
353+
target: Where the following stages run.
354+
355+
- :py:class:`~spdl.pipeline.defs.ProcessPoolExecutorConfig` — a pool of worker processes.
356+
- :py:class:`~spdl.pipeline.defs.InterpreterPoolExecutorConfig` — a pool of subinterpreters
357+
(Python 3.14+; the region's ops must avoid NumPy/PyTorch, which cannot be
358+
imported in a subinterpreter).
359+
- :py:data:`~spdl.pipeline.defs.MAIN_PROCESS` — close the current region; the
360+
following stages run in the main process.
361+
362+
A live :py:class:`~concurrent.futures.Executor` is **not** accepted — pass a spec
363+
so the pipeline stays expressible as static config. To run a single stage on a
364+
custom executor, use ``pipe(executor=...)`` instead.
365+
366+
.. note::
367+
368+
The region must be closed with ``to(MAIN_PROCESS)`` before :py:meth:`add_sink`, a
369+
stage inside a region may not use ``output_order="input"`` (order cannot be preserved
370+
across independent workers), and a subinterpreter region requires Python 3.14+. These
371+
are checked when the pipeline is built.
372+
373+
.. versionadded:: 0.6.0
374+
"""
375+
if isinstance(target, Executor):
376+
raise TypeError(
377+
"`to()` takes a serializable execution target (ProcessPoolExecutorConfig, "
378+
"InterpreterPoolExecutorConfig, or MAIN_PROCESS), not a live Executor. To run a single "
379+
"stage on a custom executor, pass it to `pipe(executor=...)` instead."
380+
)
381+
if not isinstance(
382+
target,
383+
(ProcessPoolExecutorConfig, InterpreterPoolExecutorConfig, _MainProcess),
384+
):
385+
raise TypeError(
386+
"`to()` target must be a ProcessPoolExecutorConfig, InterpreterPoolExecutorConfig, or "
387+
f"MAIN_PROCESS. Got: {type(target).__name__}."
388+
)
389+
self._process_args.append(PlacementConfig(target=target))
390+
return self
391+
249392
def path_variants(
250393
self,
251394
router: Callable,
@@ -286,14 +429,20 @@ def get_config(self) -> PipelineConfig[U]:
286429
A PipelineConfig object representing the current pipeline configuration.
287430
288431
Raises:
289-
RuntimeError: If source or sink is not set.
432+
RuntimeError: If source or sink is not set, or a subinterpreter region is used on
433+
Python < 3.14.
434+
ValueError: If an execution region opened by :py:meth:`to` is not closed with
435+
``to(MAIN_PROCESS)`` before the sink, or a stage inside a region uses
436+
``output_order="input"``.
290437
"""
291438
if (src := self._src) is None:
292439
raise RuntimeError("Source is not set. Did you call `add_source`?")
293440

294441
if (sink := self._sink) is None:
295442
raise RuntimeError("Sink is not set. Did you call `add_sink`?")
296443

444+
_validate_executor_regions(self._process_args)
445+
297446
return PipelineConfig(src, self._process_args, sink)
298447

299448
def build(

0 commit comments

Comments
 (0)