diff --git a/src/spdl/pipeline/_build.py b/src/spdl/pipeline/_build.py index 803e81d9f..9f2c0a299 100644 --- a/src/spdl/pipeline/_build.py +++ b/src/spdl/pipeline/_build.py @@ -31,6 +31,7 @@ _get_global_id, _set_global_id, AsyncQueue, + ResizableSemaphore, StageInfo, TaskHook, ) @@ -128,6 +129,9 @@ 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, + enable_adaptive_concurrency: bool = False, + _install_semaphores_for_test: bool = False, ) -> Pipeline[U]: if _DEFAULT_BUILD_CALLBACK is not None: try: @@ -147,6 +151,50 @@ 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) + ) + + # V5.1+V5.5 Diff 3a: pre-allocate the per-pipeline registries that + # `_components/_node.py` populates when semaphore installation is + # enabled (either via the public ``enable_adaptive_concurrency`` flag + # or the test-only ``_install_semaphores_for_test`` knob). These are + # then handed off to ``_PipelineImpl.__init__`` so + # ``Pipeline._resize_concurrency_async`` (Diff 3b internal) can find + # them. + # ``output_queue_by_name`` (Phase D) caches each registered stage's + # output ``AsyncQueue`` so the Diff 6 controller can read its lap + # stats — the same key set as ``semaphore_registry``. + semaphore_registry: dict[str, ResizableSemaphore] = {} + dynamic_concurrency: dict[str, int] = {} + stage_info_by_name: dict[str, StageInfo] = {} + output_queue_by_name: dict[str, AsyncQueue] = {} + + install_semaphores = enable_adaptive_concurrency or _install_semaphores_for_test + coro, queue = _build_pipeline_coro( pipeline_cfg, max_failures=max_failures, @@ -155,13 +203,24 @@ def _build_pipeline( task_hook_factory=task_hook_factory, stage_id=stage_id, background_tasks=all_bg_tasks or None, + scheduler=scheduler, + install_semaphores_for_test=install_semaphores, + semaphore_registry=semaphore_registry, + dynamic_concurrency=dynamic_concurrency, + stage_info_by_name=stage_info_by_name, + output_queue_by_name=output_queue_by_name, ) - executor = ThreadPoolExecutor( - max_workers=num_threads, - thread_name_prefix="spdl_worker_thread_", + return Pipeline( + coro, + queue, + executor, + desc=desc, + semaphore_registry=semaphore_registry, + dynamic_concurrency=dynamic_concurrency, + stage_info_by_name=stage_info_by_name, + output_queue_by_name=output_queue_by_name, ) - return Pipeline(coro, queue, executor, desc=desc) def build_pipeline( @@ -175,6 +234,9 @@ 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, + enable_adaptive_concurrency: bool = False, + _install_semaphores_for_test: bool = False, ) -> Pipeline[U]: """Build a pipeline from the config. @@ -240,6 +302,31 @@ 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. + + enable_adaptive_concurrency: If ``True``, every ``Pipe`` stage is + built with a :py:class:`ResizableSemaphore` whose initial value + matches its static ``concurrency``, and the semaphore is wired + to the V5.6 REPLACE admission gate in + :py:func:`~spdl.pipeline._components._pipe._pipe`. This enables + runtime adjustment of per-stage concurrency via the internal + :py:meth:`Pipeline._resize_concurrency_async` (intended to be + driven by an in-loop adaptive-concurrency controller running + as a :py:class:`BackgroundTask`). Default: ``False`` + (per-stage concurrency is fixed at build time, with zero + per-task overhead in the admission gate). + + _install_semaphores_for_test: **Test-only.** Same mechanical effect + as ``enable_adaptive_concurrency`` (both flip the same internal + switch), kept as a separate flag so tests can opt in without + implying the production-facing semantic. Used by the V5.5 + throughput regression test and the in-loop async-resize + regression test. Production code MUST use + ``enable_adaptive_concurrency``. """ from . import _profile @@ -255,6 +342,9 @@ def build_pipeline( task_hook_factory=task_hook_factory, stage_id=stage_id, background_tasks=background_tasks, + use_priority_scheduler=use_priority_scheduler, + enable_adaptive_concurrency=enable_adaptive_concurrency, + _install_semaphores_for_test=_install_semaphores_for_test, ) diff --git a/src/spdl/pipeline/_builder.py b/src/spdl/pipeline/_builder.py index 5ec26abb5..c57255591 100644 --- a/src/spdl/pipeline/_builder.py +++ b/src/spdl/pipeline/_builder.py @@ -291,6 +291,8 @@ 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, + enable_adaptive_concurrency: bool = False, ) -> Pipeline[U]: """Build the pipeline. @@ -328,6 +330,19 @@ 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. + + enable_adaptive_concurrency: If ``True``, every ``Pipe`` stage + is built with a :py:class:`ResizableSemaphore` so per-stage + concurrency can be adjusted at runtime via the internal + :py:meth:`Pipeline._resize_concurrency_async` (intended to + be driven by an in-loop adaptive-concurrency controller + running as a :py:class:`BackgroundTask`). Default: + ``False`` (per-stage concurrency is fixed at build time). """ return build_pipeline( self.get_config(), @@ -337,4 +352,6 @@ def build( report_stats_interval=report_stats_interval, task_hook_factory=task_hook_factory, stage_id=stage_id, + use_priority_scheduler=use_priority_scheduler, + enable_adaptive_concurrency=enable_adaptive_concurrency, ) 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..cea1fea9c 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 @@ -41,6 +42,7 @@ _pipe, ) from ._queue import AsyncQueue, get_default_queue_class +from ._semaphore import ResizableSemaphore from ._sink import _sink from ._source import _source, _source_continuous from ._variants import _path_variants_router @@ -479,11 +481,77 @@ def _convert_config( return n +def _qualified_name(info: StageInfo, branch_label: str | None = None) -> str: + """Build the qualified stage name used by ``Pipeline._resize_concurrency_async``. + + For non-MultiPipe stages, ``qualified_name = info.stage_name``. + For MultiPipe sub-pipelines, the branch label is prefixed with ``/`` + (e.g., ``"video/decode_frame"``). LCA's MultiPipe is a single SPDL + Pipe with internal dispatcher so qualified names match plain + ``stage_name`` in practice today; the addressing scheme is + forward-compatible for true SPDL fan-out. + """ + if branch_label is None: + return info.stage_name + return f"{branch_label}/{info.stage_name}" + + +def _register_semaphore( + semaphore_registry: dict[str, ResizableSemaphore] | None, + dynamic_concurrency: dict[str, int] | None, + stage_info_by_name: dict[str, StageInfo] | None, + output_queue_by_name: dict[str, AsyncQueue] | None, + qname: str, + info: StageInfo, + sem: ResizableSemaphore, + initial_value: int, + output_queue: AsyncQueue, +) -> None: + """Insert ``sem`` into the per-pipeline registries under ``qname``. + + Duplicate ``qname`` is treated as an internal error (each stage must + have a unique qualified name across the pipeline). This is the only + write site for ``_PipelineImpl._semaphore_registry`` / + ``_dynamic_concurrency`` / ``_stage_info_by_name`` / + ``_output_queue_by_name``. + + Phase D: ``output_queue_by_name`` mirrors the same key set as + ``semaphore_registry``. The captured queue handle is the SAME + instance that this stage's ``_pipe()`` coroutine writes to and + the next stage's coroutine reads from — i.e., the canonical + per-stage output queue. The Diff 6 + ``DomeVideoConcurrencyController`` reads this dict to obtain + each stage's ``StatsQueue._last_lap_stats`` (or equivalent + cached lap stats) for adaptive-tuning decisions. + """ + if semaphore_registry is None: + return + if qname in semaphore_registry: + raise RuntimeError( + f"Duplicate qualified stage name {qname!r}. " + f"This is an internal error — please report." + ) + semaphore_registry[qname] = sem + if dynamic_concurrency is not None: + dynamic_concurrency[qname] = initial_value + if stage_info_by_name is not None: + stage_info_by_name[qname] = info + if output_queue_by_name is not None: + output_queue_by_name[qname] = output_queue + + def _build_node( node: _TNodes, fc_class: type[_FailCounter], task_hook_factory: Callable[[StageInfo], list[TaskHook]], max_failures: int | Fraction, + scheduler: Any = None, + depth: int = 0, + install_semaphores_for_test: bool = False, + semaphore_registry: dict[str, ResizableSemaphore] | None = None, + dynamic_concurrency: dict[str, int] | None = None, + stage_info_by_name: dict[str, StageInfo] | None = None, + output_queue_by_name: dict[str, AsyncQueue] | None = None, ) -> None: """Build a coroutine for a single node based on its configuration type. @@ -520,6 +588,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 +657,68 @@ 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) + + # V5.1+V5.6 Diff 3a: opt the stage into the per-pipeline + # semaphore registry. When ``install_semaphores_for_test`` + # is True (test-only knob), every Pipe stage gets a + # ``ResizableSemaphore`` whose initial value matches the + # static ``args.concurrency``. The semaphore is passed + # to ``_pipe`` which uses V5.6 REPLACE semantics: the + # semaphore IS the admission gate (the static + # ``len(tasks) >= concurrency`` check is skipped). + # When the knob is off, ``semaphore=None`` and ``_pipe`` + # uses its existing static gate — ZERO per-task overhead. + pipe_semaphore: ResizableSemaphore | None = None + if install_semaphores_for_test: + qname = _qualified_name(node.info) + pipe_semaphore = ResizableSemaphore(args.concurrency) + _register_semaphore( + semaphore_registry, + dynamic_concurrency, + stage_info_by_name, + output_queue_by_name, + qname, + node.info, + pipe_semaphore, + args.concurrency, + out_q, + ) + 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, + semaphore=pipe_semaphore, ) case _PipeType.OrderedPipe: + # OrderedPipe uses an intermediate queue sized + # to ``concurrency`` and is not part of the + # Diff 3a admission-gate change. Static + # concurrency only. 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 +756,57 @@ 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, + install_semaphores_for_test: bool = False, + semaphore_registry: dict[str, ResizableSemaphore] | None = None, + dynamic_concurrency: dict[str, int] | None = None, + stage_info_by_name: dict[str, StageInfo] | None = None, + output_queue_by_name: dict[str, AsyncQueue] | None = None, ) -> None: """Recursively build coroutines for a node and all its upstream nodes. @@ -640,6 +819,27 @@ 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. + install_semaphores_for_test: V5.5 test-only knob. When True, + every Pipe stage gets a :py:class:`ResizableSemaphore` whose + initial value matches its static ``concurrency``, and the + semaphore is passed to ``_pipe()`` so its V5.6 REPLACE branch + governs admission. Production code should pass + ``enable_adaptive_concurrency=True`` to + :py:func:`build_pipeline` instead so that an in-loop + controller can call + :py:meth:`Pipeline._resize_concurrency_async`. + semaphore_registry: Per-pipeline ``dict[qualified_name, ResizableSemaphore]`` + populated when ``install_semaphores_for_test=True``. + dynamic_concurrency: Per-pipeline ``dict[qualified_name, int]`` + populated when ``install_semaphores_for_test=True``. + stage_info_by_name: Per-pipeline ``dict[qualified_name, StageInfo]`` + populated when ``install_semaphores_for_test=True``. + output_queue_by_name: Per-pipeline + ``dict[qualified_name, AsyncQueue]`` populated when + ``install_semaphores_for_test=True``. Phase D: enables the + Diff 6 controller to read each stage's lap stats. Raises: RuntimeError: If attempting to build a coroutine for a node that already has one. @@ -648,9 +848,33 @@ 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, + install_semaphores_for_test=install_semaphores_for_test, + semaphore_registry=semaphore_registry, + dynamic_concurrency=dynamic_concurrency, + stage_info_by_name=stage_info_by_name, + output_queue_by_name=output_queue_by_name, + ) - _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, + install_semaphores_for_test=install_semaphores_for_test, + semaphore_registry=semaphore_registry, + dynamic_concurrency=dynamic_concurrency, + stage_info_by_name=stage_info_by_name, + output_queue_by_name=output_queue_by_name, + ) # Used to append stage name with pipeline @@ -696,6 +920,12 @@ def _build_pipeline_node( queue_class: type[AsyncQueue] | None, task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None, stage_id: int, + scheduler: Any = None, + install_semaphores_for_test: bool = False, + semaphore_registry: dict[str, ResizableSemaphore] | None = None, + dynamic_concurrency: dict[str, int] | None = None, + stage_info_by_name: dict[str, StageInfo] | None = None, + output_queue_by_name: dict[str, AsyncQueue] | None = None, ) -> _TOutputNodes: global _PIPELINE_ID _PIPELINE_ID += 1 @@ -710,7 +940,18 @@ 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, + install_semaphores_for_test=install_semaphores_for_test, + semaphore_registry=semaphore_registry, + dynamic_concurrency=dynamic_concurrency, + stage_info_by_name=stage_info_by_name, + output_queue_by_name=output_queue_by_name, + ) return node @@ -950,7 +1191,13 @@ def _build_pipeline_coro( task_hook_factory: Callable[[StageInfo], list[TaskHook]] | None = None, stage_id: int = 0, background_tasks: Sequence[BackgroundTaskFactory] | None = None, -) -> tuple[Coroutine[None, None, None], asyncio.Queue]: + scheduler: Any = None, + install_semaphores_for_test: bool = False, + semaphore_registry: dict[str, ResizableSemaphore] | None = None, + dynamic_concurrency: dict[str, int] | None = None, + stage_info_by_name: dict[str, StageInfo] | None = None, + output_queue_by_name: dict[str, AsyncQueue] | None = None, +) -> tuple[Coroutine[None, None, None], AsyncQueue]: try: node = _build_pipeline_node( plc, @@ -959,6 +1206,12 @@ def _build_pipeline_coro( queue_class=queue_class, task_hook_factory=task_hook_factory, stage_id=stage_id, + scheduler=scheduler, + install_semaphores_for_test=install_semaphores_for_test, + semaphore_registry=semaphore_registry, + dynamic_concurrency=dynamic_concurrency, + stage_info_by_name=stage_info_by_name, + output_queue_by_name=output_queue_by_name, ) coro = _run_pipeline_coroutines(node, background_tasks=background_tasks) diff --git a/src/spdl/pipeline/_components/_pipe.py b/src/spdl/pipeline/_components/_pipe.py index 7b96dddc2..f55631995 100644 --- a/src/spdl/pipeline/_components/_pipe.py +++ b/src/spdl/pipeline/_components/_pipe.py @@ -27,6 +27,7 @@ from ._common import _EOF, _EPOCH_END, _SKIP, is_eof, is_epoch_end, StageInfo from ._hook import _stage_hooks, _task_hooks, TaskHook from ._queue import _queue_stage_hook, AsyncQueue +from ._semaphore import ResizableSemaphore # pyre-strict @@ -247,6 +248,7 @@ def _pipe( fail_counter: _FailCounter, task_hooks: list[TaskHook], op_requires_eof: bool, + semaphore: ResizableSemaphore | None = None, ) -> Coroutine: """Create a coroutine for processing data from input queue to output queue. @@ -264,6 +266,12 @@ def _pipe( task_hooks: List of hooks for monitoring task execution. op_requires_eof: If True, pass EOF token to the operation; otherwise stop processing before EOF. + semaphore: Optional :py:class:`ResizableSemaphore` admission gate + (V5.6). When provided, the static ``len(tasks) >= concurrency`` + gate is **REPLACED** by ``await semaphore.acquire()`` — + ``args.concurrency`` is ignored for admission control. + When ``None`` (default), behaviour is unchanged from + pre-V5: the static admission gate applies. Returns: A coroutine that executes the pipeline stage. @@ -294,6 +302,26 @@ def _wrap(coro: AsyncIterator[U], item: Any = None) -> Coroutine: else: raise ValueError(f"{afunc=} must be either async function or async generator.") + # V5.6: branch ONCE outside the hot loop. When `semaphore` is provided, + # the admission gate is REPLACED — `await semaphore.acquire()` becomes + # the gate, and the static `len(tasks) >= concurrency` check is + # skipped entirely. When `semaphore` is None (the default), the + # original gate applies unchanged so there is ZERO per-task overhead + # added to the existing fast path. + if semaphore is not None: + return _pipe_with_semaphore( + info, + input_queue, + output_queue, + afunc, + _wrap, + args, + fail_counter, + hooks, + op_requires_eof, + semaphore, + ) + @_queue_stage_hook(output_queue) @_stage_hooks(hooks) async def pipe() -> None: @@ -336,6 +364,85 @@ async def pipe() -> None: return pipe() +def _pipe_with_semaphore( + info: StageInfo, + input_queue: AsyncQueue, + output_queue: AsyncQueue, + # pyre-ignore[2]: afunc has a polymorphic shape (afunc / agen) + afunc: Callable[[T], Any], + # pyre-ignore[2]: _wrap closure type matches the branch in _pipe + _wrap: Callable[..., Coroutine], + args: _PipeArgs[T, U], + fail_counter: _FailCounter, + hooks: list[TaskHook], + op_requires_eof: bool, + semaphore: ResizableSemaphore, +) -> Coroutine: + """V5.6 REPLACE branch: ``semaphore.acquire()`` IS the admission gate. + + ``args.concurrency`` is ignored — the registered semaphore (whose value + is mutable via :py:meth:`Pipeline._resize_concurrency_async`) governs + the in-flight cap. Task completion calls ``semaphore.release()`` via a + done-callback so the admit cycle is symmetric with the gate. + """ + + # Define _on_done OUTSIDE the loop so the closure captures the + # single ``tasks`` set / ``semaphore`` once, and so flake8 B023 + # ("loop variable") doesn't fire. ``tasks`` is mutated in-place via + # ``add()`` / ``discard()`` / ``clear()`` (never rebound) so the + # closure always sees the current contents. + tasks: set[asyncio.Task[Any]] = set() + + def _on_done(t: asyncio.Task[Any]) -> None: + tasks.discard(t) + semaphore.release() + + @_queue_stage_hook(output_queue) + @_stage_hooks(hooks) + async def pipe() -> None: + i = 0 + while not fail_counter.too_many_failures(): + item = await input_queue.get() + + if is_epoch_end(item): + # Epoch boundary: wait for all in-flight tasks, propagate, continue + if tasks: + await asyncio.wait(tasks) + # Done callbacks have already discarded each task as it + # completed, so ``tasks`` is empty here. Clear in-place + # (no rebind) for symmetry with the static-gate branch. + tasks.clear() + await output_queue.put(_EPOCH_END) + continue + + if is_eof(item) and not op_requires_eof: + break + + # V5.6 admission gate: REPLACES `len(tasks) >= args.concurrency`. + # `acquire()` blocks here when the in-flight count reaches the + # semaphore's current value (which may have been resized via + # Pipeline._resize_concurrency_async). + await semaphore.acquire() + + task = create_task( + _wrap(afunc(item), item), + name=f"{info}:{(i := i + 1)}", + ) + tasks.add(task) + task.add_done_callback(_on_done) + + if is_eof(item): + break + + if tasks: + await asyncio.wait(tasks) + + if fail_counter.too_many_failures(): + fail_counter.raise_for_failures(info) + + return pipe() + + def _ordered_pipe( info: StageInfo, input_queue: AsyncQueue, diff --git a/src/spdl/pipeline/_components/_queue.py b/src/spdl/pipeline/_components/_queue.py index fbdcd6aca..15cf5eefd 100644 --- a/src/spdl/pipeline/_components/_queue.py +++ b/src/spdl/pipeline/_components/_queue.py @@ -196,6 +196,15 @@ def __init__( self._lap_ave_put_time = 0.0 self._lap_dur_empty = 0.0 + # Cached snapshot of the latest lap stats. ``_get_lap_stats()`` is + # destructive — it resets the lap counters — so non-callback readers + # (e.g., the LCA ``DomeVideoConcurrencyController``) cannot call it + # safely. ``_log_interval_stats()`` now writes the freshly computed + # stats here so an external reader can observe the latest interval + # without disturbing the periodic logging path. ``None`` until the + # first interval has elapsed. + self._last_lap_stats: QueuePerfStats | None = None + async def get(self) -> object: """Remove and return an item from the queue, track the time.""" with self._getc.count(): @@ -301,7 +310,14 @@ def _get_lap_stats(self) -> QueuePerfStats: ) async def _log_interval_stats(self) -> None: - await self.interval_stats_callback(self._get_lap_stats()) + stats = self._get_lap_stats() + # Cache for non-callback readers (e.g., adaptive concurrency + # controller). ``_get_lap_stats()`` resets the lap counters, so + # this is the only safe way for an external reader to observe + # the most recent interval's stats without racing the periodic + # logging path. + self._last_lap_stats = stats + await self.interval_stats_callback(stats) async def interval_stats_callback(self, stats: QueuePerfStats) -> None: """Callback for processing interval performance statistics. 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/_pipeline.py b/src/spdl/pipeline/_pipeline.py index b47dd3ad9..7e90deaa0 100644 --- a/src/spdl/pipeline/_pipeline.py +++ b/src/spdl/pipeline/_pipeline.py @@ -12,7 +12,7 @@ import time import warnings import weakref -from asyncio import AbstractEventLoop, Queue as AsyncQueue +from asyncio import AbstractEventLoop from collections.abc import Coroutine, Iterator from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager @@ -22,7 +22,9 @@ from spdl._internal import log_api_usage_once from spdl.pipeline._common._misc import create_task -from spdl.pipeline._components import is_epoch_end +from spdl.pipeline._common._types import StageInfo +from spdl.pipeline._components import AsyncQueue, is_epoch_end +from spdl.pipeline._components._semaphore import ResizableSemaphore __all__ = ["Pipeline"] @@ -206,6 +208,10 @@ def __init__( executor: ThreadPoolExecutor, *, desc: str, + semaphore_registry: dict[str, ResizableSemaphore] | None = None, + dynamic_concurrency: dict[str, int] | None = None, + stage_info_by_name: dict[str, StageInfo] | None = None, + output_queue_by_name: dict[str, AsyncQueue] | None = None, ) -> None: self._str: str = "\n".join([repr(self), desc]) @@ -213,6 +219,37 @@ def __init__( self._event_loop = _EventLoop(coro, executor) self._event_loop_state: _EventLoopState = _EventLoopState.NOT_STARTED + # V5.1 sibling registries for runtime concurrency adjustment. + # Populated at pipeline build time by ``_components/_node.py`` + # when a stage opts into the registry (e.g., via the test-only + # ``_install_semaphores_for_test`` knob on ``build_pipeline``). + # Two parallel dicts keyed on the qualified stage name (see V5.4 + # Diff 3a). NOT a mutation of ``StageInfo`` (preserves + # frozen=True / hashability for third-party code). + self._semaphore_registry: dict[str, ResizableSemaphore] = ( + semaphore_registry if semaphore_registry is not None else {} + ) + self._dynamic_concurrency: dict[str, int] = ( + dynamic_concurrency if dynamic_concurrency is not None else {} + ) + # Useful for error messages and future Track B logging hooks; not + # mutated after build time. + self._stage_info_by_name: dict[str, StageInfo] = ( + stage_info_by_name if stage_info_by_name is not None else {} + ) + # Phase D: per-stage output queue handles, keyed by qualified + # stage name (same key set as ``_semaphore_registry``). Captured + # at registry-population time in ``_components/_node.py`` so the + # Diff 6 ``DomeVideoConcurrencyController`` can read each stage's + # ``StatsQueue._last_lap_stats`` (or equivalent cached lap stats) + # to drive the adaptive concurrency loop. The queue stored here + # is the SAME instance that is passed as the OUTPUT queue to the + # stage's ``_pipe()`` coroutine and as the INPUT queue to the + # next stage — it is the canonical per-stage output queue. + self._output_queue_by_name: dict[str, AsyncQueue] = ( + output_queue_by_name if output_queue_by_name is not None else {} + ) + def __str__(self) -> str: return self._str @@ -467,9 +504,20 @@ def __init__( executor: ThreadPoolExecutor, *, desc: str, + semaphore_registry: dict[str, ResizableSemaphore] | None = None, + dynamic_concurrency: dict[str, int] | None = None, + stage_info_by_name: dict[str, StageInfo] | None = None, + output_queue_by_name: dict[str, AsyncQueue] | None = None, ) -> None: self._impl: _PipelineImpl[T] = _PipelineImpl( - coro, output_queue, executor, desc=desc + coro, + output_queue, + executor, + desc=desc, + semaphore_registry=semaphore_registry, + dynamic_concurrency=dynamic_concurrency, + stage_info_by_name=stage_info_by_name, + output_queue_by_name=output_queue_by_name, ) self._finalizer = weakref.finalize(self, _stop_impl, self._impl) @@ -552,6 +600,66 @@ def __iter__(self) -> Iterator[T]: """Call :py:meth:`~spdl.pipeline.Pipeline.get_iterator` without arguments.""" return self.get_iterator() + # -------------------------------------------------------------- + # Diff 3b — runtime concurrency adjustment (INTERNAL) + # -------------------------------------------------------------- + + async def _resize_concurrency_async( + self, + qualified_name: str, + new_value: int, + ) -> None: + """Resize the in-flight admission cap for a registered stage. + + INTERNAL — must be awaited from a coroutine running on the + pipeline's own event loop (e.g., a + :py:class:`~spdl.pipeline.BackgroundTask`). There is no + public, foreground-thread wrapper: cross-thread resize is + intentionally not exposed because the only intended caller is + an in-loop adaptive-concurrency controller. + + Atomicity: the body has zero ``await`` statements between + :py:meth:`ResizableSemaphore.resize` and the + ``_dynamic_concurrency`` dict assignment. asyncio is + single-threaded, so the entire method runs in one event-loop + turn and is therefore cancel-safe by structural invariant + (``CancelledError`` can only fire BEFORE the call begins or + AFTER it completes, never between the two writes). + + Args: + qualified_name: A fully-qualified stage name. For + non-MultiPipe stages this equals + :py:attr:`StageInfo.stage_name` (e.g., + ``"decode_single_frame"``). For MultiPipe sub-pipelines + it is ``"/"``. The set of + valid names is ``pipeline._impl._semaphore_registry``. + new_value: New admission cap; must be >= 1. Resize-up is + immediate; resize-down is graceful — in-flight tasks + finish, but no new tasks admit until in-flight drops + below ``new_value``. + + Raises: + ValueError: ``new_value < 1``. + KeyError: ``qualified_name`` is not registered. The error + message includes the list of valid names. + """ + if new_value < 1: + raise ValueError(f"new_value must be >= 1, got {new_value}") + + registry = self._impl._semaphore_registry + sem = registry.get(qualified_name) + if sem is None: + valid = sorted(registry.keys()) + raise KeyError( + f"qualified_name {qualified_name!r} not found. " + f"Valid stage names: {valid}" + ) + sem.resize(new_value) + # Keep the sibling registry in sync for observability. + # asyncio is single-threaded so this update is atomic with the + # semaphore.resize() above (no concurrent writer). + self._impl._dynamic_concurrency[qualified_name] = new_value + class PipelineIterator(Generic[T]): """PipelineIterator()""" 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/admission_gate_perf_test.py b/tests/pipeline/admission_gate_perf_test.py new file mode 100644 index 000000000..7926cecd8 --- /dev/null +++ b/tests/pipeline/admission_gate_perf_test.py @@ -0,0 +1,383 @@ +# 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 + +"""V5.5 throughput regression test for the Diff 3a admission gate REPLACE. + +Builds a small 4-stage pipeline, pushes ``NUM_ITEMS`` items through, and +compares wall time in two modes: + +- **Baseline**: ``_install_semaphores_for_test=False`` — ``_pipe()`` uses + the legacy static ``len(tasks) >= concurrency`` admission gate. +- **Treatment**: ``_install_semaphores_for_test=True`` — ``_pipe()`` uses + the V5.6 REPLACE branch with a :py:class:`ResizableSemaphore`. + +Regression > :py:data:`REGRESSION_THRESHOLD` on either p50 or p99 across +:py:data:`NUM_TRIALS` runs fails the diff (per V5.5 spec). +""" + +import asyncio +import statistics +import time +import unittest + +from spdl.pipeline import build_pipeline +from spdl.pipeline._common._types import StageInfo +from spdl.pipeline._components._queue import QueuePerfStats, StatsQueue +from spdl.pipeline.defs import Pipe, PipelineConfig, SinkConfig, SourceConfig + + +def _stage_1(x: int) -> int: + return x + 1 + + +def _stage_2(x: int) -> int: + return x * 2 + + +def _stage_3(x: int) -> int: + return x - 1 + + +def _stage_4(x: int) -> int: + return x // 2 + + +class AdmissionGateThroughputRegressionTest(unittest.TestCase): + """V5.5: admission gate REPLACE must not regress throughput by > 2%. + + A 4-stage sync pipeline is built with concurrency=4 per stage, fed + ``NUM_ITEMS`` items, and run ``NUM_TRIALS`` times in each mode. + + Note (perf-reviewer WARN-1): this regression test covers sync stages + only — extending coverage to async/sync_iter handlers is a follow-up + after Diff 3a lands. + """ + + NUM_ITEMS: int = 10_000 + NUM_TRIALS: int = 5 + REGRESSION_THRESHOLD: float = 0.02 # 2% + + def _build_config(self) -> PipelineConfig[int]: + return PipelineConfig( + src=SourceConfig(iter(range(self.NUM_ITEMS))), + pipes=[ + Pipe(_stage_1, concurrency=4), + Pipe(_stage_2, concurrency=4), + Pipe(_stage_3, concurrency=4), + Pipe(_stage_4, concurrency=4), + ], + sink=SinkConfig(buffer_size=64), + ) + + def _run_once(self, *, install_semaphores: bool) -> float: + config = self._build_config() + pipeline = build_pipeline( + config, + num_threads=8, + _install_semaphores_for_test=install_semaphores, + ) + t0 = time.perf_counter() + with pipeline.auto_stop(): + count = 0 + for _ in pipeline.get_iterator(): + count += 1 + self.assertEqual(count, self.NUM_ITEMS) + return time.perf_counter() - t0 + + def test_admission_gate_overhead_within_two_percent(self) -> None: + baseline_times = [ + self._run_once(install_semaphores=False) for _ in range(self.NUM_TRIALS) + ] + flagged_times = [ + self._run_once(install_semaphores=True) for _ in range(self.NUM_TRIALS) + ] + + baseline_p50 = statistics.median(baseline_times) + flagged_p50 = statistics.median(flagged_times) + # 5 trials → "p99" is effectively the max. + baseline_p99 = max(baseline_times) + flagged_p99 = max(flagged_times) + + p50_regression = (flagged_p50 - baseline_p50) / baseline_p50 + p99_regression = (flagged_p99 - baseline_p99) / baseline_p99 + + self.assertLess( + p50_regression, + self.REGRESSION_THRESHOLD, + f"p50 regression {p50_regression:.1%} exceeds " + f"{self.REGRESSION_THRESHOLD:.0%} threshold " + f"(baseline={baseline_p50:.3f}s, flagged={flagged_p50:.3f}s)", + ) + self.assertLess( + p99_regression, + self.REGRESSION_THRESHOLD, + f"p99 regression {p99_regression:.1%} exceeds " + f"{self.REGRESSION_THRESHOLD:.0%} threshold " + f"(baseline={baseline_p99:.3f}s, flagged={flagged_p99:.3f}s)", + ) + + +class AdmissionGateRegistryWiringTest(unittest.TestCase): + """Smoke test: ``_install_semaphores_for_test=True`` populates the registry.""" + + def test_registry_populated_for_each_pipe_stage(self) -> None: + config = PipelineConfig( + src=SourceConfig(iter(range(4))), + pipes=[ + Pipe(_stage_1, concurrency=2, name="s1"), + Pipe(_stage_2, concurrency=3, name="s2"), + ], + sink=SinkConfig(buffer_size=4), + ) + pipeline = build_pipeline( + config, + num_threads=4, + _install_semaphores_for_test=True, + ) + impl = pipeline._impl + # Both Pipe stages should be registered. The qualified name equals + # ``StageInfo.stage_name`` for non-MultiPipe stages (V5.4). + registry_keys = sorted(impl._semaphore_registry.keys()) + self.assertEqual(len(registry_keys), 2) + self.assertIn("s1", registry_keys) + self.assertIn("s2", registry_keys) + # ``_dynamic_concurrency`` mirrors the build-time concurrency. + self.assertEqual(impl._dynamic_concurrency["s1"], 2) + self.assertEqual(impl._dynamic_concurrency["s2"], 3) + # ``_stage_info_by_name`` mirrors the StageInfo for error messages. + self.assertEqual(impl._stage_info_by_name["s1"].stage_name, "s1") + self.assertEqual(impl._stage_info_by_name["s2"].stage_name, "s2") + + def test_registry_empty_when_flag_off(self) -> None: + config = PipelineConfig( + src=SourceConfig(iter(range(4))), + pipes=[ + Pipe(_stage_1, concurrency=2, name="s1"), + ], + sink=SinkConfig(buffer_size=4), + ) + pipeline = build_pipeline(config, num_threads=2) + impl = pipeline._impl + self.assertEqual(impl._semaphore_registry, {}) + self.assertEqual(impl._dynamic_concurrency, {}) + self.assertEqual(impl._stage_info_by_name, {}) + # Phase D: ``_output_queue_by_name`` mirrors the same key set, + # so it is also empty when the flag is off. + self.assertEqual(impl._output_queue_by_name, {}) + + def test_output_queue_by_name_mirrors_semaphore_registry_keys(self) -> None: + """Phase D: every entry in ``_semaphore_registry`` MUST have a + corresponding entry in ``_output_queue_by_name`` under the same + qualified name. The Diff 6 controller relies on this invariant + when iterating ``_semaphore_registry`` and looking up the queue + for each registered stage. + """ + config = PipelineConfig( + src=SourceConfig(iter(range(4))), + pipes=[ + Pipe(_stage_1, concurrency=2, name="s1"), + Pipe(_stage_2, concurrency=3, name="s2"), + Pipe(_stage_3, concurrency=4, name="s3"), + ], + sink=SinkConfig(buffer_size=4), + ) + pipeline = build_pipeline( + config, + num_threads=4, + _install_semaphores_for_test=True, + ) + impl = pipeline._impl + sem_keys = sorted(impl._semaphore_registry.keys()) + queue_keys = sorted(impl._output_queue_by_name.keys()) + self.assertEqual(sem_keys, queue_keys) + # Each value is an actual queue instance, not None. + for qname in sem_keys: + self.assertIsNotNone(impl._output_queue_by_name[qname]) + + def test_output_queue_by_name_captures_actual_pipeline_queue(self) -> None: + """Phase D: the queue stored in ``_output_queue_by_name[qname]`` + is the SAME instance that the pipeline build actually uses as + the OUTPUT queue of stage ``qname``. We verify two structural + properties: + - Each stage's queue is an :py:class:`asyncio.Queue` subclass + instance (i.e., a real pipeline queue, not a sentinel). + - Each stage gets a UNIQUE queue instance (by id), so the + controller can address each stage's queue independently. + """ + import asyncio + + config = PipelineConfig( + src=SourceConfig(iter(range(4))), + pipes=[ + Pipe(_stage_1, concurrency=2, name="first"), + Pipe(_stage_2, concurrency=2, name="middle"), + Pipe(_stage_3, concurrency=2, name="last"), + ], + sink=SinkConfig(buffer_size=4), + ) + pipeline = build_pipeline( + config, + num_threads=4, + _install_semaphores_for_test=True, + ) + impl = pipeline._impl + queues = [ + impl._output_queue_by_name["first"], + impl._output_queue_by_name["middle"], + impl._output_queue_by_name["last"], + ] + # Real queues, not sentinels. + for q in queues: + self.assertIsInstance(q, asyncio.Queue) + # Each stage has a unique queue instance — the controller can + # address each stage's lap stats independently. + queue_ids = {id(q) for q in queues} + self.assertEqual(len(queue_ids), 3) + + +class AdmissionGateBranchSemanticsTest(unittest.TestCase): + """V5.6: when a semaphore is registered, its value (not args.concurrency) + governs the admission cap — the static gate is REPLACED, not augmented. + + We verify the wired path produces correct end-to-end output. Stronger + in-flight-cap assertions live in the tests that import ``_pipe`` / + ``_pipe_with_semaphore`` directly. + """ + + def test_pipeline_with_semaphores_produces_correct_output(self) -> None: + config = PipelineConfig( + src=SourceConfig(iter(range(20))), + pipes=[ + Pipe(_stage_1, concurrency=2, name="add_one"), + Pipe(_stage_2, concurrency=2, name="double"), + ], + sink=SinkConfig(buffer_size=8), + ) + pipeline = build_pipeline( + config, num_threads=4, _install_semaphores_for_test=True + ) + with pipeline.auto_stop(): + results = sorted(pipeline.get_iterator()) + # f(x) = (x + 1) * 2 for x in range(20). + expected = sorted((x + 1) * 2 for x in range(20)) + self.assertEqual(results, expected) + + def test_semaphore_value_caps_in_flight(self) -> None: + """Sanity: in-flight count never exceeds the semaphore's value. + + Build a slow stage (sleeps) with concurrency=3, push more items + than the cap, and observe the peak in-flight count from the work + function itself. + """ + in_flight: list[int] = [0] + peak: list[int] = [0] + # The slow op runs on worker threads; we need atomic + # increments/decrements to observe the in-flight count + # accurately. This is test-instrumentation only; pipeline code + # itself uses no locks. + import threading + + lock: threading.Lock = threading.Lock() + + def slow_op(x: int) -> int: + with lock: + in_flight[0] += 1 + if in_flight[0] > peak[0]: + peak[0] = in_flight[0] + time.sleep(0.005) + with lock: + in_flight[0] -= 1 + return x + + config = PipelineConfig( + src=SourceConfig(iter(range(50))), + pipes=[ + Pipe(slow_op, concurrency=3, name="slow"), + ], + sink=SinkConfig(buffer_size=8), + ) + pipeline = build_pipeline( + config, num_threads=8, _install_semaphores_for_test=True + ) + with pipeline.auto_stop(): + count = 0 + for _ in pipeline.get_iterator(): + count += 1 + self.assertEqual(count, 50) + # The semaphore (set to args.concurrency=3) caps in-flight at 3. + self.assertLessEqual(peak[0], 3) + # Sanity: we did get parallelism (peak should reach 2 or 3). + self.assertGreaterEqual(peak[0], 2) + + +class StatsQueueLastLapStatsCacheTest(unittest.TestCase): + """Phase E: ``StatsQueue._log_interval_stats()`` caches the freshly + computed :py:class:`QueuePerfStats` on ``self._last_lap_stats`` so + non-callback readers (e.g., the LCA + ``DomeVideoConcurrencyController``) can observe the latest interval + without calling the destructive ``_get_lap_stats()``. + + These tests exercise the cache mechanism in isolation — they do NOT + drive a pipeline. The end-to-end check (controller reading non-None + ``queue_stats`` from a real pipeline) lives in + ``test_concurrency_controller.py``. + """ + + def _make_queue(self) -> StatsQueue: + info = StageInfo(pipeline_id=0, stage_id="0", stage_name="test") + return StatsQueue(info, buffer_size=4) + + def test_last_lap_stats_starts_as_none(self) -> None: + """Before the first interval fires, the cache MUST be ``None`` + (the controller's ``getattr(..., None)`` default would also + return ``None``, but we want a real attribute so static + analysers can see the type). + """ + queue = self._make_queue() + self.assertIsNone(queue._last_lap_stats) + + def test_log_interval_stats_populates_cache(self) -> None: + """Calling ``_log_interval_stats()`` once writes the freshly + computed ``QueuePerfStats`` to ``self._last_lap_stats``. + """ + queue: StatsQueue = self._make_queue() + # ``_get_lap_stats()`` reads ``self._lap_t0`` against + # ``time.monotonic()``; seed it so ``elapsed`` is positive. + queue._lap_t0 = time.monotonic() - 1.0 + + async def _run() -> None: + await queue._log_interval_stats() + + asyncio.run(_run()) + + self.assertIsNotNone(queue._last_lap_stats) + self.assertIsInstance(queue._last_lap_stats, QueuePerfStats) + + def test_log_interval_stats_overwrites_cache(self) -> None: + """A second ``_log_interval_stats()`` call replaces the cached + stats with the latest interval — verifies the cache is a + single-slot snapshot, not an accumulator. + """ + queue: StatsQueue = self._make_queue() + queue._lap_t0 = time.monotonic() - 1.0 + + async def _run_twice() -> tuple[QueuePerfStats, QueuePerfStats]: + await queue._log_interval_stats() + first = queue._last_lap_stats + assert first is not None + # Force a measurable elapsed delta on the second lap. + queue._lap_t0 = time.monotonic() - 0.5 + await queue._log_interval_stats() + second = queue._last_lap_stats + assert second is not None + return first, second + + first, second = asyncio.run(_run_twice()) + # Both cached values are real ``QueuePerfStats`` instances; the + # second call replaced the first (different object identity). + self.assertIsNot(first, second) 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/resize_concurrency_test.py b/tests/pipeline/resize_concurrency_test.py new file mode 100644 index 000000000..47ece5e1a --- /dev/null +++ b/tests/pipeline/resize_concurrency_test.py @@ -0,0 +1,327 @@ +# 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 + +"""Diff 3b tests: ``Pipeline._resize_concurrency_async`` (INTERNAL). + +Per DESIGN v6 / Decision 13, the public foreground-thread +``Pipeline.resize_concurrency`` sync wrapper, ``Pipeline.list_stages`` +debug helper, and the lifecycle gating that supported them are dropped. +Only the in-loop async source-of-truth remains. The intended caller is +an in-loop adaptive-concurrency controller running as a +:py:class:`BackgroundTask`. Coverage: + +- Async resize from a :py:class:`BackgroundTask` mutates the registered + semaphore and the sibling ``_dynamic_concurrency`` dict atomically + (single asyncio turn — no awaits between the two writes). +- ``KeyError`` on unknown stage; valid names are listed in the error. +- ``ValueError`` on ``new_value < 1``. +- Continuous-mode regression: the in-loop async resize takes effect + across an epoch sentinel boundary. +""" + +import asyncio +import inspect +import threading +import time +import unittest +from collections.abc import Iterator +from typing import Any + +import later.unittest +from spdl.pipeline import build_pipeline, Pipeline +from spdl.pipeline._bg_task import BackgroundTask +from spdl.pipeline.defs import Pipe, PipelineConfig, SinkConfig, SourceConfig + + +def _identity(x: int) -> int: + return x + + +class AsyncResizeFromBackgroundTaskTest(later.unittest.TestCase): + """Async resize via a BackgroundTask updates semaphore + dict.""" + + async def test_async_resize_from_bg_task_succeeds(self) -> None: + # Capture both the semaphore.max_value AND the sibling + # _dynamic_concurrency entry to confirm both writes landed in + # the same event-loop turn. + applied: dict[str, int] = {"sem": -1, "dict": -1} + done: threading.Event = threading.Event() + + class ResizerTask(BackgroundTask): + def __init__(self, pipeline_ref: list[Any]) -> None: + self._pipeline_ref = pipeline_ref + + async def run(self) -> None: + await asyncio.sleep(0.05) + pipeline = self._pipeline_ref[0] + await pipeline._resize_concurrency_async("s1", 7) + applied["sem"] = pipeline._impl._semaphore_registry["s1"].max_value + applied["dict"] = pipeline._impl._dynamic_concurrency["s1"] + done.set() + + pipeline_ref: list[Any] = [None] + + config = PipelineConfig( + src=SourceConfig(iter(range(1000))), + pipes=[Pipe(_identity, concurrency=3, name="s1")], + sink=SinkConfig(buffer_size=8), + ) + pipeline = build_pipeline( + config, + num_threads=4, + background_tasks=[lambda: ResizerTask(pipeline_ref)], + _install_semaphores_for_test=True, + ) + pipeline_ref[0] = pipeline + + with pipeline.auto_stop(): + it = pipeline.get_iterator() + for _ in range(10): + next(it) + await asyncio.get_running_loop().run_in_executor(None, done.wait, 5.0) + + self.assertTrue(done.is_set(), "Resizer task did not run") + self.assertEqual(applied["sem"], 7) + self.assertEqual(applied["dict"], 7) + + +class AsyncResizeValidationTest(later.unittest.TestCase): + """``_resize_concurrency_async`` raises on invalid arguments.""" + + async def _drive_with_bg_task( + self, + config: PipelineConfig[int], + bg_body: Any, # async callable taking the pipeline + result_holder: dict[str, BaseException | None], + ) -> None: + invoked: threading.Event = threading.Event() + pipeline_ref: list[Any] = [None] + + class CheckerTask(BackgroundTask): + def __init__(self, ref: list[Any]) -> None: + self._ref = ref + + async def run(self) -> None: + await asyncio.sleep(0.05) + try: + await bg_body(self._ref[0]) + except Exception as e: + result_holder["err"] = e + finally: + invoked.set() + + pipeline = build_pipeline( + config, + num_threads=4, + background_tasks=[lambda: CheckerTask(pipeline_ref)], + _install_semaphores_for_test=True, + ) + pipeline_ref[0] = pipeline + + with pipeline.auto_stop(): + it = pipeline.get_iterator() + for _ in range(10): + next(it) + await asyncio.get_running_loop().run_in_executor(None, invoked.wait, 5.0) + + self.assertTrue(invoked.is_set(), "Checker task did not run") + + async def test_unknown_qualified_name_raises_keyerror(self) -> None: + result_holder: dict[str, BaseException | None] = {"err": None} + + async def body(pipeline: Pipeline[int]) -> None: + await pipeline._resize_concurrency_async("nope", 4) + + config = PipelineConfig( + src=SourceConfig(iter(range(1000))), + pipes=[Pipe(_identity, concurrency=2, name="s1")], + sink=SinkConfig(buffer_size=4), + ) + await self._drive_with_bg_task(config, body, result_holder) + err = result_holder["err"] + self.assertIsInstance(err, KeyError) + # KeyError repr wraps the message in quotes, so str(ex) shows it. + self.assertIn("'s1'", str(err)) + + async def test_new_value_zero_raises_valueerror(self) -> None: + result_holder: dict[str, BaseException | None] = {"err": None} + + async def body(pipeline: Pipeline[int]) -> None: + await pipeline._resize_concurrency_async("s1", 0) + + config = PipelineConfig( + src=SourceConfig(iter(range(1000))), + pipes=[Pipe(_identity, concurrency=2, name="s1")], + sink=SinkConfig(buffer_size=4), + ) + await self._drive_with_bg_task(config, body, result_holder) + self.assertIsInstance(result_holder["err"], ValueError) + + async def test_new_value_negative_raises_valueerror(self) -> None: + result_holder: dict[str, BaseException | None] = {"err": None} + + async def body(pipeline: Pipeline[int]) -> None: + await pipeline._resize_concurrency_async("s1", -1) + + config = PipelineConfig( + src=SourceConfig(iter(range(1000))), + pipes=[Pipe(_identity, concurrency=2, name="s1")], + sink=SinkConfig(buffer_size=4), + ) + await self._drive_with_bg_task(config, body, result_holder) + self.assertIsInstance(result_holder["err"], ValueError) + + +class AsyncResizeAtomicityTest(unittest.TestCase): + """``_resize_concurrency_async`` is one event-loop turn (no awaits). + + The body's atomicity claim is a STRUCTURAL invariant: there is no + ``await`` between :py:meth:`ResizableSemaphore.resize` and the + ``_dynamic_concurrency`` dict assignment, so asyncio cannot schedule + any other coroutine between the two writes. We assert that by + inspecting the source for a single ``await`` token (the dispatch from + the caller is the only entry). + """ + + def test_no_await_between_semaphore_resize_and_dict_assignment( + self, + ) -> None: + source = inspect.getsource(Pipeline._resize_concurrency_async) + # Locate the two writes; assert no `await` keyword sits between + # them. ``sem.resize(`` and ``self._impl._dynamic_concurrency[`` + # are unique anchors in the body. + resize_index = source.index("sem.resize(") + dict_index = source.index("self._impl._dynamic_concurrency[") + self.assertLess(resize_index, dict_index) + between = source[resize_index:dict_index] + # Tokenise on whitespace boundaries to avoid matching ``await`` as + # a substring of a longer identifier (defensive — there are none + # today). + for token in between.split(): + self.assertNotEqual( + token, + "await", + "Found `await` between sem.resize() and dict assignment " + "— atomicity invariant of _resize_concurrency_async is " + "broken. The controller's _apply_decision contract " + "depends on this method completing in one asyncio turn.", + ) + + +class AsyncResizeContinuousModeTest(unittest.TestCase): + """In-loop async resize works across continuous-mode epoch boundaries. + + Continuous mode emits epoch-end sentinels but never EOF; the pipeline + runs until ``pipeline.stop()``. Verify: + (a) No crash on epoch sentinel propagation when the registry is non-empty. + (b) Resize at iteration N takes effect for iterations > N. + (c) In-flight items at the time of resize complete normally. + """ + + EPOCH_SIZE: int = 50 + NUM_EPOCHS: int = 3 + INITIAL_CONCURRENCY: int = 2 + RESIZE_TO: int = 5 + + def test_async_resize_works_across_epoch_boundary(self) -> None: + # Track in-flight count; reset between observation windows. + in_flight: list[int] = [0] + peak: list[int] = [0] + observe_lock: threading.Lock = threading.Lock() + resize_event: threading.Event = threading.Event() + resize_done: threading.Event = threading.Event() + + def slow_op(x: int) -> int: + with observe_lock: + in_flight[0] += 1 + if in_flight[0] > peak[0]: + peak[0] = in_flight[0] + time.sleep(0.005) + with observe_lock: + in_flight[0] -= 1 + return x + + def epoch_source() -> Iterator[int]: + for _ in range(self.NUM_EPOCHS): + yield from range(self.EPOCH_SIZE) + # Continuous-mode pipeline auto-emits an epoch sentinel + # at the end of each pass over the source iterator. + + # The resize is driven from inside the pipeline event loop via + # a BackgroundTask — that's the only supported caller for the + # async resize API. + class ResizerTask(BackgroundTask): + def __init__(self, pipeline_ref: list[Any]) -> None: + self._pipeline_ref = pipeline_ref + + async def run(self) -> None: + # Wait until the foreground signals it is ready for the + # resize (after the first epoch has been drained). + while not resize_event.is_set(): + await asyncio.sleep(0.01) + pipeline = self._pipeline_ref[0] + await pipeline._resize_concurrency_async( + "slow", AsyncResizeContinuousModeTest.RESIZE_TO + ) + resize_done.set() + + pipeline_ref: list[Any] = [None] + + config = PipelineConfig( + src=SourceConfig(epoch_source(), continuous=True), + pipes=[ + Pipe( + slow_op, + concurrency=self.INITIAL_CONCURRENCY, + name="slow", + ) + ], + sink=SinkConfig(buffer_size=4), + ) + pipeline = build_pipeline( + config, + num_threads=8, + background_tasks=[lambda: ResizerTask(pipeline_ref)], + _install_semaphores_for_test=True, + ) + pipeline_ref[0] = pipeline + + with pipeline.auto_stop(): + it = pipeline.get_iterator() + + # (a) First epoch: confirm initial cap is honoured. + for _ in range(self.EPOCH_SIZE): + next(it) + self.assertLessEqual(peak[0], self.INITIAL_CONCURRENCY) + + # (b)+(c) Trigger resize from the BG task and wait for it. + with observe_lock: + peak[0] = 0 + resize_event.set() + self.assertTrue( + resize_done.wait(timeout=5.0), + "Resizer task did not complete the in-loop resize", + ) + for _ in range(self.EPOCH_SIZE): + next(it) + self.assertLessEqual(peak[0], self.RESIZE_TO) + self.assertGreater( + peak[0], + self.INITIAL_CONCURRENCY, + "Expected resize to allow more concurrency.", + ) + + # (a) Run a third epoch to confirm no crash on sentinel propagation. + with observe_lock: + peak[0] = 0 + for _ in range(self.EPOCH_SIZE): + next(it) + self.assertLessEqual(peak[0], self.RESIZE_TO) + + # The registry survives epoch sentinel propagation. + self.assertIn("slow", pipeline._impl._semaphore_registry) 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()