diff --git a/src/spdl/pipeline/_build.py b/src/spdl/pipeline/_build.py index 803e81d9f..8c59c8440 100644 --- a/src/spdl/pipeline/_build.py +++ b/src/spdl/pipeline/_build.py @@ -128,6 +128,7 @@ def _build_pipeline( task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, background_tasks: list[BackgroundTaskFactory] | None = None, + use_priority_scheduler: bool = False, ) -> Pipeline[U]: if _DEFAULT_BUILD_CALLBACK is not None: try: @@ -147,6 +148,33 @@ def _build_pipeline( if background_tasks: all_bg_tasks.extend(background_tasks) + # Create executor before building the pipeline so the scheduler can + # reference it via _underlying_executor attribute binding. + executor = ThreadPoolExecutor( + max_workers=num_threads, + thread_name_prefix="spdl_worker_thread_", + ) + + # Construct PriorityScheduler if requested. Per V5.4 plumbing, we + # bind the underlying ThreadPoolExecutor by direct attribute + # assignment (no _bind_executor method) and wire its run loop as a + # BackgroundTask via _PrioritySchedulerBackgroundTask. + scheduler = None + if use_priority_scheduler: + from spdl.pipeline._scheduler import ( + _PrioritySchedulerBackgroundTask, + PriorityScheduler, + ) + + scheduler = PriorityScheduler(max_concurrent=num_threads) + scheduler._underlying_executor = executor + + # Capture in a default arg so the lambda refers to *this* scheduler + # instance (avoid late-binding in a loop, defensive). + all_bg_tasks.append( + lambda sched=scheduler: _PrioritySchedulerBackgroundTask(sched) + ) + coro, queue = _build_pipeline_coro( pipeline_cfg, max_failures=max_failures, @@ -155,12 +183,9 @@ def _build_pipeline( task_hook_factory=task_hook_factory, stage_id=stage_id, background_tasks=all_bg_tasks or None, + scheduler=scheduler, ) - executor = ThreadPoolExecutor( - max_workers=num_threads, - thread_name_prefix="spdl_worker_thread_", - ) return Pipeline(coro, queue, executor, desc=desc) @@ -175,6 +200,7 @@ def build_pipeline( task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, background_tasks: list[BackgroundTaskFactory] | None = None, + use_priority_scheduler: bool = False, ) -> Pipeline[U]: """Build a pipeline from the config. @@ -240,6 +266,11 @@ def build_pipeline( :py:meth:`~BackgroundTask.run` method runs alongside the pipeline stages. Tasks are cancelled when the pipeline completes. Their errors are logged but do not cause the pipeline to fail. + + use_priority_scheduler: If ``True``, enable priority-based + dispatch for sync stages via :py:class:`PriorityScheduler`. + Deeper stages (closer to sink) are given higher priority, + reducing pipeline bubble time. """ from . import _profile @@ -255,6 +286,7 @@ def build_pipeline( task_hook_factory=task_hook_factory, stage_id=stage_id, background_tasks=background_tasks, + use_priority_scheduler=use_priority_scheduler, ) diff --git a/src/spdl/pipeline/_builder.py b/src/spdl/pipeline/_builder.py index 5ec26abb5..cbd2678f5 100644 --- a/src/spdl/pipeline/_builder.py +++ b/src/spdl/pipeline/_builder.py @@ -291,6 +291,7 @@ def build( queue_class: type[AsyncQueue] | None = None, task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, + use_priority_scheduler: bool = False, ) -> Pipeline[U]: """Build the pipeline. @@ -328,6 +329,11 @@ def build( To disable hooks, provide a function that returns an empty list. stage_id: The index of the initial stage used for logging. + + use_priority_scheduler: If ``True``, enable priority-based + dispatch for sync stages via :py:class:`PriorityScheduler`. + Deeper stages (closer to sink) are given higher priority, + reducing pipeline bubble time. """ return build_pipeline( self.get_config(), @@ -337,4 +343,5 @@ def build( report_stats_interval=report_stats_interval, task_hook_factory=task_hook_factory, stage_id=stage_id, + use_priority_scheduler=use_priority_scheduler, ) diff --git a/src/spdl/pipeline/_components/__init__.py b/src/spdl/pipeline/_components/__init__.py index 437020fc0..0cbe6ced1 100644 --- a/src/spdl/pipeline/_components/__init__.py +++ b/src/spdl/pipeline/_components/__init__.py @@ -20,6 +20,7 @@ set_default_queue_class, StatsQueue, ) +from ._semaphore import ResizableSemaphore __all__ = [ "_build_pipeline_coro", @@ -30,6 +31,7 @@ "is_eof", "is_epoch_end", "PipelineFailure", + "ResizableSemaphore", "set_default_hook_class", "set_default_queue_class", "TaskHook", diff --git a/src/spdl/pipeline/_components/_node.py b/src/spdl/pipeline/_components/_node.py index 3596c5a65..c1652c246 100644 --- a/src/spdl/pipeline/_components/_node.py +++ b/src/spdl/pipeline/_components/_node.py @@ -5,11 +5,12 @@ # LICENSE file in the root directory of this source tree. import asyncio +import inspect import logging import sys from asyncio import ALL_COMPLETED, FIRST_COMPLETED, Task from collections.abc import Callable, Coroutine, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from fractions import Fraction from functools import partial from typing import Any, TypeAlias, TypeVar @@ -484,6 +485,8 @@ def _build_node( fc_class: type[_FailCounter], task_hook_factory: Callable[[StageInfo], list[TaskHook]], max_failures: int | Fraction, + scheduler: Any = None, + depth: int = 0, ) -> None: """Build a coroutine for a single node based on its configuration type. @@ -520,6 +523,17 @@ def _build_node( task_hook_factory: A factory function for creating task hooks for monitoring. max_failures: The maximum number of failures allowed before halting. + scheduler: Optional :py:class:`PriorityScheduler` instance. When + provided, sync stages get a per-stage + :py:class:`_PrioritizedExecutor` shim injected into + ``_PipeArgs.executor`` so that + :py:meth:`asyncio.AbstractEventLoop.run_in_executor` routes + through the scheduler's priority heap instead of the + underlying ``ThreadPoolExecutor``'s FIFO. + depth: Depth of this node in the pipeline graph (source-to-sink + distance). Used to compute scheduler priority as + ``priority = -depth`` (deeper stages dispatch first). Only + consumed when ``scheduler is not None``. Raises: ValueError: If an unsupported configuration type is encountered. @@ -578,14 +592,31 @@ def _build_node( in_q, out_q = node.input_queue, node.output_queue hooks = task_hook_factory(node.info) fc = fc_class(max_failures, cfg._max_failures) + + args = cfg._args + + # When a scheduler is provided, register this stage's + # priority and inject a per-stage executor (created + # by `scheduler.make_stage_executor`) so + # loop.run_in_executor() routes through the heap. + # Only sync ops are routed (async/generator ops + # bypass the executor entirely in convert_to_async). + # The factory call decouples `_node.py` from the + # concrete `_PrioritizedExecutor` class — this avoids + # a Buck dep cycle through `_scheduler.py`. + if scheduler is not None and _is_sync_op(args): + scheduler.register_stage(node.info, priority=-depth) + per_stage_exec = scheduler.make_stage_executor(node.info) + args = replace(args, executor=per_stage_exec) + match cfg._type: case _PipeType.Pipe: node._coro = _pipe( - node.info, in_q, out_q, cfg._args, fc, hooks, False + node.info, in_q, out_q, args, fc, hooks, False ) case _PipeType.OrderedPipe: node._coro = _ordered_pipe( - node.info, in_q, out_q, cfg._args, fc, hooks + node.info, in_q, out_q, args, fc, hooks ) case _: # pragma: no cover raise ValueError( @@ -623,11 +654,52 @@ def _build_node( ) +def _is_sync_op(args: _PipeArgs) -> bool: + """Whether ``convert_to_async`` will use the executor branch for ``args.op``. + + The :py:class:`PriorityScheduler` only routes work that + :py:func:`~spdl.pipeline._common._convert.convert_to_async` would + submit through a :py:class:`~concurrent.futures.Executor`. Coroutine + functions and async-gen functions bypass the executor entirely, so + they must NOT receive a :py:class:`_PrioritizedExecutor` shim. + Generator functions and process-pool branches are also excluded for + Diff 2 to keep the scope minimal. + """ + op = args.op + if inspect.iscoroutinefunction(op) or inspect.isasyncgenfunction(op): + return False + if inspect.isgeneratorfunction(op): + # Generators take the _to_async_gen branch, which uses + # loop.run_in_executor on `next` rather than the user op as a + # whole — routing through the scheduler doesn't fit cleanly. + return False + if args.executor is not None: + # User-supplied executor (e.g., ProcessPoolExecutor). Don't + # override. + return False + return True + + +def _node_depth(node: _TNodes) -> int: + """Compute a node's depth (distance from the nearest source). + + Source nodes have depth 0; each downstream node is one deeper than + the deepest of its upstream nodes. Used by :py:func:`_build_node` + to compute scheduler priorities (priority = -depth). + """ + if isinstance(node, _SourceNode): + return 0 + if not node.upstream: + return 0 + return 1 + max(_node_depth(n) for n in node.upstream) + + def _build_node_recursive( node: _TNodes, fc_class: type[_FailCounter], task_hook_factory: Callable[[StageInfo], list[TaskHook]], max_failures: int | Fraction, + scheduler: Any = None, ) -> None: """Recursively build coroutines for a node and all its upstream nodes. @@ -640,6 +712,8 @@ def _build_node_recursive( fc_class: The failure counter class for tracking task failures. task_hook_factory: A factory function for creating task hooks for monitoring. max_failures: The maximum number of failures allowed before halting. + scheduler: Optional :py:class:`PriorityScheduler` instance for + priority dispatch. Raises: RuntimeError: If attempting to build a coroutine for a node that already has one. @@ -648,9 +722,10 @@ def _build_node_recursive( return for n in node.upstream: - _build_node_recursive(n, fc_class, task_hook_factory, max_failures) + _build_node_recursive(n, fc_class, task_hook_factory, max_failures, scheduler) - _build_node(node, fc_class, task_hook_factory, max_failures) + depth = _node_depth(node) if scheduler is not None else 0 + _build_node(node, fc_class, task_hook_factory, max_failures, scheduler, depth=depth) # Used to append stage name with pipeline @@ -696,6 +771,7 @@ def _build_pipeline_node( queue_class: type[AsyncQueue] | None, task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None, stage_id: int, + scheduler: Any = None, ) -> _TOutputNodes: global _PIPELINE_ID _PIPELINE_ID += 1 @@ -710,7 +786,7 @@ def _build_pipeline_node( fc_class = _get_fail_counter() node = _convert_config(plc, q_class, _PIPELINE_ID, _MutableInt(stage_id)) _validate_continuous_mode(node) - _build_node_recursive(node, fc_class, hook_factory, max_failures) + _build_node_recursive(node, fc_class, hook_factory, max_failures, scheduler) return node @@ -950,6 +1026,7 @@ def _build_pipeline_coro( task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, background_tasks: Sequence[BackgroundTaskFactory] | None = None, + scheduler: Any = None, ) -> tuple[Coroutine[None, None, None], asyncio.Queue]: try: node = _build_pipeline_node( @@ -959,6 +1036,7 @@ def _build_pipeline_coro( queue_class=queue_class, task_hook_factory=task_hook_factory, stage_id=stage_id, + scheduler=scheduler, ) coro = _run_pipeline_coroutines(node, background_tasks=background_tasks) diff --git a/src/spdl/pipeline/_components/_semaphore.py b/src/spdl/pipeline/_components/_semaphore.py new file mode 100644 index 000000000..c19ba9f18 --- /dev/null +++ b/src/spdl/pipeline/_components/_semaphore.py @@ -0,0 +1,189 @@ +# 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 + +"""Resizable asyncio semaphore for dynamic concurrency control.""" + +__all__ = ["ResizableSemaphore"] + +import asyncio +from collections import deque + + +class ResizableSemaphore: + """asyncio.Semaphore variant whose max value can be changed at runtime. + + Thread safety: asyncio is single-threaded per event loop. All methods + MUST be called from coroutines in the same event loop. No locks needed. + + Resize semantics: + - Increase: immediately wake up to ``new_max - old_max`` blocked + waiters. + - Decrease: no preemption. Currently acquired permits continue. + New ``acquire()`` calls block until active count drops below + the new max. ``_current_value`` may go negative during drain. + + Invariant: at any moment, the number of "active" (acquired but not + yet released) permits equals ``max_value - _current_value``. When + ``_current_value`` is negative, more permits are outstanding than + the current max allows -- they drain naturally as tasks + ``release()``. + """ + + def __init__(self, value: int) -> None: + """Create a semaphore with *value* initial permits. + + Args: + value: Initial max permits. Must be >= 1. + + Raises: + ValueError: If value < 1. + """ + if value < 1: + raise ValueError(f"value must be >= 1, got {value}") + self._max_value: int = value + self._current_value: int = value + self._waiters: deque[asyncio.Future[None]] = deque() + + @property + def max_value(self) -> int: + """Current max permits (may differ from initial after resize).""" + return self._max_value + + @property + def active(self) -> int: + """Number of currently acquired (outstanding) permits. + + Can exceed ``max_value`` temporarily after a resize-down. + """ + return self._max_value - self._current_value + + async def acquire(self) -> None: + """Acquire one permit. Blocks if no permits available. + + Raises: + asyncio.CancelledError: If the waiting coroutine is + cancelled while blocked. + """ + # Fast path: permit available and no one queued ahead of us. + if self._current_value > 0 and not self._waiters: + self._current_value -= 1 + return + + fut: asyncio.Future[None] = asyncio.get_running_loop().create_future() + self._waiters.append(fut) + try: + await fut + except asyncio.CancelledError: + # PERMIT-LEAK FIX (V5.1): + # Three states are possible at this point: + # (a) fut not done: nobody granted us a permit yet. Just + # remove from the queue. + # (b) fut done with result: release()/resize() handed us a + # permit (direct hand-off — no _current_value increment + # happened). We are about to NOT enter the critical + # section, so we MUST give the permit back. Call + # release() to wake the next waiter (or restore the + # permit to the pool if no waiters remain). + # (c) fut already cancelled before we entered the await: + # same shape as (a) — `fut in self._waiters` is True + # and `self._waiters.remove(fut)` covers it. No permit + # was granted, so nothing to give back. + if fut in self._waiters: + # Case (a) or (c): pre-grant cancellation. No permit + # was handed off, so nothing to release. + self._waiters.remove(fut) + elif fut.done() and not fut.cancelled() and fut.exception() is None: + # Case (b): post-grant cancellation. release() handed us + # a permit via set_result(None) but we're not going to + # use it. Hand it back so the next waiter (or the pool) + # gets it. + self.release() + raise + # Granted via direct hand-off from release()/resize(). + # _current_value was NOT decremented (the permit transferred in + # flight from the previous holder), so we are already accounted + # for as "active". + + def release(self) -> None: + """Return a permit. Wakes one blocked waiter if any. + + Direct hand-off semantics: when waiters are queued, the permit + transfers from the releaser to the next non-cancelled waiter + without round-tripping through ``_current_value``. This avoids + a window where two concurrent callers could observe + ``_current_value > 0`` between waiter-pop and decrement. + + After a resize-down, when no waiters remain, released permits + that would push ``_current_value`` above ``max_value`` are + absorbed (clamped). This is correct: the permit belonged to + the old, larger max. + """ + # Try to hand the permit directly to the next non-cancelled + # waiter. The skip-loop drains cancelled waiters whose futures + # are already done() — protecting against the release/cancel + # race where a waiter is cancelled mid-iteration. + while self._waiters: + # pyre-ignore[1001]: Future is granted via set_result(), not awaited. + waiter = self._waiters.popleft() + if not waiter.done(): + # Permit transfers atomically: stays "in flight" with + # the new owner. Do NOT touch _current_value. + waiter.set_result(None) + return + # No live waiters; restore one permit to the pool (clamped to + # the current max so resize-down clamps don't drift over). + self._current_value = min(self._current_value + 1, self._max_value) + + def resize(self, new_max: int) -> None: + """Change the maximum number of permits. + + Args: + new_max: New maximum. Must be >= 1. + + When increasing (``new_max > old_max``): + Additional permits become immediately available. Blocked + waiters are woken (via direct hand-off) to fill the new + capacity. Any leftover permits go to the pool, clamped to + the new max. + + When decreasing (``new_max < old_max``): + No preemption -- currently active tasks continue. + ``_current_value`` is reduced by the delta, which may make + it negative. Future ``acquire()`` calls block until enough + releases bring ``_current_value`` back above 0. + + Raises: + ValueError: If new_max < 1. + """ + if new_max < 1: + raise ValueError(f"new_max must be >= 1, got {new_max}") + delta = new_max - self._max_value + self._max_value = new_max + if delta > 0: + # Increase: hand `delta` permits directly to waiters first. + # Skip cancelled waiters (their futures are already done()). + granted = 0 + while self._waiters and granted < delta: + # pyre-ignore[1001]: Future is granted via set_result(), not awaited. + waiter = self._waiters.popleft() + if not waiter.done(): + # Direct hand-off: permit transfers in flight; do + # NOT touch _current_value here. + waiter.set_result(None) + granted += 1 + # Any permits not handed off go to the pool, clamped to max. + leftover = delta - granted + if leftover > 0: + self._current_value = min( + self._current_value + leftover, + self._max_value, + ) + elif delta < 0: + # Decrease: subtract from available pool (may go negative). + # delta is already negative. + self._current_value += delta diff --git a/src/spdl/pipeline/_scheduler.py b/src/spdl/pipeline/_scheduler.py new file mode 100644 index 000000000..52fae4d8c --- /dev/null +++ b/src/spdl/pipeline/_scheduler.py @@ -0,0 +1,450 @@ +# 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 + +"""Priority-dispatch scheduler for SPDL sync stages. + +When enabled via ``use_priority_scheduler=True`` on ``build_pipeline()``, +sync ops are routed through a priority queue instead of the underlying +``ThreadPoolExecutor``'s FIFO. Priority is ``-depth``: deeper stages +(closer to the sink) are dispatched first, draining the pipeline. + +Architecture (v5) +----------------- + +The hot path is a thin :py:class:`_PrioritizedExecutor` shim that +implements the :py:class:`concurrent.futures.Executor` ABC. Each sync +stage gets its own shim, captured as the stage's ``_PipeArgs.executor``. +The shim's :py:meth:`~_PrioritizedExecutor.submit` enqueues a +:py:class:`_WorkItem` onto a shared heap inside :py:class:`PriorityScheduler`, +which dispatches in priority order to a single underlying +:py:class:`~concurrent.futures.ThreadPoolExecutor`. + +The scheduler runs as a :py:class:`~spdl.pipeline._bg_task.BackgroundTask` +on the pipeline's event loop via the +:py:class:`_PrioritySchedulerBackgroundTask` adapter. + +Cancellation semantics +---------------------- + +``cf_future.set_running_or_notify_cancel()`` is called at *dispatch* +time inside :py:meth:`PriorityScheduler.run` (NOT at submit time). This +preserves the standard ``concurrent.futures.Future`` contract: pre-dispatch +``cancel()`` returns ``True`` and the dispatch loop skips the work item; +post-dispatch ``cancel()`` returns ``False`` and the work completes +normally (matching :py:class:`~concurrent.futures.ThreadPoolExecutor`). +""" + +from __future__ import annotations + +__all__ = [ + "PriorityScheduler", + "_PrioritizedExecutor", + "_PrioritySchedulerBackgroundTask", +] + +import asyncio +import concurrent.futures +import heapq +import logging +from concurrent.futures import Executor, ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Callable, cast + +from spdl.pipeline._bg_task import BackgroundTask +from spdl.pipeline._common._types import StageInfo + +_LG: logging.Logger = logging.getLogger(__name__) + + +@dataclass(order=True) +class _WorkItem: + """Unit of work ordered by ``(priority, seq)`` for heap dispatch. + + ``cf_future`` is the :py:class:`~concurrent.futures.Future` returned + from :py:meth:`_PrioritizedExecutor.submit`. ``bridge`` is a small + mutable dict that holds the dispatched :py:class:`asyncio.Future` + and a ``dispatched`` flag so that post-dispatch cancellation can + forward to the underlying task. + """ + + priority: int + seq: int + func: Callable[..., Any] = field(compare=False) + args: tuple[Any, ...] = field(compare=False) + kwargs: dict[str, Any] = field(compare=False) + cf_future: concurrent.futures.Future[Any] = field(compare=False) + bridge: dict[str, Any] = field(compare=False) + + +class PriorityScheduler: + """Priority-based dispatch scheduler for sync pipeline stages. + + The scheduler keeps a min-heap of :py:class:`_WorkItem`s. The + :py:meth:`run` coroutine pops items in ``(priority, seq)`` order and + dispatches each to the underlying :py:class:`ThreadPoolExecutor`, + bridging the result back to the caller's ``cf_future`` via + :py:meth:`asyncio.AbstractEventLoop.call_soon_threadsafe`. + + All scheduler bookkeeping mutations happen on the event-loop thread. + + Args: + max_concurrent: Maximum simultaneous dispatches to the underlying + pool (typically ``== num_threads``). + + Note: + The underlying :py:class:`ThreadPoolExecutor` is bound by + :py:func:`spdl.pipeline._build._build_pipeline` via direct + attribute assignment to :py:attr:`_underlying_executor` after + construction. The pool is not known at scheduler construction + time because it is created inside ``_build_pipeline`` from + ``num_threads``. + """ + + def __init__( + self, + max_concurrent: int, + ) -> None: + if max_concurrent < 1: + raise ValueError(f"max_concurrent must be >= 1, got {max_concurrent}") + + self._max_concurrent = max_concurrent + + # Bound by _build_pipeline() after construction. + self._underlying_executor: ThreadPoolExecutor | None = None + + # Per-stage priority lookup, keyed by stage name (StageInfo.stage_name). + # Populated by register_stage() at build time. + self._priorities: dict[str, int] = {} + + # Min-heap of (_WorkItem) entries; ordered by (priority, seq). + self._heap: list[_WorkItem] = [] + + # Monotonic sequence number for FIFO tie-breaking. Mutated only on + # the event-loop thread (via _enqueue from submit()). + self._seq: int = 0 + + # asyncio primitives created lazily inside run() so that the + # scheduler can be constructed before the event loop exists. + self._has_work: asyncio.Event | None = None + self._dispatch_semaphore: asyncio.Semaphore | None = None + + @property + def max_concurrent(self) -> int: + """Maximum simultaneous dispatches (read-only).""" + return self._max_concurrent + + def register_stage(self, info: StageInfo, priority: int) -> None: + """Register a stage's priority for later dispatch lookup. + + Called during ``_build_node()`` for each sync stage that should + be routed through the scheduler. + + Args: + info: The stage's :py:class:`StageInfo` (stage_name is used + as the registry key). + priority: Priority value (lower = higher priority; typically + ``-depth`` so deeper stages dispatch first). + """ + self._priorities[info.stage_name] = priority + _LG.debug( + "PriorityScheduler: registered stage %s with priority=%d", + info.stage_name, + priority, + ) + + def get_priority(self, info: StageInfo) -> int: + """Look up a stage's registered priority. Defaults to 0.""" + return self._priorities.get(info.stage_name, 0) + + def make_stage_executor(self, info: StageInfo) -> "_PrioritizedExecutor": + """Construct a per-stage :py:class:`_PrioritizedExecutor` shim. + + Called by ``_build_node`` in :py:mod:`spdl.pipeline._components._node` + when wiring a sync stage. Routing the construction through this + method keeps ``_node.py`` decoupled from the concrete + ``_PrioritizedExecutor`` class, avoiding a Buck dep cycle + between ``spdl/pipeline/_components`` and ``spdl/pipeline``. + """ + return _PrioritizedExecutor(scheduler=self, info=info) + + # ------------------------------------------------------------------ + # Internal: called from _PrioritizedExecutor on the event-loop thread + # ------------------------------------------------------------------ + + def _make_work_item( + self, + *, + priority: int, + func: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + cf_future: concurrent.futures.Future[Any], + bridge: dict[str, Any], + ) -> _WorkItem: + """Construct a ``_WorkItem``. Caller is on the event-loop thread.""" + self._seq += 1 + return _WorkItem( + priority=priority, + seq=self._seq, + func=func, + args=args, + kwargs=kwargs, + cf_future=cf_future, + bridge=bridge, + ) + + def _enqueue(self, work_item: _WorkItem) -> None: + """Push a ``_WorkItem`` onto the heap. On the event-loop thread.""" + heapq.heappush(self._heap, work_item) + if self._has_work is not None: + self._has_work.set() + + def _cancel_bridge(self, bridge: dict[str, Any]) -> None: + """Forward post-dispatch cancellation to the dispatched task. + + Pre-dispatch cancellation is handled by the dispatch loop's + :py:meth:`~concurrent.futures.Future.set_running_or_notify_cancel` + check (returns ``False`` for cancelled futures, skipping dispatch). + Post-dispatch, the underlying :py:class:`ThreadPoolExecutor` work + cannot be interrupted (per Python docs); we cancel the wrapping + :py:class:`asyncio.Future` so any awaiters see ``CancelledError``. + """ + if not bridge["dispatched"]: + # Pre-dispatch: the dispatch loop will skip this item via + # set_running_or_notify_cancel(). Nothing to do here. + return + + task = bridge.get("task") + if task is not None and not task.done(): + task.cancel() + + # ------------------------------------------------------------------ + # Dispatch loop (runs as a BackgroundTask) + # ------------------------------------------------------------------ + + async def run(self) -> None: + """Dispatch loop. Pops work items in priority order and runs them. + + Runs as a :py:class:`~spdl.pipeline._bg_task.BackgroundTask` + coroutine on the pipeline's event loop. Exits via + :py:exc:`asyncio.CancelledError` when the pipeline shuts down. + """ + if self._underlying_executor is None: + raise RuntimeError( + "PriorityScheduler._underlying_executor must be bound by " + "_build_pipeline() before run() is invoked." + ) + + # Lazy init of asyncio primitives now that we have a running loop. + has_work = asyncio.Event() + self._has_work = has_work + self._dispatch_semaphore = asyncio.Semaphore(self._max_concurrent) + sem = self._dispatch_semaphore + loop = asyncio.get_running_loop() + + # If items were enqueued via _PrioritizedExecutor.submit() BEFORE + # run() was scheduled (e.g., stage tasks fire before the BG task + # in _run_pipeline_coroutines), they already sit on the heap but + # _has_work could not have been set (it didn't exist yet). Seed + # the event so the dispatch loop notices them on the first turn. + if self._heap: + has_work.set() + + while True: + await has_work.wait() + while self._heap: + await sem.acquire() + work_item = heapq.heappop(self._heap) + + cf = work_item.cf_future + # V5.2: PENDING -> RUNNING here. If pre-dispatch cancel + # already moved cf to CANCELLED, skip dispatch entirely. + if not cf.set_running_or_notify_cancel(): + sem.release() + continue + + # Dispatch onto the underlying ThreadPoolExecutor. + # loop.run_in_executor() schedules the work on the event + # loop and returns its scheduling Future immediately; we + # observe completion via add_done_callback rather than + # awaiting directly. The cast() to object discharges + # pyre's awaitable-tracking once the callback is wired. + fut: asyncio.Future[Any] = loop.run_in_executor( + self._underlying_executor, + self._invoke, + work_item, + ) + fut.add_done_callback( + lambda f, wi=work_item, s=sem: self._on_complete(f, wi, s) + ) + work_item.bridge["task"] = cast(object, fut) + work_item.bridge["dispatched"] = True + has_work.clear() + + @staticmethod + def _invoke(work_item: _WorkItem) -> Any: + """Worker-thread entrypoint. Runs the user's function.""" + return work_item.func(*work_item.args, **work_item.kwargs) + + def _on_complete( + self, + fut: asyncio.Future[Any], + work_item: _WorkItem, + sem: asyncio.Semaphore, + ) -> None: + """Bridge completion back to the caller's ``cf_future``. + + Runs on the event-loop thread (asyncio future done-callback). + """ + sem.release() + cf = work_item.cf_future + # cf is in RUNNING at this point (set at dispatch). Final state + # is FINISHED (set_result/set_exception). asyncio.CancelledError + # raised by the underlying task is surfaced as an exception on cf. + if fut.cancelled(): + cf.set_exception(asyncio.CancelledError()) + return + exc = fut.exception() + if exc is not None: + cf.set_exception(exc) + else: + cf.set_result(fut.result()) + + def _drain_pending(self) -> None: + """Cancel everything still on the heap on shutdown. + + Called from :py:class:`_PrioritySchedulerBackgroundTask` after + the dispatch loop exits. Items still on the heap have + ``cf_future`` in PENDING state; cancelling them lets any awaiting + :py:meth:`asyncio.AbstractEventLoop.run_in_executor` callers see + :py:exc:`asyncio.CancelledError` cleanly. + """ + while self._heap: + work_item = heapq.heappop(self._heap) + cf = work_item.cf_future + cf.cancel() + + +class _PrioritizedExecutor(Executor): + """Thin :py:class:`Executor` adapter that submits work through + :py:class:`PriorityScheduler`. + + One instance per sync stage — captures the stage's + :py:class:`StageInfo` so the scheduler can look up the right priority. + + NOT a real :py:class:`ThreadPoolExecutor` — it has no threads of its + own. :py:meth:`submit` enqueues into the scheduler's heap, which + dispatches to the real shared pool when capacity is available. + + Threading model (V5.3): :py:meth:`submit` is called by + :py:meth:`asyncio.AbstractEventLoop.run_in_executor` synchronously on + the pipeline's event-loop thread. We therefore can fetch the loop + via :py:func:`asyncio.get_running_loop` *per-call* — no pre-stamping + is required, which avoids races during pipeline startup before the + BG task has begun. + + Cancellation: post-dispatch ``cf_future.cancel()`` returns ``False`` + (matching :py:class:`ThreadPoolExecutor` semantics). Pre-dispatch + cancel works correctly via + :py:meth:`~concurrent.futures.Future.set_running_or_notify_cancel` + being called at dispatch time. + """ + + def __init__( + self, + scheduler: PriorityScheduler, + info: StageInfo, + ) -> None: + self._scheduler = scheduler + self._info = info + + # pyre-ignore[14]: Executor.submit's typeshed stub uses an internal + # TypeVar that cannot be matched from a subclass override; pyre flags + # any concrete signature here as Inconsistent override even when the + # runtime semantics are identical to the base class. This is the same + # workaround used by other Executor subclasses (the alternative is to + # restructure the entire dispatch API). + def submit( + self, + fn: Callable[..., Any], + /, + *args: Any, + **kwargs: Any, + ) -> concurrent.futures.Future[Any]: + # V5.3: per-call loop fetch. Works because loop.run_in_executor() + # invokes executor.submit() synchronously on the loop thread. + loop: asyncio.AbstractEventLoop = asyncio.get_running_loop() + + cf_future: concurrent.futures.Future[Any] = concurrent.futures.Future() + # cf_future is in PENDING state. set_running_or_notify_cancel() is + # called at DISPATCH time (in PriorityScheduler.run()), NOT here. + + bridge: dict[str, Any] = { + "task": None, # asyncio.Future from run_in_executor; set after dispatch + "dispatched": False, # True once the asyncio task starts + } + + def _enqueue() -> None: + # Runs on the loop thread. + work_item = self._scheduler._make_work_item( + priority=self._scheduler.get_priority(self._info), + func=fn, + args=args, + kwargs=kwargs, + cf_future=cf_future, + bridge=bridge, + ) + self._scheduler._enqueue(work_item) + + # call_soon_threadsafe is correct even when caller is the loop + # thread (defensive; also keeps submit() safe from non-loop threads + # such as unit tests). + loop.call_soon_threadsafe(_enqueue) + + def _on_cancel(cf_fut: concurrent.futures.Future[Any]) -> None: + # Called by concurrent.futures when cf_future transitions to + # a done state. We only act on the cancellation case. + if not cf_fut.cancelled(): + return + loop.call_soon_threadsafe(self._scheduler._cancel_bridge, bridge) + + cf_future.add_done_callback(_on_cancel) + return cf_future + + def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None: + """No-op: the underlying scheduler owns the worker pool lifecycle.""" + # The shared ThreadPoolExecutor is owned by the Pipeline, not by + # this shim. Shutdown happens via Pipeline.stop() and the BG task + # cancellation, which triggers _drain_pending() on the scheduler. + return None + + +class _PrioritySchedulerBackgroundTask(BackgroundTask): + """:py:class:`BackgroundTask` adapter that runs + :py:meth:`PriorityScheduler.run` on the pipeline's event loop. + + Lifecycle (from :py:class:`spdl.pipeline._bg_task.BackgroundTask`): + + - Started by ``_run_pipeline_coroutines()`` AFTER stage tasks are + created. + - Cancelled when ALL stage tasks complete (or on pipeline failure). + - Errors are logged, do NOT fail the pipeline. + + On cancellation, ensures any items left on the heap are cancelled so + that awaiting :py:meth:`asyncio.AbstractEventLoop.run_in_executor` + callers don't hang. + """ + + def __init__(self, scheduler: PriorityScheduler) -> None: + self._scheduler = scheduler + + async def run(self) -> None: + try: + await self._scheduler.run() + finally: + self._scheduler._drain_pending() diff --git a/tests/pipeline/resizable_semaphore_test.py b/tests/pipeline/resizable_semaphore_test.py new file mode 100644 index 000000000..341eeb4f0 --- /dev/null +++ b/tests/pipeline/resizable_semaphore_test.py @@ -0,0 +1,840 @@ +# 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 asyncio +import unittest + +from spdl.pipeline._components._semaphore import ResizableSemaphore + + +class ResizableSemaphoreConstructionTest(unittest.TestCase): + def test_init_valid(self) -> None: + sem = ResizableSemaphore(5) + self.assertEqual(sem.max_value, 5) + self.assertEqual(sem.active, 0) + + def test_init_one(self) -> None: + sem = ResizableSemaphore(1) + self.assertEqual(sem.max_value, 1) + self.assertEqual(sem.active, 0) + + def test_init_zero_raises(self) -> None: + with self.assertRaises(ValueError): + ResizableSemaphore(0) + + def test_init_negative_raises(self) -> None: + with self.assertRaises(ValueError): + ResizableSemaphore(-1) + + +class ResizableSemaphoreAcquireReleaseTest(unittest.TestCase): + def test_acquire_decrements_permits(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(3) + await sem.acquire() + self.assertEqual(sem.active, 1) + await sem.acquire() + self.assertEqual(sem.active, 2) + await sem.acquire() + self.assertEqual(sem.active, 3) + + asyncio.run(run()) + + def test_release_increments_permits(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(3) + await sem.acquire() + await sem.acquire() + self.assertEqual(sem.active, 2) + sem.release() + self.assertEqual(sem.active, 1) + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_acquire_blocks_when_exhausted(self) -> None: + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + acquired: asyncio.Event = asyncio.Event() + + async def try_acquire() -> None: + await sem.acquire() + acquired.set() + + task = asyncio.create_task(try_acquire()) + # Yield to let the task enter acquire() and block. + await asyncio.sleep(0) + self.assertFalse(acquired.is_set()) + + # Release unblocks the waiter. + sem.release() + await asyncio.sleep(0) + self.assertTrue(acquired.is_set()) + + # Clean up. + sem.release() + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(run()) + + def test_release_without_acquire_clamps(self) -> None: + """Release without prior acquire should not exceed max_value.""" + + async def run() -> None: + sem = ResizableSemaphore(3) + # active is 0, current_value == max_value == 3 + sem.release() + # Should clamp: active stays 0, not -1. + self.assertEqual(sem.active, 0) + self.assertEqual(sem.max_value, 3) + + asyncio.run(run()) + + def test_fifo_wake_order(self) -> None: + """Waiters should be woken in FIFO order.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + order: list[int] = [] + + async def waiter(idx: int) -> None: + await sem.acquire() + order.append(idx) + sem.release() + + t1 = asyncio.create_task(waiter(1)) + await asyncio.sleep(0) + t2 = asyncio.create_task(waiter(2)) + await asyncio.sleep(0) + t3 = asyncio.create_task(waiter(3)) + await asyncio.sleep(0) + + # Release the initial acquire — should wake waiter 1 first. + sem.release() + + # Wait for all waiters to complete. + await asyncio.wait_for(asyncio.gather(t1, t2, t3), timeout=5.0) + self.assertEqual(order, [1, 2, 3]) + + asyncio.run(run()) + + +class ResizableSemaphoreResizeUpTest(unittest.TestCase): + def test_resize_up_increases_max(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(2) + sem.resize(5) + self.assertEqual(sem.max_value, 5) + + asyncio.run(run()) + + def test_resize_up_wakes_waiters(self) -> None: + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + acquired_events: list[asyncio.Event] = [] + + async def waiter() -> None: + await sem.acquire() + evt = acquired_events[len(acquired_events)] + evt.set() + + e1 = asyncio.Event() + e2 = asyncio.Event() + acquired_events.extend([e1, e2]) + + # Simpler: track via counter + acquired_count = 0 + + async def counting_waiter() -> None: + nonlocal acquired_count + await sem.acquire() + acquired_count += 1 + + t1 = asyncio.create_task(counting_waiter()) + await asyncio.sleep(0) + t2 = asyncio.create_task(counting_waiter()) + await asyncio.sleep(0) + self.assertEqual(acquired_count, 0) + + # Resize from 1 -> 3: adds 2 permits, should wake both waiters. + sem.resize(3) + await asyncio.sleep(0) + self.assertEqual(acquired_count, 2) + self.assertEqual(sem.active, 3) + + # Clean up — release all 3 acquired permits. + sem.release() + sem.release() + sem.release() + + # Await tasks to prevent warnings. + await asyncio.wait_for(asyncio.gather(t1, t2), timeout=5.0) + + asyncio.run(run()) + + def test_resize_up_partial_wake(self) -> None: + """When resize adds fewer permits than waiters, only some wake.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + acquired: list[int] = [] + + async def waiter(idx: int) -> None: + await sem.acquire() + acquired.append(idx) + + t1 = asyncio.create_task(waiter(1)) + await asyncio.sleep(0) + t2 = asyncio.create_task(waiter(2)) + await asyncio.sleep(0) + t3 = asyncio.create_task(waiter(3)) + await asyncio.sleep(0) + + # Resize from 1 -> 2: adds 1 permit, wakes 1 waiter. + sem.resize(2) + await asyncio.sleep(0) + self.assertEqual(len(acquired), 1) + self.assertEqual(acquired[0], 1) # FIFO + + # Clean up remaining waiters. + for t in (t1, t2, t3): + t.cancel() + try: + await t + except asyncio.CancelledError: + pass + + asyncio.run(run()) + + +class ResizableSemaphoreResizeDownTest(unittest.TestCase): + def test_resize_down_no_preemption(self) -> None: + """Active tasks continue after resize down.""" + + async def run() -> None: + sem = ResizableSemaphore(3) + await sem.acquire() + await sem.acquire() + await sem.acquire() + self.assertEqual(sem.active, 3) + + # Resize down to 1. All 3 are still active. + sem.resize(1) + self.assertEqual(sem.max_value, 1) + self.assertEqual(sem.active, 3) # No preemption. + + asyncio.run(run()) + + def test_resize_down_blocks_new_acquires(self) -> None: + """After resize down, new acquires block until drain completes.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(3) + await sem.acquire() + await sem.acquire() + await sem.acquire() + + sem.resize(1) + + acquired: asyncio.Event = asyncio.Event() + + async def try_acquire() -> None: + await sem.acquire() + acquired.set() + + task = asyncio.create_task(try_acquire()) + await asyncio.sleep(0) + self.assertFalse(acquired.is_set()) + + # Release 3 permits (drain from 3 active -> 0). + # First two releases bring active from 3->2->1 (at max). + # Third release frees a permit for the waiter. + sem.release() + sem.release() + sem.release() + await asyncio.sleep(0) + self.assertTrue(acquired.is_set()) + + # Clean up. + sem.release() + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(run()) + + def test_resize_down_drain_releases_clamp(self) -> None: + """Releases during drain clamp to max_value, not old max.""" + + async def run() -> None: + sem = ResizableSemaphore(5) + await sem.acquire() + await sem.acquire() + # active=2, current_value=3 + sem.resize(2) + # current_value should now be 0 (3 - (5-2) = 0) + self.assertEqual(sem.active, 2) + self.assertEqual(sem.max_value, 2) + + # Release one: active -> 1. + sem.release() + self.assertEqual(sem.active, 1) + + # Release another: active -> 0. + sem.release() + self.assertEqual(sem.active, 0) + + # Extra release should clamp. + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + +class ResizableSemaphoreResizeEdgeCasesTest(unittest.TestCase): + def test_resize_to_same_value(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(3) + await sem.acquire() + sem.resize(3) + self.assertEqual(sem.max_value, 3) + self.assertEqual(sem.active, 1) + + asyncio.run(run()) + + def test_resize_to_zero_raises(self) -> None: + sem = ResizableSemaphore(3) + with self.assertRaises(ValueError): + sem.resize(0) + + def test_resize_to_negative_raises(self) -> None: + sem = ResizableSemaphore(3) + with self.assertRaises(ValueError): + sem.resize(-5) + + def test_resize_preserves_max_after_error(self) -> None: + """Failed resize should not change max_value.""" + sem = ResizableSemaphore(3) + try: + sem.resize(0) + except ValueError: + pass + self.assertEqual(sem.max_value, 3) + + def test_resize_while_waiters_pending_noop(self) -> None: + """Resize to same value while tasks are waiting should not wake them.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + acquired: asyncio.Event = asyncio.Event() + + async def waiter() -> None: + await sem.acquire() + acquired.set() + + task = asyncio.create_task(waiter()) + await asyncio.sleep(0) + + # Resize to same value — waiter should stay blocked. + sem.resize(1) + await asyncio.sleep(0) + self.assertFalse(acquired.is_set()) + + # Clean up. + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(run()) + + def test_multiple_resizes(self) -> None: + """Multiple sequential resizes should work correctly.""" + + async def run() -> None: + sem = ResizableSemaphore(1) + sem.resize(5) + self.assertEqual(sem.max_value, 5) + sem.resize(2) + self.assertEqual(sem.max_value, 2) + sem.resize(10) + self.assertEqual(sem.max_value, 10) + + asyncio.run(run()) + + +class ResizableSemaphoreCancellationTest(unittest.TestCase): + def test_cancelled_waiter_is_removed(self) -> None: + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + async def waiter() -> None: + await sem.acquire() + + task = asyncio.create_task(waiter()) + await asyncio.sleep(0) + + # Cancel the waiter. + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Release — should not raise even though waiter was cancelled. + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_cancel_one_of_multiple_waiters(self) -> None: + """Cancel middle waiter; remaining waiters still get served.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + order: list[int] = [] + + async def waiter(idx: int) -> None: + await sem.acquire() + order.append(idx) + sem.release() + + t1 = asyncio.create_task(waiter(1)) + await asyncio.sleep(0) + t2 = asyncio.create_task(waiter(2)) + await asyncio.sleep(0) + t3 = asyncio.create_task(waiter(3)) + await asyncio.sleep(0) + + # Cancel the second waiter. + t2.cancel() + try: + await t2 + except asyncio.CancelledError: + pass + + sem.release() + await asyncio.wait_for(asyncio.gather(t1, t3), timeout=5.0) + self.assertEqual(order, [1, 3]) + + asyncio.run(run()) + + +class ResizableSemaphorePermitLeakTest(unittest.TestCase): + """V5.1 (DESIGN.md): regression coverage for the permit-leak fix in + ``acquire()``'s cancellation handler. + + Without the fix in case (b) — when ``release()`` direct-hands a + permit via ``set_result(None)`` and the waiter is then cancelled + before resuming — the permit was lost. Each cancellation under + contention leaked one permit; eventually the semaphore deadlocked. + """ + + def test_acquire_cancel_after_grant_releases_permit(self) -> None: + """Case (b): waiter is granted a permit via direct hand-off but + cancelled before resuming. The permit must be released back so + another acquirer can make progress. Without the fix, the second + acquirer would deadlock. + """ + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() # holder takes the only permit. + + # First waiter: will be granted the permit by release(), + # then cancelled before its acquire() returns. + granted_w1: asyncio.Event = asyncio.Event() + + async def w1() -> None: + # Acquire will receive set_result(None) from the + # holder's release(); cancellation arrives during + # the await window before acquire() returns. + try: + await sem.acquire() + except asyncio.CancelledError: + granted_w1.set() + raise + + t1 = asyncio.create_task(w1()) + # Yield until W1 has enqueued its waiter future. + await asyncio.sleep(0) + self.assertEqual(len(sem._waiters), 1) + + # Direct-hand the permit to W1: this calls + # waiter.set_result(None) on W1's future. W1 is now in + # case (b): future is done with a result, but W1 hasn't + # resumed yet. + sem.release() + self.assertEqual(len(sem._waiters), 0) + + # Cancel W1 before its acquire() resumes. Permit-leak fix + # must detect case (b) and call self.release() to give + # the permit back. + t1.cancel() + try: + await t1 + except asyncio.CancelledError: + pass + self.assertTrue(granted_w1.is_set()) + + # Now a new acquirer should be able to acquire the permit + # immediately. WITHOUT THE FIX this would deadlock here. + second_acquired: asyncio.Event = asyncio.Event() + + async def w2() -> None: + await sem.acquire() + second_acquired.set() + + t2 = asyncio.create_task(w2()) + await asyncio.wait_for(t2, timeout=2.0) + self.assertTrue(second_acquired.is_set()) + + # Pool accounting: we have one active permit (W2's). + self.assertEqual(sem.active, 1) + + asyncio.run(run()) + + def test_release_grants_to_next_when_first_waiter_cancelled( + self, + ) -> None: + """Direct-hand-off race: ``release()`` pops a waiter whose + future is already cancelled (done()) and must skip it to hand + the permit to the next non-cancelled waiter. + + Sequence: + 1. sem with value=1, acquired by holder. + 2. Two waiters W1, W2 enqueued. + 3. W1.cancel() — its future becomes done() (cancelled). + 4. holder.release() — must pop W1 (skip), pop W2 (grant). + """ + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() # holder + + w1_started: asyncio.Event = asyncio.Event() + w2_started: asyncio.Event = asyncio.Event() + w2_acquired: asyncio.Event = asyncio.Event() + + async def w1() -> None: + w1_started.set() + await sem.acquire() + + async def w2() -> None: + w2_started.set() + await sem.acquire() + w2_acquired.set() + + t1 = asyncio.create_task(w1()) + await w1_started.wait() + t2 = asyncio.create_task(w2()) + await w2_started.wait() + # Force both to enqueue their waiter futures. + await asyncio.sleep(0) + await asyncio.sleep(0) + self.assertEqual(len(sem._waiters), 2) + + # Cancel W1 — its waiter future transitions to done() but + # remains in the deque (acquire()'s except handler removes + # it after the await re-raises CancelledError). + t1.cancel() + try: + await t1 + except asyncio.CancelledError: + pass + + # release() must pop W1 (skip — done()) then pop W2 (grant). + sem.release() + await asyncio.wait_for(t2, timeout=2.0) + self.assertTrue(w2_acquired.is_set()) + self.assertFalse(t2.cancelled()) + + asyncio.run(run()) + + def test_acquire_pre_cancelled_future_fast_path(self) -> None: + """Case (c) documentation: when the future has been cancelled + before the await even completes (e.g., the task was cancelled + before it ran past the ``self._waiters.append(fut)`` line), the + cancellation handler treats this the same as case (a) — the + future is still in ``self._waiters``, so ``self._waiters.remove + (fut)`` covers it. No permit was granted, so nothing to release + back. + + This test verifies the case (c) fast path via cancelling at + scheduling time rather than mid-await: pool accounting must + remain consistent (active=1 from the holder, max=1). + """ + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() # holder + + async def w1() -> None: + # Will be cancelled before release() ever fires, while + # waiter future is still pending in the deque. + await sem.acquire() + + t1 = asyncio.create_task(w1()) + # Let W1 enqueue its waiter future. + await asyncio.sleep(0) + self.assertEqual(len(sem._waiters), 1) + + # Cancel W1 immediately — future is still pending (case (a)), + # which shares the cleanup path with case (c). + t1.cancel() + try: + await t1 + except asyncio.CancelledError: + pass + + # The waiter future must be cleaned up from the deque so + # subsequent release() doesn't try to grant to a dead + # waiter. + self.assertEqual(len(sem._waiters), 0) + + # Pool accounting: holder still owns the permit. + self.assertEqual(sem.active, 1) + self.assertEqual(sem.max_value, 1) + + # Release returns the permit cleanly. + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + +class ResizableSemaphoreConcurrencyTest(unittest.TestCase): + def test_concurrent_acquire_release(self) -> None: + """Many tasks acquiring and releasing concurrently.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(3) + completed = 0 + + async def worker() -> None: + nonlocal completed + await sem.acquire() + # Yield to let other tasks proceed. + await asyncio.sleep(0) + sem.release() + completed += 1 + + tasks = [asyncio.create_task(worker()) for _ in range(20)] + await asyncio.wait_for(asyncio.gather(*tasks), timeout=5.0) + self.assertEqual(completed, 20) + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_concurrent_acquire_with_resize(self) -> None: + """Resize during concurrent acquire/release operations.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(2) + completed = 0 + + async def worker() -> None: + nonlocal completed + await sem.acquire() + await asyncio.sleep(0) + sem.release() + completed += 1 + + tasks = [asyncio.create_task(worker()) for _ in range(10)] + + # Yield a few times to let some workers start. + await asyncio.sleep(0) + await asyncio.sleep(0) + sem.resize(5) # Expand. + await asyncio.sleep(0) + sem.resize(1) # Contract. + + await asyncio.wait_for(asyncio.gather(*tasks), timeout=5.0) + self.assertEqual(completed, 10) + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_active_never_exceeds_max_under_load(self) -> None: + """active should never exceed max_value when max is stable.""" + + async def run() -> None: + max_permits = 4 + sem: ResizableSemaphore = ResizableSemaphore(max_permits) + max_seen = 0 + + async def worker() -> None: + nonlocal max_seen + await sem.acquire() + current = sem.active + if current > max_seen: + max_seen = current + await asyncio.sleep(0) + sem.release() + + tasks = [asyncio.create_task(worker()) for _ in range(30)] + await asyncio.wait_for(asyncio.gather(*tasks), timeout=5.0) + self.assertLessEqual(max_seen, max_permits) + + asyncio.run(run()) + + +class ResizableSemaphoreResizeReleaseInteractionTest(unittest.TestCase): + """Interaction between ``resize()`` and ``release()`` direct hand-off. + + These tests close gaps in the existing suite around the boundary + between resize-down accounting and release-direct-handoff accounting. + """ + + def test_resize_down_then_release_with_waiter_caps_at_new_max(self) -> None: + """resize-down → release direct-handoff to a waiter must not exceed new max.""" + + async def run() -> None: + # Arrange + sem: ResizableSemaphore = ResizableSemaphore(3) + await sem.acquire() + await sem.acquire() + await sem.acquire() + + acquired_after_resize: list[int] = [] + + async def waiter(idx: int) -> None: + await sem.acquire() + acquired_after_resize.append(idx) + + # Enqueue a waiter while the pool is exhausted. + t = asyncio.create_task(waiter(1)) + await asyncio.sleep(0) + + # Resize down: max -> 1; current_value goes to -2. + sem.resize(1) + + # Act: three releases drain the over-fill, then one more + # crosses the threshold and direct-hands to the waiter. + sem.release() + sem.release() + sem.release() + await asyncio.sleep(0) + self.assertEqual(acquired_after_resize, [1]) + + # Assert: with the waiter holding the only permit (active=1) + # and max_value=1, sem must report active==max_value. + self.assertEqual(sem.active, 1) + self.assertEqual(sem.max_value, 1) + + # Cleanup + sem.release() + await asyncio.wait_for(t, timeout=2.0) + # Pool returned to "no active permits" without exceeding max. + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_concurrent_resize_calls_converge_to_last_value(self) -> None: + """Two coroutines calling resize() concurrently — the last wins. + + asyncio is single-threaded but multiple coroutines can issue + resize() in the same tick. The final ``max_value`` must equal + the last resize call's value, and accounting must remain + consistent. + """ + + async def run() -> None: + # Arrange + sem: ResizableSemaphore = ResizableSemaphore(2) + await sem.acquire() + + async def resizer(target: int) -> None: + # Yield once to interleave the two resize calls. + await asyncio.sleep(0) + sem.resize(target) + + # Act: dispatch two concurrent resizes. + await asyncio.gather(resizer(5), resizer(3)) + + # Assert: final value is whichever resizer ran last. Both are + # valid (>= 1); accounting must be consistent. + self.assertIn(sem.max_value, (3, 5)) + # active == max_value - current_value; the holder still owns + # one permit, so active is 1 regardless of resize order. + self.assertEqual(sem.active, 1) + + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + +class ResizableSemaphorePropertiesTest(unittest.TestCase): + def test_max_value_reflects_resize(self) -> None: + sem = ResizableSemaphore(3) + self.assertEqual(sem.max_value, 3) + sem.resize(7) + self.assertEqual(sem.max_value, 7) + sem.resize(1) + self.assertEqual(sem.max_value, 1) + + def test_active_reflects_acquire_release(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(5) + self.assertEqual(sem.active, 0) + await sem.acquire() + self.assertEqual(sem.active, 1) + await sem.acquire() + self.assertEqual(sem.active, 2) + sem.release() + self.assertEqual(sem.active, 1) + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_active_exceeds_max_after_resize_down(self) -> None: + """active can temporarily exceed max_value after resize down.""" + + async def run() -> None: + sem = ResizableSemaphore(5) + for _ in range(5): + await sem.acquire() + self.assertEqual(sem.active, 5) + + sem.resize(2) + self.assertEqual(sem.max_value, 2) + self.assertEqual(sem.active, 5) # Exceeds max_value. + + # Drain. + for _ in range(5): + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pipeline/scheduler_test.py b/tests/pipeline/scheduler_test.py new file mode 100644 index 000000000..5c0aa5ec6 --- /dev/null +++ b/tests/pipeline/scheduler_test.py @@ -0,0 +1,647 @@ +# 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 + +"""Unit tests for ``PriorityScheduler`` + ``_PrioritizedExecutor`` (v5).""" + +import asyncio +import concurrent.futures +import threading +import time +import unittest +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import later.unittest +from spdl.pipeline import PipelineBuilder +from spdl.pipeline._common._types import StageInfo +from spdl.pipeline._scheduler import ( + _PrioritizedExecutor, + _PrioritySchedulerBackgroundTask, + PriorityScheduler, +) + + +def _make_info(name: str, stage_id: str = "0", concurrency: int = 1) -> StageInfo: + return StageInfo( + pipeline_id=0, + stage_id=stage_id, + stage_name=name, + concurrency=concurrency, + ) + + +class PrioritySchedulerConstructionTest(unittest.TestCase): + """V5: PriorityScheduler is priority-only (no `adapt` flag).""" + + def test_max_concurrent_property(self) -> None: + scheduler = PriorityScheduler(max_concurrent=4) + self.assertEqual(scheduler.max_concurrent, 4) + + def test_max_concurrent_must_be_positive(self) -> None: + with self.assertRaises(ValueError): + PriorityScheduler(max_concurrent=0) + with self.assertRaises(ValueError): + PriorityScheduler(max_concurrent=-1) + + def test_underlying_executor_unbound_at_init(self) -> None: + # Bound by _build_pipeline() via direct attribute assignment. + scheduler = PriorityScheduler(max_concurrent=2) + self.assertIsNone(scheduler._underlying_executor) + + +class RegisterStageTest(unittest.TestCase): + """V5: register_stage takes (StageInfo, priority) only.""" + + def test_register_stage_stores_priority(self) -> None: + scheduler = PriorityScheduler(max_concurrent=4) + info = _make_info("decode") + scheduler.register_stage(info, priority=-3) + self.assertEqual(scheduler.get_priority(info), -3) + + def test_get_priority_default_zero_for_unregistered(self) -> None: + scheduler = PriorityScheduler(max_concurrent=4) + info = _make_info("never_registered") + self.assertEqual(scheduler.get_priority(info), 0) + + def test_register_overwrite(self) -> None: + scheduler = PriorityScheduler(max_concurrent=4) + info = _make_info("s") + scheduler.register_stage(info, priority=-1) + scheduler.register_stage(info, priority=-5) + self.assertEqual(scheduler.get_priority(info), -5) + + +class PrioritizedExecutorContractTest(unittest.TestCase): + """V5 m2: _PrioritizedExecutor must NOT be a ProcessPoolExecutor. + + convert_to_async branches on isinstance(executor, ProcessPoolExecutor) — + we want the default path so loop.run_in_executor() invokes our submit(). + """ + + def test_isinstance_is_not_process_pool(self) -> None: + scheduler = PriorityScheduler(max_concurrent=2) + info = _make_info("t") + shim = _PrioritizedExecutor(scheduler, info) + self.assertNotIsInstance(shim, concurrent.futures.ProcessPoolExecutor) + + def test_is_executor_subclass(self) -> None: + # Confirms the shim implements the Executor ABC so it can be + # passed to loop.run_in_executor. + scheduler = PriorityScheduler(max_concurrent=2) + info = _make_info("t") + shim = _PrioritizedExecutor(scheduler, info) + self.assertIsInstance(shim, concurrent.futures.Executor) + + def test_shutdown_is_noop(self) -> None: + scheduler = PriorityScheduler(max_concurrent=2) + shim = _PrioritizedExecutor(scheduler, _make_info("t")) + # Should not raise; the underlying pool is owned elsewhere. + shim.shutdown() + shim.shutdown(wait=False) + shim.shutdown(wait=True, cancel_futures=True) + + +class PrioritizedExecutorPerCallLoopFetchTest(later.unittest.TestCase): + """V5.3: stage tasks may submit() before the BG task runs. + + Per-call asyncio.get_running_loop() must not require pre-stamping. + """ + + async def test_submit_works_without_bg_task_running(self) -> None: + with ThreadPoolExecutor(max_workers=2) as pool: + scheduler = PriorityScheduler(max_concurrent=2) + scheduler._underlying_executor = pool + info = _make_info("t") + scheduler.register_stage(info, priority=0) + shim = _PrioritizedExecutor(scheduler, info) + + # submit() succeeds without scheduler.run() being invoked. + cf = shim.submit(lambda: 42) + self.assertIsInstance(cf, concurrent.futures.Future) + # The work item is enqueued via call_soon_threadsafe; let it + # land before checking heap state. + await asyncio.sleep(0) + self.assertEqual(len(scheduler._heap), 1) + + # Now start the dispatch loop in the background. + run_task = asyncio.create_task(scheduler.run()) + try: + result = await asyncio.wait_for(asyncio.wrap_future(cf), timeout=2.0) + self.assertEqual(result, 42) + finally: + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + +class CancelPreDispatchTest(later.unittest.TestCase): + """V5.2: pre-dispatch cancel returns True; dispatch loop skips item.""" + + async def test_cancel_before_dispatch_skips(self) -> None: + with ThreadPoolExecutor(max_workers=1) as pool: + scheduler = PriorityScheduler(max_concurrent=1) + scheduler._underlying_executor = pool + info = _make_info("t") + scheduler.register_stage(info, priority=0) + shim = _PrioritizedExecutor(scheduler, info) + + ran: threading.Event = threading.Event() + + def work() -> int: + ran.set() + return 1 + + cf = shim.submit(work) + # Cancel BEFORE the dispatch loop pops the item. + self.assertTrue(cf.cancel()) + self.assertTrue(cf.cancelled()) + + # Now run the dispatch loop briefly. It should pop the cancelled + # item, see set_running_or_notify_cancel returns False, and skip. + run_task = asyncio.create_task(scheduler.run()) + try: + await asyncio.sleep(0.1) + self.assertFalse(ran.is_set()) + finally: + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + +class CancelPostDispatchBestEffortTest(later.unittest.TestCase): + """V5.2: post-dispatch cancel returns False (matches stdlib semantics).""" + + async def test_cancel_after_dispatch_returns_false(self) -> None: + with ThreadPoolExecutor(max_workers=1) as pool: + scheduler = PriorityScheduler(max_concurrent=1) + scheduler._underlying_executor = pool + info = _make_info("t") + scheduler.register_stage(info, priority=0) + shim = _PrioritizedExecutor(scheduler, info) + + started: threading.Event = threading.Event() + release: threading.Event = threading.Event() + + def slow_work() -> int: + started.set() + release.wait(timeout=2.0) + return 99 + + run_task = asyncio.create_task(scheduler.run()) + try: + cf = shim.submit(slow_work) + # Wait for dispatch to begin. + await asyncio.get_running_loop().run_in_executor( + None, started.wait, 2.0 + ) + # cf is now RUNNING; cancel() should return False per + # ThreadPoolExecutor semantics. + self.assertFalse(cf.cancel()) + # Let the worker complete. + release.set() + result = await asyncio.wait_for(asyncio.wrap_future(cf), timeout=2.0) + self.assertEqual(result, 99) + finally: + release.set() + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + +class PriorityDispatchOrderTest(later.unittest.TestCase): + """Items dispatch in (priority, seq) order.""" + + async def test_lower_priority_value_dispatches_first(self) -> None: + # Use max_concurrent=1 so dispatch is strictly serialized. + with ThreadPoolExecutor(max_workers=1) as pool: + scheduler = PriorityScheduler(max_concurrent=1) + scheduler._underlying_executor = pool + info_a = _make_info("a") + info_b = _make_info("b") + info_c = _make_info("c") + scheduler.register_stage(info_a, priority=0) # lowest priority + scheduler.register_stage(info_b, priority=-2) # higher + scheduler.register_stage(info_c, priority=-5) # highest + + shim_a = _PrioritizedExecutor(scheduler, info_a) + shim_b = _PrioritizedExecutor(scheduler, info_b) + shim_c = _PrioritizedExecutor(scheduler, info_c) + + results: list[str] = [] + + def work(label: str) -> str: + results.append(label) + return label + + # Hold the dispatch loop until all 3 are enqueued. + cf_a = shim_a.submit(work, "a") + cf_b = shim_b.submit(work, "b") + cf_c = shim_c.submit(work, "c") + + # Yield so the call_soon_threadsafe enqueues land. + await asyncio.sleep(0) + self.assertEqual(len(scheduler._heap), 3) + + run_task = asyncio.create_task(scheduler.run()) + try: + await asyncio.gather( + asyncio.wrap_future(cf_a), + asyncio.wrap_future(cf_b), + asyncio.wrap_future(cf_c), + ) + # Highest priority (lowest value) first. + self.assertEqual(results, ["c", "b", "a"]) + finally: + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + +class FifoTiebreakTest(later.unittest.TestCase): + """Same priority -> FIFO via monotonic seq.""" + + async def test_same_priority_is_fifo(self) -> None: + with ThreadPoolExecutor(max_workers=1) as pool: + scheduler = PriorityScheduler(max_concurrent=1) + scheduler._underlying_executor = pool + info = _make_info("t") + scheduler.register_stage(info, priority=0) + shim = _PrioritizedExecutor(scheduler, info) + + results: list[int] = [] + + def work(val: int) -> int: + results.append(val) + return val + + cfs = [shim.submit(work, i) for i in range(5)] + await asyncio.sleep(0) + + run_task = asyncio.create_task(scheduler.run()) + try: + await asyncio.gather(*[asyncio.wrap_future(cf) for cf in cfs]) + self.assertEqual(results, [0, 1, 2, 3, 4]) + finally: + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + +class MaxConcurrentDispatchTest(later.unittest.TestCase): + """At most max_concurrent items dispatched simultaneously.""" + + async def test_max_concurrent_two(self) -> None: + with ThreadPoolExecutor(max_workers=8) as pool: + scheduler = PriorityScheduler(max_concurrent=2) + scheduler._underlying_executor = pool + info = _make_info("t") + scheduler.register_stage(info, priority=0) + shim = _PrioritizedExecutor(scheduler, info) + + active = 0 + max_active = 0 + # threading.Lock() returns _thread.LockType, which Pyre cannot + # narrow from a captured local. Store as Any to silence captured- + # variable annotation warnings without an explicit ignore. + lock: Any = threading.Lock() + + def tracked(val: int) -> int: + nonlocal active, max_active + with lock: + active += 1 + max_active = max(max_active, active) + time.sleep(0.05) + with lock: + active -= 1 + return val + + cfs = [shim.submit(tracked, i) for i in range(8)] + + run_task = asyncio.create_task(scheduler.run()) + try: + await asyncio.gather(*[asyncio.wrap_future(cf) for cf in cfs]) + self.assertLessEqual(max_active, 2) + finally: + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + +class ExceptionPropagationTest(later.unittest.TestCase): + """Exception in func -> set_exception on cf_future.""" + + async def test_exception_propagates(self) -> None: + with ThreadPoolExecutor(max_workers=1) as pool: + scheduler = PriorityScheduler(max_concurrent=1) + scheduler._underlying_executor = pool + info = _make_info("t") + scheduler.register_stage(info, priority=0) + shim = _PrioritizedExecutor(scheduler, info) + + def boom(x: int) -> int: + raise ValueError(f"boom-{x}") + + run_task = asyncio.create_task(scheduler.run()) + try: + cf = shim.submit(boom, 7) + with self.assertRaises(ValueError) as cm: + await asyncio.wait_for(asyncio.wrap_future(cf), timeout=2.0) + self.assertIn("boom-7", str(cm.exception)) + finally: + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + +class DrainPendingOnShutdownTest(later.unittest.TestCase): + """V5.2: _drain_pending cancels everything still on the heap.""" + + async def test_drain_cancels_pending_items(self) -> None: + with ThreadPoolExecutor(max_workers=1) as pool: + scheduler = PriorityScheduler(max_concurrent=1) + scheduler._underlying_executor = pool + info = _make_info("t") + scheduler.register_stage(info, priority=0) + shim = _PrioritizedExecutor(scheduler, info) + + cf1 = shim.submit(lambda: 1) + cf2 = shim.submit(lambda: 2) + await asyncio.sleep(0) + self.assertEqual(len(scheduler._heap), 2) + + scheduler._drain_pending() + self.assertEqual(len(scheduler._heap), 0) + self.assertTrue(cf1.cancelled()) + self.assertTrue(cf2.cancelled()) + + +class PrioritySchedulerBackgroundTaskTest(later.unittest.TestCase): + """The BG-task adapter runs scheduler.run() and drains on cancel.""" + + async def test_bg_task_runs_scheduler_and_drains(self) -> None: + with ThreadPoolExecutor(max_workers=1) as pool: + scheduler = PriorityScheduler(max_concurrent=1) + scheduler._underlying_executor = pool + info = _make_info("t") + scheduler.register_stage(info, priority=0) + shim = _PrioritizedExecutor(scheduler, info) + + bg = _PrioritySchedulerBackgroundTask(scheduler) + run_task = asyncio.create_task(bg.run()) + + cf = shim.submit(lambda: "ok") + result = await asyncio.wait_for(asyncio.wrap_future(cf), timeout=2.0) + self.assertEqual(result, "ok") + + # Submit one more then cancel BG task BEFORE it dispatches. + cf2 = shim.submit(lambda: time.sleep(10)) + run_task.cancel() + try: + await run_task + except asyncio.CancelledError: + pass + + # _drain_pending in the BG-task `finally` block should have + # cancelled the un-dispatched item. + self.assertTrue(cf2.cancelled()) + + +class PipelineWithPrioritySchedulerEndToEndTest(unittest.TestCase): + """End-to-end: build(use_priority_scheduler=True) produces correct output.""" + + def test_pipeline_with_scheduler_produces_output(self) -> None: + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(lambda x: x * 2, concurrency=2) + .pipe(lambda x: x + 1, concurrency=2) + .add_sink(3) + .build(num_threads=4, use_priority_scheduler=True) + ) + + results: list[int] = [] + with pipeline.auto_stop(): + for item in pipeline: + results.append(item) + + self.assertEqual(sorted(results), sorted(x * 2 + 1 for x in range(10))) + + def test_pipeline_without_scheduler_unchanged(self) -> None: + pipeline = ( + PipelineBuilder() + .add_source(range(10)) + .pipe(lambda x: x * 2, concurrency=2) + .add_sink(3) + .build(num_threads=4, use_priority_scheduler=False) + ) + + results: list[int] = [] + with pipeline.auto_stop(): + for item in pipeline: + results.append(item) + + self.assertEqual(sorted(results), sorted(x * 2 for x in range(10))) + + +class PipelineSchedulerBypassTest(unittest.TestCase): + """V5: async/generator stages bypass the scheduler entirely.""" + + def test_async_stage_bypasses_scheduler(self) -> None: + async def async_double(x: int) -> int: + return x * 2 + + pipeline = ( + PipelineBuilder() + .add_source(range(5)) + .pipe(async_double, concurrency=2) + .add_sink(3) + .build(num_threads=4, use_priority_scheduler=True) + ) + + results: list[int] = [] + with pipeline.auto_stop(): + for item in pipeline: + results.append(item) + + self.assertEqual(sorted(results), sorted(x * 2 for x in range(5))) + + def test_generator_stage_bypasses_scheduler(self) -> None: + def gen_double(x: int) -> Iterator[int]: + yield x * 2 + + pipeline = ( + PipelineBuilder() + .add_source(range(5)) + .pipe(gen_double, concurrency=2) + .add_sink(3) + .build(num_threads=4, use_priority_scheduler=True) + ) + + results: list[int] = [] + with pipeline.auto_stop(): + for item in pipeline: + results.append(item) + + self.assertEqual(sorted(results), sorted(x * 2 for x in range(5))) + + +class PipeArgsRevertTest(unittest.TestCase): + """V5: _PipeArgs has no nice/_depth fields anymore (Diff 2 reverts v2.1).""" + + def test_pipe_args_minimal_fields(self) -> None: + from spdl.pipeline.defs import _PipeArgs + + args = _PipeArgs(op=lambda x: x) + # The v2.1 spec added .nice and ._depth — v5 removes both. + self.assertFalse(hasattr(args, "nice")) + self.assertFalse(hasattr(args, "_depth")) + + +class ToAsyncRevertTest(unittest.TestCase): + """V5: _to_async no longer takes scheduler/stage_name params.""" + + def test_to_async_signature(self) -> None: + import inspect as _inspect + + from spdl.pipeline._common._convert import _to_async + + sig = _inspect.signature(_to_async) + params = list(sig.parameters.keys()) + # Only func + executor should remain. + self.assertEqual(params, ["func", "executor"]) + + def test_convert_to_async_signature(self) -> None: + import inspect as _inspect + + from spdl.pipeline._common._convert import convert_to_async + + sig = _inspect.signature(convert_to_async) + params = list(sig.parameters.keys()) + self.assertEqual(params, ["op", "executor"]) + + +class NodeDepthComputationTest(unittest.TestCase): + """V5: depth is computed at build time via _node_depth (not on _PipeArgs).""" + + def test_linear_pipeline_depths(self) -> None: + from spdl.pipeline._components._node import ( + _convert_config, + _MutableInt, + _node_depth, + ) + from spdl.pipeline._components._queue import AsyncQueue + from spdl.pipeline.defs import ( + Pipe, + PipelineConfig, + SinkConfig, + SourceConfig, + ) + + plc = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[ + Pipe(lambda x: x, name="s0"), + Pipe(lambda x: x, name="s1"), + Pipe(lambda x: x, name="s2"), + ], + sink=SinkConfig(3), + ) + + sink_node = _convert_config(plc, AsyncQueue, 0, _MutableInt(0)) + # source(0) -> s0(1) -> s1(2) -> s2(3) -> sink(4) + self.assertEqual(_node_depth(sink_node), 4) + + def test_zero_pipe_pipeline_depth(self) -> None: + """No pipes between source and sink: depth(sink) == 1.""" + # Arrange + from spdl.pipeline._components._node import ( + _convert_config, + _MutableInt, + _node_depth, + ) + from spdl.pipeline._components._queue import AsyncQueue + from spdl.pipeline.defs import PipelineConfig, SinkConfig, SourceConfig + + plc = PipelineConfig( + src=SourceConfig([1, 2, 3]), + pipes=[], + sink=SinkConfig(3), + ) + + # Act + sink_node = _convert_config(plc, AsyncQueue, 0, _MutableInt(0)) + + # Assert: source(0) -> sink(1). + self.assertEqual(_node_depth(sink_node), 1) + + def test_merge_pipeline_depth_takes_max_branch(self) -> None: + """For a merge node, depth == 1 + max(depth(branch_i)).""" + # Arrange + from spdl.pipeline._components._node import ( + _convert_config, + _MutableInt, + _node_depth, + ) + from spdl.pipeline._components._queue import AsyncQueue + from spdl.pipeline.defs import ( + Merge, + Pipe, + PipelineConfig, + SinkConfig, + SourceConfig, + ) + + # Branch A: 1 pipe (depth 2 at branch's sink) + branch_a = PipelineConfig( + src=SourceConfig([1]), + pipes=[Pipe(lambda x: x, name="a0")], + sink=SinkConfig(3), + ) + # Branch B: 3 pipes (depth 4 at branch's sink) + branch_b = PipelineConfig( + src=SourceConfig([1]), + pipes=[ + Pipe(lambda x: x, name="b0"), + Pipe(lambda x: x, name="b1"), + Pipe(lambda x: x, name="b2"), + ], + sink=SinkConfig(3), + ) + merged = PipelineConfig( + src=Merge([branch_a, branch_b]), + pipes=[], + sink=SinkConfig(3), + ) + + # Act + sink_node = _convert_config(merged, AsyncQueue, 0, _MutableInt(0)) + + # Assert: deepest branch contributes; downstream merge + sink each + # add one. Branch B reaches depth 4 at its own sink; merge wraps + # both and the outer sink follows. + depth = _node_depth(sink_node) + self.assertGreaterEqual(depth, 5) + + +if __name__ == "__main__": + unittest.main()