|
7 | 7 | # pyre-strict |
8 | 8 |
|
9 | 9 | import logging |
| 10 | +import sys |
10 | 11 | from collections.abc import AsyncIterable, Callable, Iterable, Sequence |
11 | 12 | from concurrent.futures import Executor |
12 | 13 | from fractions import Fraction |
|
15 | 16 | from spdl._internal import log_api_usage_once |
16 | 17 | from spdl.pipeline._components import AsyncQueue, StageInfo, TaskHook |
17 | 18 | from spdl.pipeline.defs import ( |
| 19 | + _MainProcess, |
| 20 | + _PipeType, |
18 | 21 | _TPipeInputs, |
19 | 22 | Aggregate, |
20 | 23 | AggregateConfig, |
21 | 24 | Aggregator, |
22 | 25 | Disaggregate, |
23 | 26 | DisaggregateConfig, |
| 27 | + InterpreterPoolExecutorConfig, |
24 | 28 | PathVariants, |
25 | 29 | PathVariantsConfig, |
26 | 30 | Pipe, |
27 | 31 | PipeConfig, |
28 | 32 | PipelineConfig, |
| 33 | + PlacementConfig, |
| 34 | + ProcessPoolExecutorConfig, |
29 | 35 | SinkConfig, |
30 | 36 | SourceConfig, |
31 | 37 | ) |
|
45 | 51 | U_ = TypeVar("U_") |
46 | 52 |
|
47 | 53 |
|
| 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 | + |
48 | 127 | ################################################################################ |
49 | 128 | # Builder |
50 | 129 | ################################################################################ |
@@ -81,7 +160,11 @@ def __init__(self) -> None: |
81 | 160 |
|
82 | 161 | self._src: SourceConfig[T] | None = None |
83 | 162 | self._process_args: list[ |
84 | | - PipeConfig | AggregateConfig | DisaggregateConfig | PathVariantsConfig |
| 163 | + PipeConfig |
| 164 | + | AggregateConfig |
| 165 | + | DisaggregateConfig |
| 166 | + | PathVariantsConfig |
| 167 | + | PlacementConfig |
85 | 168 | ] = [] |
86 | 169 | self._sink: SinkConfig[U] | None = None |
87 | 170 |
|
@@ -246,6 +329,66 @@ def disaggregate(self) -> "PipelineBuilder[T, U]": |
246 | 329 | self._process_args.append(Disaggregate()) |
247 | 330 | return self |
248 | 331 |
|
| 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 | + |
249 | 392 | def path_variants( |
250 | 393 | self, |
251 | 394 | router: Callable, |
@@ -286,14 +429,20 @@ def get_config(self) -> PipelineConfig[U]: |
286 | 429 | A PipelineConfig object representing the current pipeline configuration. |
287 | 430 |
|
288 | 431 | 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"``. |
290 | 437 | """ |
291 | 438 | if (src := self._src) is None: |
292 | 439 | raise RuntimeError("Source is not set. Did you call `add_source`?") |
293 | 440 |
|
294 | 441 | if (sink := self._sink) is None: |
295 | 442 | raise RuntimeError("Sink is not set. Did you call `add_sink`?") |
296 | 443 |
|
| 444 | + _validate_executor_regions(self._process_args) |
| 445 | + |
297 | 446 | return PipelineConfig(src, self._process_args, sink) |
298 | 447 |
|
299 | 448 | def build( |
|
0 commit comments