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 to the author's max_batch_cost token
budget (scalar cost, one-dimensional packing — no shape buckets, no graph ladder;
that lands in PR4). 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(items) / on_stop (close); eager
  drain-on-completion (no timer; max_wait_ms is an opt-in knob); coalesce by
  scalar cost up to max_batch_cost (None = pass-through). Per-item idempotent
  finalizer; cancellation + repeat-cancel retirement; worker supervisor; optional
  max_outstanding_cost admission; cross-loop-safe shutdown. submit(items, costs).
- AsyncVisionEncoder now routes preprocessed items through the batcher (was a
  single-worker serial executor); the A5 preprocess barrier is unchanged; build is
  build(model_id) and the hardcoded image_token_id is validated at load.

Tests: batcher concurrency suite (eager-drain, cost budget, pass-through,
cancellation/retirement, supervisor, lifecycle, on_stop) + batcher-backed glue.

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,37 +1,41 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Serial async glue between the worker's event loop and a ``VisionEncoderBackend``.
-
-This is the **eager** 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).
-
-Request-level atomicity: 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 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 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**: 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 scalar ``cost``)
+  to a ``ThreadedMicroBatcher``, which coalesces across concurrent ``encode`` calls
+  by cost and runs ``backend.forward_batch`` 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
+teardown. ``load`` fails fast: it re-raises a build error and resolves the image
+placeholder id once, so a misconfigured encoder errors at startup, not on the
+first request.
 """
 
 from __future__ import annotations
 
 import asyncio
 import logging
 from concurrent.futures import ThreadPoolExecutor
-from typing import Generic, List
+from typing import Generic, List, Optional
 
 import torch
 
+from dynamo.vllm.multimodal_utils.threaded_micro_batcher import ThreadedMicroBatcher
 from dynamo.vllm.multimodal_utils.vision_encoder_backend import (
     ItemT,
     Preprocessed,
@@ -43,12 +47,16 @@
 
 
 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``.
+        preprocess_concurrency: Worker threads for off-loop ``preprocess``.
+        name: Base name for the actor thread / preprocess pool.
     """
 
     def __init__(
@@ -63,31 +71,38 @@ def __init__(
         self._backend = backend
         self._preprocess_concurrency = preprocess_concurrency
         self._name = name
-        self._actor: ThreadPoolExecutor | None = None  # build + every forward
-        self._pool: ThreadPoolExecutor | None = None  # off-loop preprocess
+        self._batcher: Optional[ThreadedMicroBatcher] = None
+        self._pool: Optional[ThreadPoolExecutor] = None
 
     # ---- lifecycle ---------------------------------------------------------
 
     def load(self, model_id: 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 hardcoded image token id
-        so a misconfigured encoder errors at startup. Single-shot: a second
-        ``load()`` raises rather than orphaning the first actor thread and model.
+        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
+        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).result()
+            self._batcher = ThreadedMicroBatcher(
+                self._backend.forward_batch,
+                max_batch_cost=self._backend.max_batch_cost,
+                on_start=lambda: self._backend.build(model_id),
+                on_stop=self._backend.close,
+                name=self._name,
+            )
+            self._batcher.start()  # runs backend.build() on the actor thread
             self.validate()
         except BaseException:
             self.shutdown()
@@ -110,48 +125,40 @@ 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)`` CPU tensor per raw input,
-        in order. Raises if any image's preprocess fails (no GPU work) or if the
-        forward fails.
+        Returns one ``(n_visual_tokens, lm_hidden_dim)`` tensor per raw input, in
+        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).
+        # 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 until CUDA-graph batching
-        # is supported).
-        return await loop.run_in_executor(
-            self._actor, self._backend.forward_batch, items
-        )
+        costs = [p.cost for p in preprocessed]
+        return await self._batcher.submit(items, costs)
 
     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)