diff --git a/src/spdl/pipeline/_components/_node.py b/src/spdl/pipeline/_components/_node.py index eaf7898b5..157955d70 100644 --- a/src/spdl/pipeline/_components/_node.py +++ b/src/spdl/pipeline/_components/_node.py @@ -26,6 +26,7 @@ PathVariantsConfig, PipeConfig, PipelineConfig, + PlacementConfig, SinkConfig, SourceConfig, ) @@ -312,6 +313,7 @@ def _convert_pipes( | DisaggregateConfig | PathVariantsConfig | _SubprocessPipelineConfig + | PlacementConfig ], n: _TNodes, q_class: type[AsyncQueue], @@ -340,6 +342,12 @@ def _convert_pipes( match cfg: case PathVariantsConfig(): n = _convert_path_variants(cfg, n, q_class, pipeline_id, stage_id, idx) + case PlacementConfig(): + # Region markers are build-time directives resolved before the pipeline is + # built; they must never reach node construction. + raise ValueError( + "PlacementConfig region markers must be resolved before building nodes." + ) case _: in_q = _get_output_queue(n, idx) info = _get_stage_info(cfg, pipeline_id, stage_id) diff --git a/src/spdl/pipeline/_profile.py b/src/spdl/pipeline/_profile.py index 10009f529..30d435288 100644 --- a/src/spdl/pipeline/_profile.py +++ b/src/spdl/pipeline/_profile.py @@ -25,6 +25,7 @@ PathVariantsConfig, PipeConfig, PipelineConfig, + PlacementConfig, SinkConfig, SourceConfig, ) @@ -263,7 +264,9 @@ def _profile_pipeline( raise ValueError(f"Unexpected source type {type(cfg.src)}") for pipe in cfg.pipes: - if isinstance(pipe, (PathVariantsConfig, _SubprocessPipelineConfig)): + if isinstance( + pipe, (PathVariantsConfig, _SubprocessPipelineConfig, PlacementConfig) + ): _LG.warning( "Skipping %s stage in profiling (not supported).", type(pipe).__name__ ) diff --git a/src/spdl/pipeline/defs/__init__.py b/src/spdl/pipeline/defs/__init__.py index 3feaea7c9..f67126bf5 100644 --- a/src/spdl/pipeline/defs/__init__.py +++ b/src/spdl/pipeline/defs/__init__.py @@ -26,6 +26,7 @@ """ from ._defs import ( + _MainProcess, _PipeArgs, _PipeType, _SubprocessPipelineConfig, @@ -36,6 +37,9 @@ Collate, Disaggregate, DisaggregateConfig, + ExecutorConfig, + InterpreterPoolExecutorConfig, + MAIN_PROCESS, Merge, MergeConfig, PathVariants, @@ -43,11 +47,14 @@ Pipe, PipeConfig, PipelineConfig, + PlacementConfig, + ProcessPoolExecutorConfig, SinkConfig, SourceConfig, ) __all__ = [ + "_MainProcess", "_PipeArgs", "_PipeType", "_SubprocessPipelineConfig", @@ -62,6 +69,9 @@ "Collate", "Disaggregate", "DisaggregateConfig", + "ExecutorConfig", + "InterpreterPoolExecutorConfig", + "MAIN_PROCESS", "Merge", "MergeConfig", "PathVariants", @@ -69,6 +79,8 @@ "Pipe", "PipeConfig", "PipelineConfig", + "PlacementConfig", + "ProcessPoolExecutorConfig", "SinkConfig", "SourceConfig", ] diff --git a/src/spdl/pipeline/defs/_defs.py b/src/spdl/pipeline/defs/_defs.py index 3d934659e..954223a0a 100644 --- a/src/spdl/pipeline/defs/_defs.py +++ b/src/spdl/pipeline/defs/_defs.py @@ -40,6 +40,12 @@ "Collate", "DisaggregateConfig", "_SubprocessPipelineConfig", + "_MainProcess", + "ExecutorConfig", + "ProcessPoolExecutorConfig", + "InterpreterPoolExecutorConfig", + "MAIN_PROCESS", + "PlacementConfig", "PipelineConfig", "SinkConfig", "SourceConfig", @@ -436,6 +442,110 @@ class _SubprocessPipelineConfig: """The picklable submit-side handle (``_SubprocessPipelineHandle``) for the worker pool.""" +################################################################################ +# Executor placement (region markers for `.to()`) +################################################################################ + + +@dataclass(frozen=True) +class ExecutorConfig: + """**[Experimental]** Base spec for a worker-pool executor. + + A serializable description of a pool of workers, materialized into a live executor by + the pipeline. Subclassed by :py:class:`ProcessPoolExecutorConfig` and + :py:class:`InterpreterPoolExecutorConfig`; those subclasses are used as ``.to()`` + placement targets (via :py:class:`PlacementConfig`). + + .. versionadded:: 0.6.0 + """ + + max_workers: int | None = None + """Number of workers. If ``None``, defaults to the number of CPUs.""" + + initializer: Callable[..., object] | None = None + """Callable run once in each worker before it processes any work.""" + + initargs: tuple[Any, ...] = () + """Positional arguments passed to ``initializer``.""" + + +@dataclass(frozen=True) +class ProcessPoolExecutorConfig(ExecutorConfig): + """**[Experimental]** A worker-pool executor backed by subprocesses. + + Used as a ``.to()`` target to run a region's stages in a pool of worker *processes* + as one nested pipeline — the op->op handoff stays in the worker, so intermediate + values are not copied back to the main process and need not be picklable. The region + ends at the next ``.to()`` target (see :py:data:`MAIN_PROCESS`). + + .. versionadded:: 0.6.0 + """ + + mp_context: str | None = None + """Multiprocessing start method (e.g. ``"spawn"``, ``"fork"``, ``"forkserver"``), + as accepted by :py:func:`multiprocessing.get_context`. ``None`` uses the default + context.""" + + +@dataclass(frozen=True) +class InterpreterPoolExecutorConfig(ExecutorConfig): + """**[Experimental]** A worker-pool executor backed by subinterpreters. + + Like :py:class:`ProcessPoolExecutorConfig`, but the workers are Python + *subinterpreters* (:py:mod:`concurrent.interpreters`) sharing the process rather than + separate processes. There is no ``mp_context`` because no new process is started. + + .. note:: + + Requires Python 3.14 or later; using this target on an older interpreter is an + error. NumPy and PyTorch cannot be imported inside a subinterpreter, so a region + whose stages need them must use :py:class:`ProcessPoolExecutorConfig` instead. + + .. versionadded:: 0.6.0 + """ + + +@dataclass(frozen=True) +class _MainProcess: + """Sentinel ``.to()`` target for the main process. Use the :py:data:`MAIN_PROCESS` + singleton rather than instantiating this type.""" + + def __repr__(self) -> str: + return "MAIN_PROCESS" + + +MAIN_PROCESS: _MainProcess = _MainProcess() +"""**[Experimental]** The main-process execution target. Pass to +:py:meth:`spdl.pipeline.PipelineBuilder.to` +to close a worker-pool region and bring subsequent stages back to the main process. + +.. versionadded:: 0.6.0 +""" + + +@dataclass(frozen=True) +class PlacementConfig: + """**[Experimental]** A region marker designating where the subsequent stages execute. + + Sits among the stage configs in :py:attr:`PipelineConfig.pipes`: every stage after + this marker (until the next :py:class:`PlacementConfig`) runs on :py:attr:`target`. + :py:meth:`spdl.pipeline.PipelineBuilder.to` appends one of these. A pipeline + implicitly starts on the main process, so a marker is needed only to enter a + worker-pool region and (with :py:data:`MAIN_PROCESS`) to leave it. + + .. versionadded:: 0.6.0 + """ + + target: "ProcessPoolExecutorConfig | InterpreterPoolExecutorConfig | _MainProcess" + """Where the stages following this marker execute.""" + + name: str = "placement" + """Name of the marker (used only for display).""" + + def __repr__(self) -> str: + return f"{self.name}({self.target!r})" + + ################################################################################ # PathVariants ################################################################################ @@ -607,6 +717,7 @@ class PipelineConfig(Generic[U]): | DisaggregateConfig[Any] | PathVariantsConfig[Any] | _SubprocessPipelineConfig + | PlacementConfig ] """Pipe configurations.""" diff --git a/tests/pipeline/executor_config_test.py b/tests/pipeline/executor_config_test.py new file mode 100644 index 000000000..c9b3c0691 --- /dev/null +++ b/tests/pipeline/executor_config_test.py @@ -0,0 +1,64 @@ +# 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 spdl.pipeline.defs import ( + InterpreterPoolExecutorConfig, + MAIN_PROCESS, + PlacementConfig, + ProcessPoolExecutorConfig, +) + +# NOTE: The behavior of these config types (how a region's worker pool is built and run from a +# spec), including that a ``PipelineConfig`` accepts region markers in its ``pipes``, is covered +# end-to-end by ``marked_region_fuse_test`` and ``builder_to_test``. The tests here cover only what +# those cannot: the hand-written ``__repr__`` methods and the deliberate structural contract of +# ``InterpreterPoolExecutorConfig`` (that it rejects an ``mp_context``). Plain-dataclass mechanics +# (default values, field storage, ``frozen``, auto-generated ``__eq__``) are intentionally not +# re-tested. + + +class TestInterpreterPoolExecutorConfig(unittest.TestCase): + """Verify the InterpreterPoolExecutorConfig spec.""" + + def test_rejects_mp_context(self) -> None: + """InterpreterPoolExecutorConfig rejects ``mp_context``; ProcessPoolExecutorConfig honors it. + + A subinterpreter shares the process, so there is no multiprocessing start method to choose. + Passing ``mp_context`` must raise rather than be silently accepted and ignored -- unlike + :py:class:`ProcessPoolExecutorConfig`, which starts real processes and stores the chosen + start method. + """ + self.assertEqual( + ProcessPoolExecutorConfig(mp_context="spawn").mp_context, "spawn" + ) + with self.assertRaises(TypeError): + InterpreterPoolExecutorConfig(mp_context="spawn") # pyre-ignore[28] + + +class TestMainProcess(unittest.TestCase): + """Verify the MAIN_PROCESS sentinel target.""" + + def test_repr(self) -> None: + """The sentinel's custom ``__repr__`` renders its public name for readable configs.""" + self.assertEqual(repr(MAIN_PROCESS), "MAIN_PROCESS") + + +class TestPlacementConfig(unittest.TestCase): + """Verify the PlacementConfig region marker.""" + + def test_repr_includes_target(self) -> None: + """The marker's custom ``__repr__`` surfaces its target for readable pipeline dumps.""" + self.assertEqual( + repr(PlacementConfig(target=MAIN_PROCESS)), "placement(MAIN_PROCESS)" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pipeline/shutdown_hook_test.py b/tests/pipeline/shutdown_hook_test.py index aed7983b6..076b46f7b 100644 --- a/tests/pipeline/shutdown_hook_test.py +++ b/tests/pipeline/shutdown_hook_test.py @@ -214,13 +214,29 @@ def _child_exit_code( class InterpreterExitTopologyTest(unittest.TestCase): """Held, unstopped pipelines exit cleanly at interpreter exit, across topologies.""" + @unittest.skipUnless( + "forkserver" in mp.get_all_start_methods(), + "forkserver start method is unavailable on this platform (e.g. Windows)", + ) def test_plain_forkserver(self) -> None: """Plain thread-only pipeline exits cleanly (forkserver).""" - self.assertEqual(_child_exit_code(_scenario_plain, "forkserver"), 0) + exit_code = _child_exit_code(_scenario_plain, "forkserver") + self.assertIsNotNone( + exit_code, + "Child process hung (did not exit within timeout); the pipeline's " + "non-daemon event-loop thread likely blocked interpreter shutdown.", + ) + self.assertEqual(exit_code, 0) def test_plain_spawn(self) -> None: """Plain thread-only pipeline exits cleanly (spawn).""" - self.assertEqual(_child_exit_code(_scenario_plain, "spawn"), 0) + exit_code = _child_exit_code(_scenario_plain, "spawn") + self.assertIsNotNone( + exit_code, + "Child process hung (did not exit within timeout); the pipeline's " + "non-daemon event-loop thread likely blocked interpreter shutdown.", + ) + self.assertEqual(exit_code, 0) # Note: the `.to(ProcessPoolExecutorConfig)` region topology's interpreter-exit teardown is