Add semaphore registry and admission gate REPLACE for adaptive concurrency (#1391)

Summary:

This is Diff 3a in the T262755626 stack. Adds the foundation for runtime concurrency adjustment in SPDL pipelines:

1. `_PipelineImpl._semaphore_registry: dict[str, ResizableSemaphore]` — sibling registry keyed on the qualified stage name. Populated at pipeline build time when a stage opts in. NO modification to `StageInfo` (preserves frozen=True / hashability for third-party code).
2. `_PipelineImpl._dynamic_concurrency: dict[str, int]` — current value per stage; updated by future `Pipeline.resize_concurrency()` calls.
3. `_PipelineImpl._stage_info_by_name: dict[str, StageInfo]` — for error messages and future Track B logging hooks.
4. `_pipe()` admission gate REPLACE (V5.6): when a `ResizableSemaphore` is registered for a stage, `await semaphore.acquire()` becomes the admission gate, REPLACING the static `len(tasks) >= concurrency` check. When no semaphore is registered (the default), behaviour is unchanged — ZERO per-task overhead added to the existing fast path. The branch is taken ONCE outside the hot loop (dispatching to `_pipe_with_semaphore`).
5. `build_pipeline(_install_semaphores_for_test: bool = False)` test-only knob: when True, every Pipe stage is built with a `ResizableSemaphore` and registered. Production code MUST NOT pass this — it bypasses the selective opt-in semantics of `Pipeline.resize_concurrency` (Diff 3b).
6. V5.5 throughput regression test (`admission_gate_perf_test.py`): 4-stage sync pipeline, N=10000 items, 5 trials. Regression > 2% on either p50 or p99 fails the diff. Includes registry wiring smoke tests and an in-flight-cap assertion.

Note (perf-reviewer WARN-1): the V5.5 perf test covers sync stages only — extending to async/sync_iter handlers is a follow-up after Diff 3a lands.

Note (perf-reviewer WARN-2): qualified stage names == `info.stage_name` for non-MultiPipe stages. LCA's MultiPipe is a single SPDL Pipe with internal dispatcher; the `_qualified_name(branch_label=...)` scheme is forward-compatible for true SPDL fan-out.

The user-facing API (`Pipeline.resize_concurrency`) lands in Diff 3b.

Differential Revision: D102929216

Author: justinvjoseph

src/spdl/pipeline/_build.py

@@ -31,6 +31,7 @@
     _get_global_id,
     _set_global_id,
     AsyncQueue,
+    ResizableSemaphore,
     StageInfo,
     TaskHook,
 )
@@ -129,6 +130,7 @@ def _build_pipeline(
     stage_id: int = 0,
     background_tasks: list[BackgroundTaskFactory] | None = None,
     use_priority_scheduler: bool = False,
+    _install_semaphores_for_test: bool = False,
 ) -> Pipeline[U]:
     if _DEFAULT_BUILD_CALLBACK is not None:
         try:
@@ -175,6 +177,18 @@ def _build_pipeline(
             lambda sched=scheduler: _PrioritySchedulerBackgroundTask(sched)
         )
 
+    # V5.1+V5.5 Diff 3a: pre-allocate the per-pipeline registries that
+    # `_components/_node.py` populates when ``_install_semaphores_for_test``
+    # is True. These are then handed off to ``_PipelineImpl.__init__`` so
+    # ``Pipeline.resize_concurrency`` (Diff 3b) 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] = {}
+
     coro, queue = _build_pipeline_coro(
         pipeline_cfg,
         max_failures=max_failures,
@@ -184,9 +198,23 @@ def _build_pipeline(
         stage_id=stage_id,
         background_tasks=all_bg_tasks or None,
         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,
     )
 
-    return Pipeline(coro, queue, executor, desc=desc)
+    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,
+    )
 
 
 def build_pipeline(
@@ -201,6 +229,7 @@ def build_pipeline(
     stage_id: int = 0,
     background_tasks: list[BackgroundTaskFactory] | None = None,
     use_priority_scheduler: bool = False,
+    _install_semaphores_for_test: bool = False,
 ) -> Pipeline[U]:
     """Build a pipeline from the config.
 
@@ -271,6 +300,16 @@ def build_pipeline(
             dispatch for sync stages via :py:class:`PriorityScheduler`.
             Deeper stages (closer to sink) are given higher priority,
             reducing pipeline bubble time.
+
+        _install_semaphores_for_test: **Test-only.** When ``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`. Used by
+            the V5.5 throughput regression test and the
+            :py:meth:`Pipeline.resize_concurrency` continuous-mode test.
+            Production code MUST NOT pass this — it bypasses the
+            selective opt-in semantics of ``Pipeline.resize_concurrency``.
     """
     from . import _profile
 
@@ -287,6 +326,7 @@ def build_pipeline(
         stage_id=stage_id,
         background_tasks=background_tasks,
         use_priority_scheduler=use_priority_scheduler,
+        _install_semaphores_for_test=_install_semaphores_for_test,
     )
 
 

src/spdl/pipeline/_components/_node.py

@@ -42,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
@@ -480,13 +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``.
+
+    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.
 
@@ -609,12 +674,49 @@ def _build_node(
                         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, 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, args, fc, hooks
                             )
@@ -700,6 +802,11 @@ def _build_node_recursive(
     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.
 
@@ -714,6 +821,22 @@ def _build_node_recursive(
         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 use
+            :py:meth:`Pipeline.resize_concurrency` (Diff 3b) instead.
+        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.
@@ -722,10 +845,33 @@ def _build_node_recursive(
         return
 
     for n in node.upstream:
-        _build_node_recursive(n, fc_class, task_hook_factory, max_failures, scheduler)
+        _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,
+        )
 
     depth = _node_depth(node) if scheduler is not None else 0
-    _build_node(node, fc_class, task_hook_factory, max_failures, scheduler, depth=depth)
+    _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
@@ -772,6 +918,11 @@ def _build_pipeline_node(
     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
@@ -786,7 +937,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, scheduler)
+    _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
 
 
@@ -1027,7 +1189,12 @@ def _build_pipeline_coro(
     stage_id: int = 0,
     background_tasks: Sequence[BackgroundTaskFactory] | None = None,
     scheduler: Any = None,
-) -> tuple[Coroutine[None, None, None], asyncio.Queue]:
+    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,
@@ -1037,6 +1204,11 @@ def _build_pipeline_coro(
             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)