feat(multimodal): ThreadedMicroBatcher + batcher-backed AsyncVisionEncoder (cross-request batching)

PR3 of the custom vision-encoder series, stacked on PR2. Replaces the serial
direct-call glue with cross-request batching. handlers.py is unchanged — the
AsyncVisionEncoder public surface is the stable boundary.

- ThreadedMicroBatcher (generic, torch-free): one pinned non-daemon actor thread
  running on_start (build) / every fn / on_stop (close); eager drain-on-completion
  (no timer); coalesce by cost within bucket_key. max_batch_cost is optional —
  None (default) = pass-through (no cap; the whole drained bucket runs as one fn
  call), an int caps the batch cost. Per-item idempotent finalizer; cancellation +
  repeat-cancel retirement; worker supervisor; optional max_outstanding_cost
  admission; cross-loop-safe shutdown. submit(items, costs, bucket_keys) —
  cost/bucket_key precomputed off-thread (A4).
- AsyncVisionEncoder now routes preprocessed items through the batcher (was a
  single-worker serial executor); the A5 preprocess barrier is unchanged.

No graph ladder yet (batches to the cost budget, or passes through) — the
buckets/target_bucket machinery lands in PR4.

Tests: batcher concurrency suite (eager-drain, cost budget, pass-through,
bucket_key isolation, cancellation/retirement, supervisor, lifecycle, on_stop) +
batcher-backed glue (A5 barrier, single actor thread, fail-fast load + pool reaping).

Note: based on the series' validated merge-base; rebase onto main before merge.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: furionw

components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py

@@ -1,27 +1,28 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Serial async glue (L3) between the worker's event loop and a ``VisionEncoderBackend``.
-
-This is the **eager-milestone** glue: it proves the splice path end to end with a
-**direct call — no micro-batcher**. Per request it preprocesses the images off the
-event loop, enforces request-level atomicity, and runs the author's
-``forward_batch`` on a single dedicated **actor thread**, serialized — there is no
-cross-request coalescing. A follow-up swaps this body for a
-``ThreadedMicroBatcher`` (cross-request batching); the public surface
-(``load`` / ``encode`` / ``get_image_placeholder_token_id`` / ``shutdown``) is
-identical, so the worker integration does not change.
-
-Why a single actor thread (not ``asyncio.to_thread``): build and every
-``forward_batch`` run on the **same** thread, so an author that captures a CUDA
-graph in ``build`` can replay it from ``forward_batch`` — the affinity the batched
-version also guarantees. ``max_workers=1`` serializes forwards (FIFO), so
-concurrent ``encode`` calls run one forward at a time without interleaving.
-
-Request-level atomicity (design A5): a gather-barrier sits between preprocess and
-the forward — ``encode`` waits for *every* image's preprocess to settle and runs
-the forward only if **all** succeed; on any failure it does no GPU work and raises
-the request-level error, so a text-only LM never sees a partial result.
+"""Async glue (L3) between the worker's event loop and a ``VisionEncoderBackend``.
+
+``AsyncVisionEncoder`` is the **Dynamo-owned** layer the worker talks to. It
+turns the author's synchronous, thread-affine backend (L2) into an awaitable
+``encode(raws) -> list[tensor]`` by:
+
+- running ``backend.preprocess`` **off the event loop** on a bounded
+  ``ThreadPoolExecutor`` (CPU-heavy fetch / resize / patchify must not serialize
+  on the GPU actor thread);
+- enforcing **request-level atomicity** (A5): a gather-barrier between preprocess
+  and submit — ``encode`` waits for *every* image's preprocess to settle and only
+  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 ``cost`` /
+  ``bucket_key``) to a ``ThreadedMicroBatcher``, which **coalesces across
+  concurrent ``encode`` calls** up to the backend's ``max_batch_cost`` token
+  budget and runs ``backend.forward_batch`` on the single actor thread.
+
+This replaces the earlier serial (direct-call) glue with cross-request batching;
+the public surface (``load`` / ``encode`` / ``get_image_placeholder_token_id`` /
+``shutdown``) is unchanged, so the worker integration does not move.
 """
 
 from __future__ import annotations
@@ -33,6 +34,7 @@
 
 import torch
 
