[pipeline] Add batched routing to path_variants (facebookresearch#1595)

Add a `batched` option to `PathVariants` and
`PipelineBuilder.path_variants`.

When set, the router receives a whole batch (a list) and returns one
path index per element; the batch is partitioned into per-path
sub-batches (order preserved), each path processes its sub-batch as a
list, and the fan-in merge concatenates the sub-batches back into one
batch.

Per-item routing pays the router + fan-out/fan-in machinery once per
element. For workloads that route in bulk -- e.g. a cache hit/miss split
over an already-aggregated batch -- that fixed per-item cost dominates.
Batched routing amortizes it over the whole batch: the router runs once
per batch, each branch runs its ops once per batch, and one list (not
one item) crosses each queue per batch.

Semantics:
- Each input batch contributes exactly one sub-batch (possibly empty) to
every path queue, so the queues stay in lockstep and the merge
recombines an input batch's sub-batches by reading one list from each
path.
- A branch may drop elements by returning a shorter list, and must
tolerate an empty input list (a path that received no elements for a
batch).
- A batched router must return exactly one index per element; otherwise
the stage fails.

`batched` defaults to `False`, so existing per-item behavior is
byte-identical (not backward-incompatible).

Author: mthrok

src/spdl/pipeline/_builder.py

@@ -386,17 +386,30 @@ def path_variants(
         router: Callable,
         paths: Sequence,
         name: str | None = None,
+        batched: bool = False,
     ) -> "PipelineBuilder[T, U]":
         """Route items to different processing paths based on a router function.
 
+        .. versionadded:: 0.6.0
+           The ``batched`` argument.
+
         Args:
-            router: A callable that takes an item and returns an int index
-                selecting which path the item should be routed to.
+            router: A callable that selects the path for each input. In per-item
+                mode (default) it takes an item and returns an int index. In
+                ``batched`` mode it takes a batch (list) and returns one int index
+                per element (same length as the batch).
             paths: A sequence of paths, where each path is a sequence of
                 pipe configs.
             name: Optional name for the stage.
+            batched: If True, route whole batches instead of single items: the
+                batch is partitioned into per-path sub-batches (each path op takes
+                and returns a list) and merged back into one batch. Aggregate the
+                source into batches upstream of this stage. See
+                :py:func:`~spdl.pipeline.defs.PathVariants`.
         """
-        self._process_args.append(PathVariants(router, paths, name=name))
+        self._process_args.append(
+            PathVariants(router, paths, name=name, batched=batched)
+        )
         return self
 
     def add_sink(

src/spdl/pipeline/_components/_node.py

@@ -46,7 +46,7 @@
 from ._sink import _sink
 from ._source import _source, _source_continuous
 from ._subprocess_pipe import _subprocess_pipeline
-from ._variants import _path_variants_router
+from ._variants import _batched_path_variants_merge, _path_variants_router
 
 T = TypeVar("T")
 S = TypeVar("S")
@@ -69,14 +69,17 @@
 # pyre-strict
 
 
+@dataclass
 class _PathVariantsMergeConfig:
     """Internal config for the fan-in merge node of PathVariants.
 
     This is not user-facing. It is used as the ``cfg`` of the merge ``_Node``
     so that ``_build_node`` can dispatch to the correct coroutine builder.
     """
 
-    pass
+    batched: bool = False
+    """Whether the parent PathVariants stage routes batches (see
+    :py:func:`~spdl.pipeline._components._variants._batched_path_variants_merge`)."""
 
 
 # Used to express the upstream relation ship of coroutines,
@@ -434,7 +437,7 @@ def _convert_path_variants(
     merge_out_q = q_class(merge_info, buffer_size=_BUFFER_SIZE)
     merge_node = _FanInNode(
         merge_info,
-        _PathVariantsMergeConfig(),
+        _PathVariantsMergeConfig(batched=cfg.batched),
         end_nodes,
         input_queues=[n.output_queue for n in end_nodes],
         output_queue=merge_out_q,
@@ -570,6 +573,7 @@ def _build_node(
                 node.output_queues,
                 node.cfg.router,
                 hooks,
+                batched=node.cfg.batched,
             )
         case _FanInNode():
             hooks = task_hook_factory(node.info)
@@ -585,8 +589,16 @@ def _build_node(
                         cfg.op,
                     )
                 case _PathVariantsMergeConfig():
+                    # A batched stage recombines each input batch's per-path
+                    # sub-batches into one batch; the per-item stage passes items
+                    # through in arrival order (default merge).
                     node._coro = _merge(
-                        node.info, node.input_queues, node.output_queue, fc, hooks, None
+                        node.info,
+                        node.input_queues,
+                        node.output_queue,
+                        fc,
+                        hooks,
+                        _batched_path_variants_merge if cfg.batched else None,
                     )
                 case _:  # pragma: no cover
                     raise ValueError(

src/spdl/pipeline/_components/_variants.py

@@ -9,6 +9,7 @@
 # pyre-strict
 
 __all__ = [
+    "_batched_path_variants_merge",
     "_path_variants_router",
 ]
 
@@ -21,17 +22,20 @@
 
 from spdl.pipeline._common._convert import _to_async
 
-from ._common import _EOF, _EPOCH_END, _ShieldedHook, is_eof, is_epoch_end
+from ._common import _EOF, _EPOCH_END, _ShieldedHook, is_eof, is_epoch_end, StageInfo
 from ._hook import _stage_hooks, TaskHook
 from ._queue import AsyncQueue
 
 _LG: logging.Logger = logging.getLogger(__name__)
 
 
 def _make_async_router(
-    router: Callable[[Any], int] | Callable[[Any], Awaitable[int]],
-) -> Callable[[Any], Awaitable[int]]:
-    """Wrap a sync router with run_in_executor; pass through async routers."""
+    router: Callable[[Any], Any] | Callable[[Any], Awaitable[Any]],
+) -> Callable[[Any], Awaitable[Any]]:
+    """Wrap a sync router with run_in_executor; pass through async routers.
+
+    The router returns an ``int`` (per-item mode) or a ``Sequence[int]`` (batched
+    mode); this wrapper is agnostic to which."""
     if inspect.iscoroutinefunction(router):
         return router
     call = getattr(router, "__call__", None)
@@ -88,14 +92,25 @@ async def _queue_stage_hook(queues: Sequence[AsyncQueue]) -> AsyncGenerator[None
 def _path_variants_router(
     input_queue: AsyncQueue,
     path_queues: Sequence[AsyncQueue],
-    router: Callable[[Any], int] | Callable[[Any], Awaitable[int]],
+    router: Callable[[Any], Any] | Callable[[Any], Awaitable[Any]],
     task_hooks: list[TaskHook],
+    batched: bool = False,
 ) -> Coroutine[None, None, None]:
     """Create a coroutine that routes items to per-path queues.
 
-    The router reads items from ``input_queue``, calls ``router(item)`` to
-    determine the target path index, and puts the item on the corresponding
-    queue in ``path_queues``.
+    In the default (per-item) mode the router reads items from ``input_queue``,
+    calls ``router(item)`` to determine the target path index, and puts the item
+    on the corresponding queue in ``path_queues``.
+
+    In ``batched`` mode each item read from ``input_queue`` is a whole batch (a
+    list). ``router(batch)`` returns one path index per element; the batch is
+    partitioned into per-path sub-batches (preserving order) and each sub-batch
+    is put on its path's queue as a single list. A sub-batch is put on **every**
+    path queue -- an empty list where a path received no elements -- so that each
+    input batch contributes exactly one list to every path queue. This keeps the
+    downstream fan-in merge in lockstep, letting it recombine the sub-batches of
+    one input batch by reading one list from each path (see
+    :py:func:`_batched_path_variants_merge`).
 
     Sync routers are wrapped with ``run_in_executor`` to avoid blocking the
     event loop. Async routers are awaited directly.
@@ -111,14 +126,47 @@ def _path_variants_router(
     Args:
         input_queue: The queue to consume items from.
         path_queues: Per-path output queues, one per variant path.
-        router: Callable that maps an item to a path index.
+        router: Callable that maps an item (or a batch, if ``batched``) to a path
+            index (or a sequence of per-element path indices, if ``batched``).
         task_hooks: Hooks for monitoring.
+        batched: If True, route whole batches (see above) instead of single items.
 
     Returns:
         A coroutine that executes the router stage.
     """
     num_paths: int = len(path_queues)
-    arouter: Callable[[Any], Awaitable[int]] = _make_async_router(router)
+    arouter: Callable[[Any], Awaitable[Any]] = _make_async_router(router)
+
+    async def _route_item(item: Any) -> None:
+        idx = await arouter(item)
+        if idx < 0 or idx >= num_paths:
+            raise IndexError(
+                f"Router returned index {idx}, but there are only "
+                f"{num_paths} paths (valid range: [0, {num_paths}))."
+            )
+        await path_queues[idx].put(item)
+
+    async def _route_batch(batch: Any) -> None:
+        indices = await arouter(batch)
+        if len(indices) != len(batch):
+            raise ValueError(
+                f"Batched router returned {len(indices)} indices for a batch of "
+                f"{len(batch)} items; it must return exactly one index per item."
+            )
+        parts: list[list[Any]] = [[] for _ in range(num_paths)]
+        for item, idx in zip(batch, indices):
+            if idx < 0 or idx >= num_paths:
+                raise IndexError(
+                    f"Router returned index {idx}, but there are only "
+                    f"{num_paths} paths (valid range: [0, {num_paths}))."
+                )
+            parts[idx].append(item)
+        # Emit to every path (empty lists included) so each input batch contributes
+        # exactly one list per path, keeping the fan-in merge in lockstep.
+        for q, part in zip(path_queues, parts):
+            await q.put(part)
+
+    _route: Callable[[Any], Awaitable[None]] = _route_batch if batched else _route_item
 
     @_queue_stage_hook(path_queues)
     # pyrefly: ignore [not-callable]
@@ -134,12 +182,38 @@ async def _router() -> None:
                     await q.put(_EPOCH_END)
                 continue
 
-            idx = await arouter(item)
-            if idx < 0 or idx >= num_paths:
-                raise IndexError(
-                    f"Router returned index {idx}, but there are only "
-                    f"{num_paths} paths (valid range: [0, {num_paths}))."
-                )
-            await path_queues[idx].put(item)
+            await _route(item)
 
     return _router()
+
+
+async def _batched_path_variants_merge(
+    info: StageInfo,
+    input_queues: Sequence[asyncio.Queue],
+    output_queue: asyncio.Queue,
+) -> None:
+    """Fan-in merge for a ``batched`` path_variants stage.
+
+    Each input batch was partitioned by the router into one sub-batch (list) per
+    path, so the queues stay in lockstep: reading one list from every input queue
+    yields the sub-batches of a single original batch. They are concatenated (in
+    path order) back into one batch and emitted, so the stage is batch-in /
+    batch-out and downstream sees whole batches -- not the individual items.
+
+    A fully-empty recombined batch (every path routed nothing, or every element
+    was dropped) is skipped rather than emitted downstream.
+
+    ``_EPOCH_END`` and ``_EOF`` are broadcast by the router to every path queue,
+    so they surface on all input queues together; the merge forwards one epoch
+    boundary and stops on end-of-stream.
+    """
+    while True:
+        items = [await q.get() for q in input_queues]
+        if any(is_eof(it) for it in items):
+            return
+        if any(is_epoch_end(it) for it in items):
+            await output_queue.put(_EPOCH_END)
+            continue
+        combined = [item for sub_batch in items for item in sub_batch]
+        if combined:
+            await output_queue.put(combined)