From a09f9ab615c24aa5a7900e91243938b6a97b627e Mon Sep 17 00:00:00 2001 From: Moto Hira Date: Mon, 29 Jun 2026 09:11:26 -0700 Subject: [PATCH 1/2] Fix unreliable PipelineFailure propagation with thread output queue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: ## Problem With `use_thread_output_queue=True`, the sink's output queue routes through `loop.run_in_executor(...)`, perturbing event-loop scheduling during shutdown. A failing stage raises inside `_queue_stage_hook`, which does `await queue.put(_EOF)` (a suspension point) *before* re-raising. When the downstream sink completes first, `_run_pipeline_coroutines` → `_cancel_orphaned` cancels the still-suspended upstream task before it re-raises; the `CancelledError` masks the real error and `PipelineFailure` was dropped on ~50% of runs. ## Root-cause fix Instead of recovering the masked error after the fact, prevent the masking: `_cancel_recursive` now skips cancelling a task that has already failed (detected via the terminal exception `_queue_stage_hook` records before its EOF handoff). The failed task resumes and re-raises naturally; its upstream producers are still cancelled. `_gather_error` reverts to its simple form (`done` + not-cancelled + `isinstance Exception`). The `use_thread_output_queue=False` path is unaffected — the skip is a no-op unless a stage actually failed. [Session trajectory link](https://www.internalfb.com/intern/devai/devmate/inspector/?id=1b804c2a-bac7-49ce-ac13-ebcb631cd954) Differential Revision: D110018369 --- src/spdl/pipeline/_components/_node.py | 23 +++++++- src/spdl/pipeline/_components/_queue.py | 30 +++++++++- tests/pipeline/pipeline_builder_test.py | 74 +++++++++++++++++++++++++ tests/pipeline/pipeline_node_test.py | 65 ++++++++++++++++++++++ 4 files changed, 188 insertions(+), 4 deletions(-) diff --git a/src/spdl/pipeline/_components/_node.py b/src/spdl/pipeline/_components/_node.py index eaf7898b5..978946580 100644 --- a/src/spdl/pipeline/_components/_node.py +++ b/src/spdl/pipeline/_components/_node.py @@ -41,7 +41,12 @@ _ordered_pipe, _pipe, ) -from ._queue import _ThreadBasedAsyncQueue, AsyncQueue, get_default_queue_class +from ._queue import ( + _get_stage_exc, + _ThreadBasedAsyncQueue, + AsyncQueue, + get_default_queue_class, +) from ._sink import _sink from ._source import _source, _source_continuous from ._subprocess_pipe import _subprocess_pipeline @@ -780,7 +785,15 @@ def _start_tasks(node: _TNodes) -> set[Task]: def _cancel_recursive(node: _TNodes) -> None: - node.task.cancel() + # Do not cancel a task that has already failed and recorded its terminal + # exception (it is only suspended at its EOF handoff and will re-raise once + # resumed). Cancelling it would replace the real exception with + # CancelledError and mask the failure -- a race exposed by the thread-backed + # sink queue, which lets a downstream stage complete (and trigger this + # cancellation) before the failed task is resumed. Its upstream producers + # are still cancelled so they are not left orphaned. + if _get_stage_exc(node.task) is None: + node.task.cancel() for n in node.upstream: _cancel_recursive(n) @@ -822,7 +835,11 @@ def _gather_error( task = node.task errs = [] - if not task.cancelled() and (err := task.exception()) is not None: + if ( + task.done() + and not task.cancelled() + and isinstance(err := task.exception(), Exception) + ): errs.append((task.get_name(), err)) for n in node.upstream: diff --git a/src/spdl/pipeline/_components/_queue.py b/src/spdl/pipeline/_components/_queue.py index 4a90dd506..f6ab3ef4d 100644 --- a/src/spdl/pipeline/_components/_queue.py +++ b/src/spdl/pipeline/_components/_queue.py @@ -21,6 +21,8 @@ __all__ = [ "_queue_stage_hook", + "_get_stage_exc", + "_STAGE_EXC_ATTR", "AsyncQueue", "_ThreadBasedAsyncQueue", "StatsQueue", @@ -31,6 +33,21 @@ _LG: logging.Logger = logging.getLogger(__name__) +# Attribute used to stash a stage's terminal exception onto its asyncio Task. +# See `_queue_stage_hook` / `_get_stage_exc` for why this is needed. +_STAGE_EXC_ATTR: str = "_spdl_stage_exc" + + +def _get_stage_exc(task: "asyncio.Task[Any]") -> Exception | None: + """Return the terminal exception recorded by `_queue_stage_hook`, if any. + + A failing stage records its exception onto its own Task before yielding at the + EOF handoff. This marks the task as already-failed so that orphan cancellation + during shutdown can skip it (and let it re-raise naturally) instead of masking + the real error with `CancelledError`. + """ + return getattr(task, _STAGE_EXC_ATTR, None) + class AsyncQueue(asyncio.Queue): """Extends :py:class:`asyncio.Queue` with init/finalize logic. @@ -78,7 +95,18 @@ async def _queue_stage_hook(queue: AsyncQueue) -> AsyncGenerator[None, None]: async with queue.stage_hook(): try: yield - except Exception: + except Exception as e: + # Record the terminal exception onto the running Task *before* the + # `await queue.put(_EOF)` suspension point, marking this task as + # already-failed. During shutdown a downstream stage can complete + # first and trigger orphan cancellation of this still-suspended + # Task; the resulting CancelledError would replace the real error + # before it is re-raised, masking the failure (only observed with + # the thread-backed sink queue, which perturbs scheduling order). + # `_cancel_recursive` reads this mark and skips cancelling the task + # so it resumes and re-raises its real exception. + if (task := asyncio.current_task()) is not None: + setattr(task, _STAGE_EXC_ATTR, e) await queue.put(_EOF) raise else: diff --git a/tests/pipeline/pipeline_builder_test.py b/tests/pipeline/pipeline_builder_test.py index 4edaef925..2f96abd45 100644 --- a/tests/pipeline/pipeline_builder_test.py +++ b/tests/pipeline/pipeline_builder_test.py @@ -3355,3 +3355,77 @@ def failing_range(n): self.assertTrue( any(note.startswith("Pipeline stage:") for note in notes), ) + + +class TestPipelineThreadOutputQueueFailure(unittest.TestCase): + """Exception propagation must be reliable with the thread-backed sink queue. + + With ``use_thread_output_queue=True`` the sink's output queue routes through + ``run_in_executor``, which perturbs event-loop scheduling during shutdown. This + used to let a failing stage be orphan-cancelled before it could re-raise, so the + ``PipelineFailure`` was dropped on ~50% of runs. These tests stress the failure + paths to confirm the error surfaces on every iteration. + """ + + _STRESS_RUNS: int = 50 + + def _assert_always_fails(self, build_pipeline) -> None: + for _ in range(self._STRESS_RUNS): + pipeline = build_pipeline() + with self.assertRaises(PipelineFailure): + with pipeline.auto_stop(): + list(pipeline.get_iterator(timeout=30)) + + def test_max_failures(self) -> None: + """Exceeding max_failures reliably fails with the thread output queue.""" + + def fail_odd(x): + if x % 2: + raise ValueError(f"Only even numbers are allowed. {x}") + return x + + def build(): + return ( + PipelineBuilder() + .add_source(range(10)) + .pipe(fail_odd) + .add_sink(1) + .build(num_threads=1, max_failures=3, use_thread_output_queue=True) + ) + + self._assert_always_fails(build) + + def test_type_error(self) -> None: + """A stage TypeError reliably surfaces with the thread output queue.""" + + async def wrong_sig(i, _): + return i + + def build(): + return ( + PipelineBuilder() + .add_source(range(10)) + # pyre-ignore[6] + .pipe(wrong_sig) + .add_sink(1) + .build(num_threads=1, use_thread_output_queue=True) + ) + + self._assert_always_fails(build) + + def test_source_failure(self) -> None: + """A source exception reliably surfaces with the thread output queue.""" + + def failure_source(): + raise RuntimeError("Foo") + yield None + + def build(): + return ( + PipelineBuilder() + .add_source(failure_source()) + .add_sink(1) + .build(num_threads=1, use_thread_output_queue=True) + ) + + self._assert_always_fails(build) diff --git a/tests/pipeline/pipeline_node_test.py b/tests/pipeline/pipeline_node_test.py index eab1bb24b..dd4ae6cb5 100644 --- a/tests/pipeline/pipeline_node_test.py +++ b/tests/pipeline/pipeline_node_test.py @@ -21,6 +21,7 @@ _SourceNode, _start_tasks, ) +from spdl.pipeline._components._queue import _STAGE_EXC_ATTR from spdl.pipeline.defs import SinkConfig, SourceConfig @@ -261,6 +262,70 @@ async def run() -> None: asyncio.run(run()) + def test_cancel_recursive_skips_already_failed(self) -> None: + """`_cancel_recursive` skips an already-failed task (so its real + exception is not masked) while still cancelling its upstream.""" + # A task that has already failed and recorded its terminal exception + # (i.e. it is suspended at its EOF handoff, mid-failure) must NOT be + # cancelled, otherwise CancelledError masks the real exception. Its + # upstream producers are still cancelled. + # + # A -> B (already failed, suspended) + + async def run() -> None: + a = _node("A", []) + b = _node("B", [a]) # sleeps: stands in for "suspended at put(_EOF)" + tasks = _start_tasks(b) + await asyncio.sleep(0) + + # Mark B as already-failed, as `_queue_stage_hook` does before its + # `await queue.put(_EOF)` suspension point. + setattr(b.task, _STAGE_EXC_ATTR, DummyException("fail B")) + + _cancel_recursive(b) + await asyncio.sleep(0) + + self.assertFalse(b.task.cancelled()) # skipped: already failed + self.assertTrue(a.task.cancelled()) # upstream still cancelled + + b.task.cancel() # cleanup + await asyncio.wait(tasks) + + asyncio.run(run()) + + def test_cancel_orphaned_skips_already_failed_upstream(self) -> None: + """Orphan cancellation skips an already-failed upstream task (so its + exception is not masked) while still cancelling further upstream.""" + # When a downstream task is done, orphan cancellation must skip an + # upstream task that has already failed (and is only suspended at its + # EOF handoff), so its exception is not masked by CancelledError. + # + # A -> B (already failed, suspended) -> C (done) + + async def run() -> None: + a = _node("A", []) + b = _node("B", [a]) # sleeps: stands in for "suspended at put(_EOF)" + c = _node("C", [b]) + tasks = _start_tasks(c) + await asyncio.sleep(0) + + # B recorded its terminal exception but has not re-raised yet. + setattr(b.task, _STAGE_EXC_ATTR, DummyException("fail B")) + # C has finished (e.g. it read B's EOF). + c.task.cancel() + await asyncio.wait([c.task]) + + _cancel_orphaned(c) + await asyncio.sleep(0) + + self.assertFalse(b.task.cancelled()) # skipped: already failed + self.assertTrue(a.task.cancelled()) # upstream still cancelled + + b.task.cancel() # cleanup + await asyncio.wait(tasks) + + asyncio.run(run()) + def test_gather_error_with_cancelled(self) -> None: # A1 B1 # | | From 92cf24bf72d84a7a1f20700bb26442ddda8609ca Mon Sep 17 00:00:00 2001 From: Moto Hira Date: Mon, 29 Jun 2026 09:11:26 -0700 Subject: [PATCH 2/2] spdl: default use_thread_output_queue to True MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Make the thread-backed sink output queue the default across the public pipeline build entry points by flipping the default of `use_thread_output_queue` from `False` to `True`. This affects `build_pipeline`, `PipelineBuilder.build`, `run_pipeline_in_subprocess`, and `run_pipeline_in_subinterpreter` (plus the internal forwarders that thread the argument through). When enabled, the sink hands the final batch from the background event loop to the foreground consumer thread via a `queue.Queue`-backed queue instead of `asyncio.run_coroutine_threadsafe`, cutting per-batch handoff latency from ~200-400us to ~10us. This changes a public default value in a backward-incompatible way, so the title is tagged `[BC-breaking]` and the affected docstrings carry a `.. versionchanged:: 0.6.0` directive. `profile_pipeline` reads the sink output queue's occupancy/lap stats, which only the asyncio stats queue collects, so it now builds its internal pipelines with `use_thread_output_queue=False` explicitly. Known issue — must be resolved before landing: pipeline-level exception propagation (e.g. `PipelineFailure` raised when a stage exceeds its failure threshold, or any raising stage) is currently unreliable (~50%) when the thread-backed output queue is used. This is a pre-existing race in the thread-output-queue path that is merely exposed by making it the default; it is documented and tracked separately. See the test plan for how to validate the fix. Differential Revision: D109953066 --- src/spdl/pipeline/_build.py | 15 +++++++++------ src/spdl/pipeline/_builder.py | 7 +++++-- src/spdl/pipeline/_components/_node.py | 6 +++--- src/spdl/pipeline/_profile.py | 7 ++++++- 4 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/spdl/pipeline/_build.py b/src/spdl/pipeline/_build.py index ccf91388a..86b4b7100 100644 --- a/src/spdl/pipeline/_build.py +++ b/src/spdl/pipeline/_build.py @@ -137,7 +137,7 @@ def _build_pipeline( task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, background_tasks: list[BackgroundTaskFactory] | None = None, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, fuse_subprocess_stages: bool = False, ) -> Pipeline[U]: if _DEFAULT_BUILD_CALLBACK is not None: @@ -197,7 +197,7 @@ def build_pipeline( task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, background_tasks: list[BackgroundTaskFactory] | None = None, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, fuse_subprocess_stages: bool = False, ) -> Pipeline[U]: """Build a pipeline from the config. @@ -269,7 +269,10 @@ def build_pipeline( :py:class:`queue.Queue`-backed queue for the final handoff from the background event loop to the foreground consumer thread. This bypasses ``asyncio.run_coroutine_threadsafe``, reducing per-batch latency from - ~200-400us to ~10us. Default: ``False``. + ~200-400us to ~10us. Default: ``True``. + + .. versionchanged:: 0.6.0 + ``use_thread_output_queue`` now defaults to ``True`` (was ``False``). fuse_subprocess_stages: If ``True``, fuse runs of two or more adjacent pipe stages that share the same process-pool (or interpreter-pool) executor instance into a single @@ -332,7 +335,7 @@ def __init__( queue_class: type[AsyncQueue] | None, task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, background_tasks: list[BackgroundTaskFactory] | None = None, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, ) -> None: self.config = config self.num_threads = num_threads @@ -404,7 +407,7 @@ def run_pipeline_in_subprocess( queue_class: type[AsyncQueue] | None = None, task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, background_tasks: list[BackgroundTaskFactory] | None = None, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, fuse_subprocess_stages: bool = False, **kwargs: Any, ) -> Iterable[T]: @@ -675,7 +678,7 @@ def run_pipeline_in_subinterpreter( queue_class: type[AsyncQueue] | None = None, task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, background_tasks: list[BackgroundTaskFactory] | None = None, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, **kwargs: Any, ) -> Iterable[T]: """**[Experimental]** Run the given Pipeline in a subinterpreter, and iterate on the result. diff --git a/src/spdl/pipeline/_builder.py b/src/spdl/pipeline/_builder.py index 0501aa215..8c194bc9b 100644 --- a/src/spdl/pipeline/_builder.py +++ b/src/spdl/pipeline/_builder.py @@ -291,7 +291,7 @@ def build( queue_class: type[AsyncQueue] | None = None, task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, fuse_subprocess_stages: bool = False, ) -> Pipeline[U]: """Build the pipeline. @@ -333,7 +333,10 @@ def build( use_thread_output_queue: If ``True``, replace the sink's output queue with a ``queue.Queue``-backed queue for lower-latency batch handoff. - Default: ``False``. + Default: ``True``. + + .. versionchanged:: 0.6.0 + ``use_thread_output_queue`` now defaults to ``True`` (was ``False``). fuse_subprocess_stages: If ``True``, fuse runs of two or more adjacent pipe stages that share the same process-pool (or interpreter-pool) executor instance into a diff --git a/src/spdl/pipeline/_components/_node.py b/src/spdl/pipeline/_components/_node.py index 978946580..679e5af87 100644 --- a/src/spdl/pipeline/_components/_node.py +++ b/src/spdl/pipeline/_components/_node.py @@ -446,7 +446,7 @@ def _convert_config( pipeline_id: int, stage_id: _MutableInt, disable_sink: bool = False, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, ) -> _TOutputNodes: """Convert a :py:class:`~spdl.pipeline.defs.PipelineConfig` into a linked list of :py:class:`~spdl.pipeline._components._node._Node` objects. @@ -723,7 +723,7 @@ def _build_pipeline_node( queue_class: type[AsyncQueue] | None, task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None, stage_id: int, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, ) -> _TOutputNodes: global _PIPELINE_ID _PIPELINE_ID += 1 @@ -997,7 +997,7 @@ def _build_pipeline_coro( task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, background_tasks: Sequence[BackgroundTaskFactory] | None = None, - use_thread_output_queue: bool = False, + use_thread_output_queue: bool = True, ) -> tuple[Coroutine[None, None, None], asyncio.Queue]: try: node = _build_pipeline_node( diff --git a/src/spdl/pipeline/_profile.py b/src/spdl/pipeline/_profile.py index 10009f529..0ddef6e6d 100644 --- a/src/spdl/pipeline/_profile.py +++ b/src/spdl/pipeline/_profile.py @@ -211,7 +211,12 @@ def _profile_pipe( cfg_ = _build_pipeline_config(inputs, pipe, max(concurrencies)) outputs = [] for concurrency in concurrencies: - pipeline = _build._build_pipeline(cfg_, num_threads=concurrency) + # Profiling reads the sink output queue's lap stats (occupancy_rate), which only + # the stats-collecting AsyncQueue provides. Force the asyncio-backed output queue + # here; the thread-backed output queue (the default) does not collect stats. + pipeline = _build._build_pipeline( + cfg_, num_threads=concurrency, use_thread_output_queue=False + ) with hook_.stage_profile_hook(pipe.name, concurrency): qps_, outputs = _run(pipeline)