+from dynamo.vllm.multimodal_utils.threaded_micro_batcher import ThreadedMicroBatcher
 from dynamo.vllm.multimodal_utils.vision_encoder_backend import (
     ItemT,
     Preprocessed,
@@ -44,12 +46,11 @@
 
 
 class AsyncVisionEncoder(Generic[RawT, ItemT]):
-    """Drive a ``VisionEncoderBackend`` from the async request path, serially.
+    """Drive a ``VisionEncoderBackend`` from the worker's async request path.
 
     The worker calls ``load`` once at startup and ``await``s ``encode`` per
     request; ``shutdown`` on teardown. All model knowledge lives in ``backend``;
-    this class owns the preprocess pool, the A5 barrier, and the single actor
-    thread that runs ``build`` / ``forward_batch`` / ``close``.
+    this class owns the preprocess pool, the A5 barrier, and the micro-batcher.
 
     Args:
         backend: The author-written ``VisionEncoderBackend``.
@@ -69,32 +70,38 @@ def __init__(
         self._backend = backend
         self._preprocess_concurrency = preprocess_concurrency
         self._name = name
-        self._actor: Optional[ThreadPoolExecutor] = None  # build + every forward
-        self._pool: Optional[ThreadPoolExecutor] = None  # off-loop preprocess
+        self._batcher: Optional[ThreadedMicroBatcher] = None
+        self._pool: Optional[ThreadPoolExecutor] = None
 
     # ---- lifecycle ---------------------------------------------------------
 
     def load(self, model_id: str, device: str) -> None:
-        """Run ``backend.build`` on the actor thread and fail fast.
+        """Start the actor thread (running ``backend.build`` on it) and fail fast.
 
         Re-raises any build error, then ``validate``s the placeholder id so a
         misconfigured encoder errors at startup instead of on the first request.
         Single-shot: a second ``load()`` raises rather than orphaning the first
-        actor thread and model.
+        batcher's (non-daemon) worker thread and model.
         """
-        if self._actor is not None or self._pool is not None:
+        if self._batcher is not None or self._pool is not None:
             raise RuntimeError("AsyncVisionEncoder.load() called twice")
+        # Construct the pool + batcher INSIDE the try so a constructor failure
+        # (e.g. a backend exposing a max_batch_cost the batcher rejects) still
+        # reaps the pool via shutdown() instead of leaking it. shutdown() is
+        # None-safe on the not-yet-assigned member.
         try:
-            # One actor thread so build + every forward share a thread; a single
-            # worker also serializes forwards (FIFO) — no cross-request batching.
-            self._actor = ThreadPoolExecutor(
-                max_workers=1, thread_name_prefix=f"{self._name}-actor"
-            )
             self._pool = ThreadPoolExecutor(
                 max_workers=self._preprocess_concurrency,
                 thread_name_prefix=f"{self._name}-pre",
             )
-            self._actor.submit(self._backend.build, model_id, device).result()
+            self._batcher = ThreadedMicroBatcher(
+                self._backend.forward_batch,
+                max_batch_cost=self._backend.max_batch_cost,
+                on_start=lambda: self._backend.build(model_id, device),
+                on_stop=self._backend.close,
+                name=self._name,
+            )
+            self._batcher.start()  # runs backend.build() on the actor thread
             self.validate()
         except BaseException:
             self.shutdown()
@@ -112,48 +119,41 @@ def get_image_placeholder_token_id(self) -> int:
     # ---- request path ------------------------------------------------------
 
     async def encode(self, raws: List[RawT]) -> List[torch.Tensor]:
-        """Preprocess (off-loop, A5 barrier) then run a single serial forward.
+        """Preprocess (off-loop, A5 barrier) then batched-encode; all-or-nothing.
 
         Returns one ``(n_visual_tokens, lm_hidden_dim)`` tensor per raw input, in
-        order. Raises if any image's preprocess fails (no GPU work) or if the
-        forward fails.
+        order. Raises if any image's preprocess fails (submitting nothing) or if
+        the batched forward fails.
         """
-        if self._actor is None or self._pool is None:
+        if self._batcher is None or self._pool is None:
             raise RuntimeError("AsyncVisionEncoder.encode() called before load()")
         if not raws:
             return []
         loop = asyncio.get_running_loop()
         # A5 barrier: preprocess all images concurrently, wait for EVERY one to
-        # settle, and run the forward only if all succeeded. return_exceptions=True
-        # makes the gather a true barrier (no short-circuit), so a failed sibling
-        # cannot leave a half-run request.
+        # settle, and submit only if all succeeded. return_exceptions=True makes
+        # the gather a true barrier (it never short-circuits), so a failed sibling
+        # cannot leave a half-submitted request — we submit nothing on any error.
         tasks = [
             loop.run_in_executor(self._pool, self._backend.preprocess, raw)
             for raw in raws
         ]
         settled = await asyncio.gather(*tasks, return_exceptions=True)
         errors = [r for r in settled if isinstance(r, BaseException)]
         if errors:
+            # Fail the whole request atomically; no item was submitted (no GPU
+            # work). Surface the first failure.
             raise errors[0]
         preprocessed: List[Preprocessed] = list(settled)  # type: ignore[arg-type]
         items = [p.item for p in preprocessed]
-        # Direct, serialized forward on the actor thread (eager; target_bucket
-        # defaults to None — there is no graph ladder in this milestone).
-        return await loop.run_in_executor(
-            self._actor, self._backend.forward_batch, items
-        )
+        costs = [p.cost for p in preprocessed]
+        bucket_keys = [p.bucket_key for p in preprocessed]
+        return await self._batcher.submit(items, costs, bucket_keys)
 
     def shutdown(self) -> None:
-        """Run ``backend.close`` on the actor thread, then stop both pools. Safe
-        before ``load`` and idempotent."""
-        if self._actor is not None:
-            try:
-                self._actor.submit(self._backend.close).result(timeout=10)
-            except BaseException:  # noqa: BLE001 — teardown best-effort
-                logger.exception(
-                    "AsyncVisionEncoder(%s): backend.close raised during teardown",
-                    self._name,
-                )
-            self._actor.shutdown(wait=False)
+        """Stop the actor thread (running ``backend.close`` on it) and the
+        preprocess pool. Safe before ``load`` and idempotent."""
+        if self._batcher is not None:
+            self._batcher.shutdown()  # runs backend.close() on the actor thread
         if self._pool is not None:
             self._pool.shutdown(wait=False)