diff --git a/src/spdl/pipeline/_build.py b/src/spdl/pipeline/_build.py index ccf91388a..b2f14499c 100644 --- a/src/spdl/pipeline/_build.py +++ b/src/spdl/pipeline/_build.py @@ -139,6 +139,7 @@ def _build_pipeline( background_tasks: list[BackgroundTaskFactory] | None = None, use_thread_output_queue: bool = False, fuse_subprocess_stages: bool = False, + mp_context: str | None = None, ) -> Pipeline[U]: if _DEFAULT_BUILD_CALLBACK is not None: try: @@ -146,44 +147,66 @@ def _build_pipeline( except Exception: _LG.exception("Build callback failed.") - pools: list[Any] = [] - 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, pools = _fuse_subprocess_stages( - pipeline_cfg, report_stats_interval=report_stats_interval, stacklevel=4 + # Both ``_fuse_subprocess_stages`` and ``_hoist_process_pools`` eagerly spawn worker + # processes. They are handed to the returned Pipeline, which owns and reaps them. But if + # anything between spawning them and returning the Pipeline raises, ownership never + # transfers, so reap them here before propagating rather than leaking the worker processes + # and their pipe fds. (Each helper already reaps the pools it spawned if it raises partway; + # this guards the window *between* the helpers and the ``Pipeline`` construction.) + fuse_pools: list[Any] = [] + worker_pools: list[Any] = [] + try: + 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. + # stacklevel=4: _fuse_subprocess_stages -> _build_pipeline -> build_pipeline -> user. + pipeline_cfg, fuse_pools = _fuse_subprocess_stages( + pipeline_cfg, report_stats_interval=report_stats_interval, stacklevel=4 + ) + + # Treat any user-provided ``ProcessPoolExecutor`` as a specification and replace it with + # an eagerly-spawned, pipeline-owned worker pool (a ``_RemoteExecutor`` submit proxy + # backed by a ``_WorkerPool``). Spawning the workers now -- at build time, before the + # event-loop thread starts -- avoids forking a worker lazily mid-load from this + # multi-threaded process, which can corrupt the child. (Inside the + # ``run_pipeline_in_subprocess`` worker the process pools were already hoisted in the + # main process, so no ``ProcessPoolExecutor`` remains and this is a no-op.) + pipeline_cfg, worker_pools = _hoist_process_pools(pipeline_cfg, mp_context) + + desc = repr(pipeline_cfg) + + _LG.debug("%s", desc) + + # Merge per-pipeline background tasks with defaults + all_bg_tasks: list[BackgroundTaskFactory] = [] + default_bg = get_default_background_tasks() + if default_bg: + all_bg_tasks.extend(default_bg) + if background_tasks: + all_bg_tasks.extend(background_tasks) + + coro, queue = _build_pipeline_coro( + pipeline_cfg, + max_failures=max_failures, + report_stats_interval=report_stats_interval, + queue_class=queue_class, + task_hook_factory=task_hook_factory, + stage_id=stage_id, + background_tasks=all_bg_tasks or None, + use_thread_output_queue=use_thread_output_queue, ) - desc = repr(pipeline_cfg) - - _LG.debug("%s", desc) - - # Merge per-pipeline background tasks with defaults - all_bg_tasks: list[BackgroundTaskFactory] = [] - default_bg = get_default_background_tasks() - if default_bg: - all_bg_tasks.extend(default_bg) - if background_tasks: - all_bg_tasks.extend(background_tasks) - - coro, queue = _build_pipeline_coro( - pipeline_cfg, - max_failures=max_failures, - report_stats_interval=report_stats_interval, - queue_class=queue_class, - task_hook_factory=task_hook_factory, - stage_id=stage_id, - background_tasks=all_bg_tasks or None, - use_thread_output_queue=use_thread_output_queue, - ) - - executor = ThreadPoolExecutor( - max_workers=num_threads, - thread_name_prefix="spdl_worker_thread_", - ) - return Pipeline(coro, queue, executor, desc=desc, pools=pools) + executor = ThreadPoolExecutor( + max_workers=num_threads, + thread_name_prefix="spdl_worker_thread_", + ) + return Pipeline( + coro, queue, executor, desc=desc, pools=[*fuse_pools, *worker_pools] + ) + except BaseException: + _shutdown_pools(worker_pools) + _shutdown_pipeline_pools(fuse_pools) + raise def build_pipeline( @@ -199,6 +222,7 @@ def build_pipeline( background_tasks: list[BackgroundTaskFactory] | None = None, use_thread_output_queue: bool = False, fuse_subprocess_stages: bool = False, + mp_context: str | None = None, ) -> Pipeline[U]: """Build a pipeline from the config. @@ -287,6 +311,21 @@ def build_pipeline( .. versionadded:: 0.6.0 The ``fuse_subprocess_stages`` argument. + + mp_context: The multiprocessing start method (as accepted by + :py:func:`multiprocessing.get_context`, e.g. ``"spawn"`` or ``"forkserver"``) + used to spawn the worker processes for any stage configured with a + :py:class:`~concurrent.futures.ProcessPoolExecutor`. If ``None`` (default), the + platform default start method is used. + + .. note:: + + The ``"fork"`` start method can deadlock or corrupt a worker when the pipeline + is built from a process that already has other threads running. Prefer + ``"spawn"`` or ``"forkserver"`` in that case. + + .. versionadded:: 0.6.0 + The ``mp_context`` argument. """ from . import _profile @@ -304,6 +343,7 @@ def build_pipeline( background_tasks=background_tasks, use_thread_output_queue=use_thread_output_queue, fuse_subprocess_stages=fuse_subprocess_stages, + mp_context=mp_context, ) @@ -512,10 +552,12 @@ def run_pipeline_in_subprocess( ``InterpreterPoolExecutor`` are explicitly supported, even though these executors are not picklable. - Such an executor must be **freshly constructed** — handed over without any work + Such an executor should be **freshly constructed** — handed over without any work submitted yet — because its workers are (re)created as part of running the pipeline in - the subprocess (the whole point of moving execution there). Passing one that has already - spawned workers (i.e. been used) lifts it mid-lifecycle and raises :py:exc:`ValueError`. + the subprocess (the whole point of moving execution there). The executor is treated as a + *specification*: its own workers are not used. Passing one that has already spawned + workers (i.e. been used) emits a :py:exc:`RuntimeWarning` and continues; its workers are + not used and you remain responsible for shutting it down. - ``ThreadPoolExecutor`` / ``InterpreterPoolExecutor``: their constructor arguments are serialized and an equivalent executor (same type, same ``max_workers``) is diff --git a/src/spdl/pipeline/_builder.py b/src/spdl/pipeline/_builder.py index 0501aa215..0282b8647 100644 --- a/src/spdl/pipeline/_builder.py +++ b/src/spdl/pipeline/_builder.py @@ -171,6 +171,19 @@ def pipe( into asynchronous one. If ``None``, the default executor is used. It is invalid to provide this argument when the given op is already async. + + A :py:class:`~concurrent.futures.ProcessPoolExecutor` is treated as a + *specification*: the pipeline reads its worker count and initializer and runs + its own equivalent worker pool, which it spawns eagerly when the pipeline is + built and shuts down when the pipeline stops. Pass a freshly constructed + executor — one that has already spawned workers or had work submitted triggers + a :py:class:`RuntimeWarning`, its own workers are not used, and you remain + responsible for shutting it down. Use the ``mp_context`` argument of + :py:meth:`build` to control the start method of the spawned workers. + + .. versionchanged:: 0.6.0 + A ``ProcessPoolExecutor`` is now recreated and owned by the pipeline (its + workers are spawned eagerly at build time) rather than used directly. 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. @@ -293,6 +306,7 @@ def build( stage_id: int = 0, use_thread_output_queue: bool = False, fuse_subprocess_stages: bool = False, + mp_context: str | None = None, ) -> Pipeline[U]: """Build the pipeline. @@ -348,6 +362,16 @@ def build( .. versionadded:: 0.6.0 The ``fuse_subprocess_stages`` argument. + + mp_context: The multiprocessing start method (e.g. ``"spawn"`` or + ``"forkserver"``) used to spawn the worker processes for any stage configured + with a :py:class:`~concurrent.futures.ProcessPoolExecutor`. If ``None`` + (default), the platform default start method is used. Prefer ``"spawn"`` or + ``"forkserver"`` over ``"fork"`` when building from a process that already has + other threads running, as ``"fork"`` can deadlock or corrupt a worker. + + .. versionadded:: 0.6.0 + The ``mp_context`` argument. """ return build_pipeline( self.get_config(), @@ -359,4 +383,5 @@ def build( stage_id=stage_id, use_thread_output_queue=use_thread_output_queue, fuse_subprocess_stages=fuse_subprocess_stages, + mp_context=mp_context, ) diff --git a/src/spdl/pipeline/_executor_proxy.py b/src/spdl/pipeline/_executor_proxy.py index 81e5ca68b..b37ea5be5 100644 --- a/src/spdl/pipeline/_executor_proxy.py +++ b/src/spdl/pipeline/_executor_proxy.py @@ -28,6 +28,7 @@ import sys import threading +import warnings from collections.abc import Callable from concurrent.futures import Executor, ThreadPoolExecutor from dataclasses import replace @@ -41,9 +42,9 @@ ) __all__ = [ - "_ensure_executor_unused", "_make_config_executors_picklable", "_rewrite_config_executors", + "_warn_if_executor_used", ] @@ -179,16 +180,22 @@ def _interpreter_pool_kwargs(executor: Any) -> dict[str, Any]: _EXECUTOR_ARG_EXTRACTORS[InterpreterPoolExecutor] = _interpreter_pool_kwargs -def _ensure_executor_unused(executor: Executor) -> None: - """Reject a stdlib pool executor that has already been used. +def _warn_if_executor_used(executor: Executor) -> None: + """Warn if a stdlib pool executor handed to the pipeline has already been used. - Moving a pipeline to a subprocess recreates the executor's workers as part of running it - there (thread / interpreter pools are reconstructed in the subprocess; a - :py:class:`~concurrent.futures.ProcessPoolExecutor`'s workers are hoisted into the main - process). The executor must therefore be freshly constructed — one that has already - spawned workers or had work submitted is being lifted mid-lifecycle, which is not the - supported contract: the whole point is that execution, including the pool's worker startup, - happens as part of the run. Construct the executor and hand it over without using it first. + The pipeline treats a passed executor as a *specification*: it reads the executor's type + and constructor arguments and runs its own equivalent workers (thread / interpreter pools + are reconstructed; a :py:class:`~concurrent.futures.ProcessPoolExecutor`'s workers are + spawned eagerly in the main process). The executor should therefore be freshly + constructed. One that has already spawned workers or had work submitted is being handed + over mid-lifecycle: its own workers are not used by the pipeline, and the caller remains + responsible for shutting it down. Construct the executor and hand it over without using it + first. + + This helper runs deep inside the config-rewriting machinery and is reached via several + distinct, variable-depth build paths, so no fixed ``stacklevel`` reliably points at the + user's ``pipe()``/``build()`` call. The warning therefore carries a self-contained message; + its ``stacklevel`` is a best-effort constant rather than a per-call-site value. """ if ( getattr( @@ -199,10 +206,14 @@ def _ensure_executor_unused(executor: Executor) -> None: ) # Process pool: spawned worker processes or getattr(executor, "_queue_count", 0) # Process pool: work already submitted ): - raise ValueError( - "run_pipeline_in_subprocess() requires a freshly constructed executor with no " - "work submitted yet: its workers are (re)created when the pipeline runs in the " - "subprocess. Construct the executor and pass it without using it first." + warnings.warn( + "An executor passed to the SPDL pipeline has already spawned workers or had work " + "submitted. The pipeline treats the executor as a specification and runs its own " + "equivalent workers, so the passed executor's own workers are not used and you " + "remain responsible for shutting it down. Construct the executor and pass it " + "without using it first.", + RuntimeWarning, + stacklevel=2, ) @@ -217,7 +228,7 @@ def _maybe_proxy(executor: Executor | None) -> Executor | _ExecutorProxy | None: extractor = _EXECUTOR_ARG_EXTRACTORS.get(type(executor)) if extractor is None: return executor - _ensure_executor_unused(executor) + _warn_if_executor_used(executor) return _ExecutorProxy(type(executor), extractor(executor)) diff --git a/src/spdl/pipeline/_fuse.py b/src/spdl/pipeline/_fuse.py index a5c041b54..04c1c66b0 100644 --- a/src/spdl/pipeline/_fuse.py +++ b/src/spdl/pipeline/_fuse.py @@ -45,7 +45,7 @@ from spdl.pipeline._common._convert import _is_process_pool from spdl.pipeline._components import _get_global_id, _set_global_id -from spdl.pipeline._executor_proxy import _ensure_executor_unused +from spdl.pipeline._executor_proxy import _warn_if_executor_used from spdl.pipeline._subprocess_pipeline_pool import _SubprocessPipelinePool from spdl.pipeline.defs._defs import ( _PipeType, @@ -343,7 +343,7 @@ def _worker_initializer( 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) + _warn_if_executor_used(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 () diff --git a/src/spdl/pipeline/_subprocess_worker_pool.py b/src/spdl/pipeline/_subprocess_worker_pool.py index ea282389c..289f59973 100644 --- a/src/spdl/pipeline/_subprocess_worker_pool.py +++ b/src/spdl/pipeline/_subprocess_worker_pool.py @@ -44,8 +44,8 @@ from typing import Any, TypeVar from spdl.pipeline._executor_proxy import ( - _ensure_executor_unused, _rewrite_config_executors, + _warn_if_executor_used, ) from spdl.pipeline.defs import PipelineConfig @@ -349,7 +349,7 @@ def convert(executor: Any) -> Any: key = id(executor) if key in seen: return seen[key] - _ensure_executor_unused(executor) + _warn_if_executor_used(executor) if not ctx_box: ctx = mp.get_context(mp_context) if ctx.get_start_method() == "fork" and threading.active_count() > 1: diff --git a/src/spdl/pipeline/defs/_defs.py b/src/spdl/pipeline/defs/_defs.py index 8b0ee72c3..42fa0e173 100644 --- a/src/spdl/pipeline/defs/_defs.py +++ b/src/spdl/pipeline/defs/_defs.py @@ -791,6 +791,18 @@ def Pipe( into asynchronous one. If ``None``, the default executor is used. It is invalid to provide this argument when the given op is already async. + + A :py:class:`~concurrent.futures.ProcessPoolExecutor` is treated as a + *specification*: the pipeline reads its worker count and initializer and runs its + own equivalent worker pool, which it spawns eagerly when the pipeline is built and + shuts down when the pipeline stops. Pass a freshly constructed executor — one that + has already spawned workers or had work submitted triggers a + :py:class:`RuntimeWarning`, its own workers are not used, and you remain + responsible for shutting it down. + + .. versionchanged:: 0.6.0 + A ``ProcessPoolExecutor`` is now recreated and owned by the pipeline (its + workers are spawned eagerly at build time) rather than used directly. 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/executor_ownership_test.py b/tests/pipeline/executor_ownership_test.py new file mode 100644 index 000000000..3b5e77aed --- /dev/null +++ b/tests/pipeline/executor_ownership_test.py @@ -0,0 +1,157 @@ +# 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 concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor + +from spdl.pipeline import PipelineBuilder + + +def _double(i: int) -> int: + return i * 2 + + +class TestProcessPoolOwnership(unittest.TestCase): + """The pipeline recreates and owns a process-pool executor passed to a stage.""" + + def test_process_pool_recreated_and_owned(self) -> None: + """A ProcessPoolExecutor becomes an eagerly-spawned, pipeline-owned worker pool.""" + executor = ProcessPoolExecutor(max_workers=2) + try: + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(_double, executor=executor) + .add_sink(100) + .build(num_threads=1, mp_context="spawn") + ) + # The user's executor is used only as a spec — its own workers never spawn. + self.assertEqual(len(executor._processes), 0) + # The pipeline owns an equivalent worker pool, spawned eagerly at build time. + self.assertEqual(len(pipeline._impl._pools), 1) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator()) + self.assertEqual(results, [i * 2 for i in range(10)]) + + # The user's executor stayed a pure spec: even after the pipeline ran to + # completion, it never spawned a worker or had work submitted to it. + self.assertEqual(len(executor._processes), 0) + self.assertEqual(executor._queue_count, 0) + + # Stopping reaps the owned pool exactly once. + self.assertEqual(len(pipeline._impl._pools), 0) + finally: + executor.shutdown() + + def test_used_process_pool_warns(self) -> None: + """Passing an already-used ProcessPoolExecutor warns at build time.""" + executor = ProcessPoolExecutor(max_workers=2) + try: + # Submitting work spawns workers, marking the executor as used. + executor.submit(_double, 1).result() + + builder = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(_double, executor=executor) + .add_sink(100) + ) + with self.assertWarns(RuntimeWarning): + pipeline = builder.build(num_threads=1, mp_context="spawn") + # Clean up the owned pool that the build spawned. + pipeline.stop() + finally: + executor.shutdown() + + def test_shared_process_pool_owned_once(self) -> None: + """A process pool attached to multiple stages maps to a single owned pool.""" + executor = ProcessPoolExecutor(max_workers=2) + try: + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(_double, executor=executor) + .pipe(_double, executor=executor) + .add_sink(100) + .build(num_threads=1, mp_context="spawn") + ) + self.assertEqual(len(pipeline._impl._pools), 1) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator()) + self.assertEqual(results, [i * 4 for i in range(10)]) + finally: + executor.shutdown() + + +class TestThreadPoolUnchanged(unittest.TestCase): + """Thread-pool executors are not recreated or owned by this change.""" + + def test_thread_pool_not_owned(self) -> None: + """A ThreadPoolExecutor is used directly and not added to the owned pools.""" + executor = ThreadPoolExecutor(max_workers=2) + try: + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(_double, executor=executor) + .add_sink(100) + .build(num_threads=1) + ) + self.assertEqual(len(pipeline._impl._pools), 0) + + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator()) + self.assertEqual(results, [i * 2 for i in range(10)]) + finally: + executor.shutdown() + + +class TestContinuousSourceOwnership(unittest.TestCase): + """Pipeline ownership is what cleans up a process pool in continuous mode. + + With a continuous source the pipeline is started once and reused across epochs, + so there is no ``auto_stop`` block scoping its lifetime. Ownership is what makes + this safe: because the pipeline (not the caller) owns the recreated worker pool, + stopping the pipeline reaps the workers even though the caller never wrapped + iteration in ``auto_stop``. + """ + + def test_continuous_source_without_auto_stop(self) -> None: + """A continuous-source process-pool pipeline reaps its owned pool on stop, no auto_stop.""" + executor = ProcessPoolExecutor(max_workers=2) + try: + pipeline = ( + PipelineBuilder() + .add_source(range(5), continuous=True) + .pipe(_double, executor=executor) + .add_sink(100) + .build(num_threads=1, mp_context="spawn") + ) + # The pipeline owns an eagerly-spawned pool; the user executor is untouched. + self.assertEqual(len(pipeline._impl._pools), 1) + self.assertEqual(len(executor._processes), 0) + + # Iterate several epochs without auto_stop and without an explicit start: each + # ``for`` pass consumes one epoch and the background thread auto-starts on the first. + for _ in range(3): + epoch = sorted(item for item in pipeline) + self.assertEqual(epoch, [i * 2 for i in range(5)]) + + # The owned pool's workers ran the whole time; the user executor never did. + self.assertEqual(len(executor._processes), 0) + self.assertEqual(executor._queue_count, 0) + + # Stopping the pipeline (the caller never opened an ``auto_stop`` block) reaps the + # owned pool exactly once: ``_shutdown_pools`` clears the owned-pool list. This is a + # deterministic check of the ownership contract, independent of GC/finalizer timing. + pipeline.stop() + self.assertEqual(len(pipeline._impl._pools), 0) + finally: + executor.shutdown() diff --git a/tests/pipeline/pipeline_builder_test.py b/tests/pipeline/pipeline_builder_test.py index 4edaef925..62dc69e14 100644 --- a/tests/pipeline/pipeline_builder_test.py +++ b/tests/pipeline/pipeline_builder_test.py @@ -2943,12 +2943,12 @@ def test_run_in_subprocess_with_interpreterpool(self) -> None: results = list(run_pipeline_in_subprocess(config, num_threads=1)) self.assertEqual(sorted(results), [2 * i for i in range(10)]) - def test_used_thread_pool_is_rejected(self) -> None: - """A ThreadPoolExecutor that already ran work is rejected (must be freshly built).""" + def test_used_thread_pool_warns(self) -> None: + """A ThreadPoolExecutor that already ran work warns (it is used only as a spec).""" tpe = ThreadPoolExecutor(max_workers=2) try: tpe.submit(_sync_double, 1).result(timeout=30) # spawns a worker thread - with self.assertRaises(ValueError): + with self.assertWarns(RuntimeWarning): _make_config_executors_picklable(_config_with_executor(tpe)) finally: tpe.shutdown() @@ -3032,13 +3032,14 @@ def test_remote_executor_exposes_max_workers(self) -> None: finally: _shutdown_pools([pool]) - def test_hoist_rejects_used_process_pool(self) -> None: - """A ProcessPoolExecutor that already ran work is rejected (must be freshly built).""" + def test_hoist_warns_on_used_process_pool(self) -> None: + """A ProcessPoolExecutor that already ran work warns (it is used only as a spec).""" ppe = ProcessPoolExecutor(max_workers=2) try: ppe.submit(_sync_double, 1).result(timeout=30) # spawns worker processes - with self.assertRaises(ValueError): - _hoist_process_pools(_config_with_executor(ppe)) + with self.assertWarns(RuntimeWarning): + _, pools = _hoist_process_pools(_config_with_executor(ppe)) + _shutdown_pools(pools) finally: ppe.shutdown()