perf(workflow): qualify integrated encoder ensemble

components/src/dynamo/vllm/multimodal_utils/custom_encoder/async_encoder.py

@@ -15,9 +15,9 @@
   submits if **all** succeed; on any failure it submits nothing (zero GPU work)
   and raises the request-level error, so a text-only LM never sees a partial
   result;
-- handing the preprocessed items (with their off-thread-computed scalar ``cost``)
-  to a ``ThreadedMicroBatcher``, which coalesces across concurrent ``encode`` calls
-  by cost and runs ``backend.forward_batch`` on the single actor thread.
+- handing the preprocessed items (with their scalar ``cost`` and opaque
+  compatibility key) to a ``ThreadedMicroBatcher``, which batches concurrent
+  ``encode`` calls by key, cost, and item count on the single actor thread.
 
 The backend's ``build`` runs on the batcher's actor thread (so a CUDA graph it
 captures is replayed on the same thread) and its ``close`` runs there at
@@ -118,6 +118,7 @@ def load(self, model_id: str) -> None:
             self._batcher = ThreadedMicroBatcher(
                 self._backend.forward_batch,
                 max_batch_cost=self._backend.max_batch_cost,
+                max_batch_items=self._backend.max_batch_items,
                 on_start=lambda: self._backend.build(model_id),
                 on_stop=self._backend.close,
                 name=self._name,
@@ -176,7 +177,8 @@ async def encode(self, raws: List[RawT]) -> List[ArtifactT]:
         preprocessed: List[Preprocessed[ItemT]] = settled  # type: ignore[assignment]
         items = [p.item for p in preprocessed]
         costs = [p.cost for p in preprocessed]
-        return await self._batcher.submit(items, costs)
+        bucket_keys = [p.bucket_key for p in preprocessed]
+        return await self._batcher.submit(items, costs, bucket_keys)
 
     def shutdown(self) -> None:
         """Stop the actor thread (running ``backend.close`` on it) and the

components/src/dynamo/vllm/multimodal_utils/custom_encoder/backend/base.py

@@ -21,10 +21,11 @@
   to peak; if ``buckets`` is set (once CUDA-graph batching is supported), capture
   one CUDA graph per rung here so it is bound to the thread that later replays it
   in ``forward_batch``. Pick the device yourself (``"cuda"`` / the current device).
-- ``preprocess(raw) -> Preprocessed{item, cost}`` — **off the actor thread,
+- ``preprocess(raw) -> Preprocessed{item, cost, bucket_key}`` — **off the actor thread,
   concurrent.** Deterministic, thread-safe, CUDA-free (fetch / resize / patchify
   on CPU/pinned memory). ``cost`` is a **scalar** — how much the item adds toward
-  ``max_batch_cost`` (e.g. its visual-token count). Raise to reject a bad input —
+  ``max_batch_cost`` (e.g. its patch count). ``bucket_key`` is an opaque
+  compatibility partition; items with different keys are never mixed. Raise to reject a bad input —
   it fails only that image, before any GPU work. **Off by default:** override
   ``preprocess`` *and* set ``preprocess_concurrency > 0`` together to enable this
   pool. With the defaults (identity passthrough, ``preprocess_concurrency = 0``)
@@ -44,16 +45,17 @@
 - ``max_batch_cost`` — the scalar dispatch ceiling the batcher packs up to; a
   *chosen* budget (a token budget when ``cost`` is a token count). ``None`` (the
   default) ⇒ **pass-through**: no cap (the author owns sizing).
+- ``max_batch_items`` — an optional independent item-count ceiling for backends
+  with finite graph-batch capacity. ``None`` leaves the count unconstrained.
 - ``buckets`` — sorted graph ladder, forward-compatible (unused until CUDA-graph
   batching is supported). ``None``/empty ⇒ eager.
 - ``preprocess_concurrency`` — size of the off-thread pool Dynamo runs
   ``preprocess`` on. ``0`` (the **default**) ⇒ no preprocess phase: raws go
   straight to ``forward_batch``. Set ``> 0`` (with an overridden ``preprocess``)
   for off-loop fetch / resize / patchify.
 
-Batching is **one-dimensional**: Dynamo packs by scalar ``cost`` up to
-``max_batch_cost`` and never inspects item shape — the author owns any
-shape/padding concerns inside ``forward_batch``.
+Inside each opaque compatibility partition, Dynamo packs by scalar ``cost`` and
+item count. It never inspects item shape; the author owns the key and padding.
 
 Raising errors
 --------------
@@ -93,7 +95,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Generic, List, Optional, Sequence, TypeVar
+from typing import Generic, Hashable, List, Optional, Sequence, TypeVar
 
 RawT = TypeVar("RawT")  # raw input the author preprocesses (e.g. an image URL)
 ItemT = TypeVar("ItemT")  # opaque payload preprocess() hands to forward_batch()
@@ -102,7 +104,7 @@
 
 @dataclass(frozen=True)
 class Preprocessed(Generic[ItemT]):
-    """The result of ``preprocess(raw)``: an opaque item plus its batching cost.
+    """The result of ``preprocess(raw)``: an opaque item plus batching policy.
 
     ``cost`` is computed **once, off the actor thread**, so the batcher never
     evaluates model policy (it stays torch-free) and packs purely by this scalar.
@@ -114,10 +116,14 @@ class Preprocessed(Generic[ItemT]):
             **pass-through mode** (``max_batch_cost`` is ``None``) the batcher
             never reads it, so a pass-through author can leave it at the default
             ``1``.
+        bucket_key: Optional opaque compatibility partition. Items with different
+            keys are never passed to the same ``forward_batch`` call. ``None``
+            preserves the single-partition behavior.
     """
 
     item: ItemT
     cost: int = 1
+    bucket_key: Hashable | None = None
 
 
 class VisionEncoderBackend(ABC, Generic[RawT, ItemT, ArtifactT]):
@@ -127,8 +133,9 @@ class VisionEncoderBackend(ABC, Generic[RawT, ItemT, ArtifactT]):
     on a dedicated actor thread (``ThreadedMicroBatcher``) and exposes the async
     request API (``AsyncVisionEncoder``). Subclasses implement ``build`` and
     ``forward_batch``; ``preprocess`` (default identity passthrough),
-    ``max_batch_cost``, ``buckets``, and ``preprocess_concurrency`` are overridden
-    only as needed. Artifact interpretation belongs to the selected adapter.
+    ``max_batch_cost``, ``max_batch_items``, ``buckets``, and
+    ``preprocess_concurrency`` are overridden only as needed. Artifact
+    interpretation belongs to the selected adapter.
     """
 
     #: Scalar dispatch ceiling: the batcher packs items up to this summed ``cost``
@@ -137,6 +144,11 @@ class VisionEncoderBackend(ABC, Generic[RawT, ItemT, ArtifactT]):
     #: (the author owns sizing; ``cost`` is ignored).
     max_batch_cost: Optional[int] = None
 
+    #: Optional independent item-count ceiling per ``forward_batch`` call. Use
+    #: this for a finite graph-count capacity; keep scalar compute/resource work
+    #: in ``max_batch_cost``.
+    max_batch_items: Optional[int] = None
+
     #: Sorted graph ladder (the captured rungs), **forward-compatible** — unused
     #: until CUDA-graph batching is supported. ``None``/empty ⇒ eager.
     buckets: Optional[Sequence[int]] = None

components/src/dynamo/vllm/multimodal_utils/custom_encoder/batcher.py

@@ -14,9 +14,9 @@
 - **eager batching by default**: whenever the worker is free it drains everything
   queued and runs it as one ``fn`` call, then repeats (no timer) — a lone item
   runs the next loop, and batch size auto-scales with load. Pass ``max_batch_cost``
-  to **micro-batch** instead: the drained items are split into batches whose
-  summed per-item ``cost`` stays within that budget (one-dimensional packing by
-  ``cost`` alone; item shape is never inspected).
+  to **micro-batch** instead: the drained items are partitioned by opaque
+  compatibility key, then split by summed ``cost`` and ``max_batch_items``.
+  Item shape is never inspected. There is no timed coalescing hold.
 
 The caller speaks in opaque items plus an optional per-item scalar ``cost``
 (computed off-thread, see ``Preprocessed``), so all model knowledge stays in the
@@ -38,9 +38,10 @@
 import logging
 import queue
 import threading
+from collections import deque
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import Callable, Generic, List, Optional, TypeVar
+from typing import Callable, Deque, Generic, Hashable, List, Optional, TypeVar
 
 logger = logging.getLogger(__name__)
 
@@ -80,6 +81,7 @@ class _Work(Generic[T]):
 
     item: T
     cost: int
+    bucket_key: Hashable | None
     request: _Request
     index: int
 
@@ -93,6 +95,8 @@ class ThreadedMicroBatcher(Generic[T, R]):
         max_batch_cost: Max summed ``cost`` of a single ``fn`` batch (>= 1).
             ``None`` (default) ⇒ **pass-through**: no cap — the whole drained set
             runs as one ``fn`` call (``cost`` ignored).
+        max_batch_items: Optional independent item-count cap (>= 1). ``None``
+            leaves item count unconstrained.
         on_start: Optional callable run once on the worker thread before serving
             (model build / warmup); its failure surfaces from ``start()``.
         on_stop: Optional callable run once on the worker thread at teardown (after
@@ -107,15 +111,19 @@ def __init__(
         fn: Callable[[List[T]], List[R]],
         *,
         max_batch_cost: Optional[int] = None,
+        max_batch_items: Optional[int] = None,
         on_start: Optional[Callable[[], None]] = None,
         on_stop: Optional[Callable[[], None]] = None,
         name: str = "micro-batcher",
         join_timeout_s: float = 10.0,
     ) -> None:
         if max_batch_cost is not None and max_batch_cost < 1:
             raise ValueError("max_batch_cost must be >= 1 (or None for pass-through)")
+        if max_batch_items is not None and max_batch_items < 1:
+            raise ValueError("max_batch_items must be >= 1 (or None for no cap)")
         self._fn = fn
         self._max_batch_cost = max_batch_cost
+        self._max_batch_items = max_batch_items
         self._on_start = on_start
         self._on_stop = on_stop
         self._name = name
@@ -182,13 +190,14 @@ async def submit(
         self,
         items: List[T],
         costs: Optional[List[int]] = None,
+        bucket_keys: Optional[List[Hashable | None]] = None,
     ) -> List[R]:
         """Submit a group of items; await one result per item, in order.
 
         ``costs`` is computed off-thread by the caller (see ``Preprocessed``);
         when omitted it defaults to ``1`` per item (plain count-based batching).
-        Batching is one-dimensional — the batcher packs by ``cost`` alone and
-        never inspects item shape.
+        ``bucket_keys`` partition incompatible items before the batcher packs by
+        cost and item count. ``None`` preserves one compatibility partition.
 
         Cancelling the await tombstones this request. Work that has not passed the
         final pre-``fn`` dispatch check is dropped; a committed ``fn`` still finishes.
@@ -201,6 +210,12 @@ async def submit(
             costs = [1] * len(items)
         elif len(costs) != len(items):
             raise ValueError(f"costs has {len(costs)} entries for {len(items)} items")
+        if bucket_keys is None:
+            bucket_keys = [None] * len(items)
+        elif len(bucket_keys) != len(items):
+            raise ValueError(
+                f"bucket_keys has {len(bucket_keys)} entries for {len(items)} items"
+            )
         for c in costs:
             if not isinstance(c, int) or isinstance(c, bool) or c < 1:
                 raise ValueError(f"cost must be a positive int, got {c!r}")
@@ -209,13 +224,19 @@ async def submit(
                     f"item cost {c} exceeds max_batch_cost {self._max_batch_cost}; "
                     "it has no batch it can fit"
                 )
+        for key in bucket_keys:
+            try:
+                hash(key)
+            except TypeError as exc:
+                raise ValueError(f"bucket_key must be hashable, got {key!r}") from exc
         request: _Request = _Request(
             completion=concurrent.futures.Future(),
             results=[None] * len(items),
             remaining=len(items),
         )
         works = [
-            _Work(item, c, request, i) for i, (item, c) in enumerate(zip(items, costs))
+            _Work(item, c, key, request, i)
+            for i, (item, c, key) in enumerate(zip(items, costs, bucket_keys))
         ]
         # State check + admission + queue-commit under one lock so a concurrent
         # shutdown() cannot strand the request and capacity cannot leak.
@@ -337,7 +358,7 @@ def _collect(self) -> Optional[List[_Work]]:
         return works
 
     def _dispatch(self, works: List[_Work]) -> None:
-        """Split live items by cost budget, run ``fn`` (one-dimensional packing).
+        """Partition by compatibility key, then split by cost and item count.
 
         Tombstoned (failed / done) items are dropped before batching —
         an already-failed request never reaches ``fn``."""
@@ -349,21 +370,35 @@ def _dispatch(self, works: List[_Work]) -> None:
                 live.append(work)
         if not live:
             return
-        if self._max_batch_cost is None:
-            # Pass-through: no cost cap — the whole drained set is one batch.
-            self._run_batch(live)
-            return
-        batch: List[_Work] = []
-        batch_cost = 0
+        groups: dict[Hashable | None, Deque[_Work]] = {}
         for work in live:
-            if batch and batch_cost + work.cost > self._max_batch_cost:
+            groups.setdefault(work.bucket_key, deque()).append(work)
+
+        pending = list(groups.values())
+        while pending:
+            next_round: List[Deque[_Work]] = []
+            for group in pending:
+                batch: List[_Work] = []
+                batch_cost = 0
+                while group:
+                    work = group[0]
+                    if batch and (
+                        (
+                            self._max_batch_cost is not None
+                            and batch_cost + work.cost > self._max_batch_cost
+                        )
+                        or (
+                            self._max_batch_items is not None
+                            and len(batch) >= self._max_batch_items
+                        )
+                    ):
+                        break
+                    batch.append(group.popleft())
+                    batch_cost += work.cost
                 self._run_batch(batch)
-                batch, batch_cost = [], 0
-            batch.append(work)
-            batch_cost += work.cost
-        # `batch` always holds at least the final work here (live is non-empty and
-        # every work is appended after any mid-loop flush), so flush unconditionally.
-        self._run_batch(batch)
+                if group:
+                    next_round.append(group)
+            pending = next_round
 
     def _run_batch(self, batch: List[_Work]) -> None:
         # Re-filter immediately before fn: a request may have been failed by a

components/src/dynamo/vllm/tests/multimodal_utils/custom_encoder/backend/test_vllm_base.py

@@ -66,8 +66,7 @@ def test_preprocessed_cost_defaults_to_1():
 
 
 def test_preprocessed_has_no_bucket_key():
-    # Batching is one-dimensional (scalar cost only) — there is no bucket_key.
-    assert not hasattr(Preprocessed(item="x"), "bucket_key")
+    assert Preprocessed(item="x").bucket_key is None
 
 
 def test_abc_cannot_be_instantiated():
@@ -95,5 +94,6 @@ def test_default_attrs_and_close_noop():
     # preprocess pool.
     assert e.buckets is None
     assert e.max_batch_cost is None
+    assert e.max_batch_items is None
     assert e.preprocess_concurrency == 0
     assert e.close() is None

components/src/dynamo/vllm/tests/multimodal_utils/custom_encoder/test_vllm_async_encoder.py

@@ -46,6 +46,7 @@ def __init__(self, *, fail_on=None):
         self.close_calls = 0
         self.model_id = None
         self.forward_threads: list[int] = []
+        self.forward_batches: list[list[str]] = []
 
     def build(self, model_id):
         self.build_thread = threading.get_ident()
@@ -58,6 +59,7 @@ def preprocess(self, raw):
 
     def forward_batch(self, items, target_bucket=None):
         self.forward_threads.append(threading.get_ident())
+        self.forward_batches.append(list(items))
         return [torch.full((2, 4), float(len(str(it)))) for it in items]
 
     def close(self):
@@ -77,6 +79,23 @@ async def test_encode_returns_one_tensor_per_raw():
         enc.shutdown()
 
 
+async def test_encode_partitions_preprocessed_compatibility_keys(monkeypatch):
+    be = _FakeBackend()
+    monkeypatch.setattr(
+        be,
+        "preprocess",
+        lambda raw: Preprocessed(item=raw, cost=1, bucket_key=raw[0]),
+    )
+    enc = AsyncVisionEncoder(be)
+    enc.load("m")
+    try:
+        out = await enc.encode(["a", "bbbb", "aa"])
+        assert be.forward_batches == [["a", "aa"], ["bbbb"]]
+        assert [float(tensor[0, 0]) for tensor in out] == [1.0, 4.0, 2.0]
+    finally:
+        enc.shutdown()
+
+
 async def test_preprocess_barrier_fails_atomically_with_no_gpu_work():
     be = _FakeBackend(fail_on={"bad"})
     enc = AsyncVisionEncoder(be)