Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
furionw merged 1 commit into from
Jul 15, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions components/src/dynamo/vllm/handlers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1089,8 +1089,8 @@ def _load_custom_encoder(self, config: Config) -> None:
f"VisionEncoderBackend subclass, got {backend_cls!r}."
)
# The author writes the VisionEncoderBackend; Dynamo wraps it in the
# AsyncVisionEncoder glue, which owns the preprocess pool and actor
# thread. load() runs backend.build() on the actor thread
# AsyncVisionEncoder glue, which owns the preprocess pool and
# ThreadedMicroBatcher actor thread. load() runs backend.build() there
# (the backend picks its own device) and cleans that thread up on failure.
encoder = AsyncVisionEncoder(backend_cls())
encoder.load(config.model)
Expand Down Expand Up @@ -2925,8 +2925,8 @@ async def _assemble_custom_encoder_prompt(
# failure becomes a structured request error instead of escaping the
# request coroutine and tearing down the stream.
try:
# encode() preprocesses off-thread and serializes forwards on one
# dedicated actor thread.
# AsyncVisionEncoder preprocesses off-thread; its ThreadedMicroBatcher
# coalesces concurrent calls onto one dedicated actor thread.
img_tensors: list[torch.Tensor] = await self._custom_encoder.encode(
image_urls
)
Expand Down
172 changes: 90 additions & 82 deletions components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,60 +1,67 @@
# SPDX-FileCopyrightText: Copyright (c) 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

import torch

from dynamo.vllm.multimodal_utils.threaded_micro_batcher import ThreadedMicroBatcher
from dynamo.vllm.multimodal_utils.vision_encoder_backend import (
ItemT,
Preprocessed,
RawT,
VisionEncoderBackend,
)

logger = logging.getLogger(__name__)


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 optional preprocess pool, the request-atomicity barrier,
and the single actor thread that runs ``build`` / ``forward_batch`` / ``close``.
and the micro-batcher.

``preprocess_concurrency`` (backend-declared, constructor-overridable) is not
just a pool size — it gates the whole preprocess **phase**: ``0`` means
``preprocess`` is never called and raws go straight to ``forward_batch``. A
backend that overrides ``preprocess`` while the effective concurrency is ``0``
is rejected at construction (the override would silently never run).

Args:
backend: The author-written ``VisionEncoderBackend``.
preprocess_concurrency: Off-loop ``preprocess`` pool size; ``None`` ⇒ use
the backend's value (default 0 ⇒ no pool / passthrough).
name: Base name for the actor thread / preprocess pool.
"""

def __init__(
Expand Down Expand Up @@ -91,25 +98,33 @@ def __init__(
self._backend = backend
self._preprocess_concurrency = conc
self._name = name
self._actor: ThreadPoolExecutor | None = None # build + every forward
self._pool: ThreadPoolExecutor | None = None # off-loop preprocess (or None)
self._batcher: ThreadedMicroBatcher | None = None
self._pool: ThreadPoolExecutor | None = 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:
if self._batcher is not None:
raise RuntimeError("AsyncVisionEncoder.load() called twice")
# Construct INSIDE the try so a later start()/build failure reaps the pool
# via shutdown() (None-safe on the not-yet-assigned members). The batcher
# is built before the pool so that a config it rejects (e.g. a backend
# max_batch_cost < 1) raises from its ctor before any pool is spawned —
# nothing to reap in that case.
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._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,
)
# No pool when concurrency is 0 — preprocess is skipped (passthrough).
self._pool = (
Expand All @@ -120,7 +135,7 @@ def load(self, model_id: str) -> None:
if self._preprocess_concurrency > 0
else None
)
self._actor.submit(self._backend.build, model_id).result()
self._batcher.start() # runs backend.build() on the actor thread
self.validate()
except BaseException:
self.shutdown()
Expand All @@ -144,62 +159,55 @@ def get_image_placeholder_token_id(self) -> int:

async def encode(self, raws: List[RawT]) -> List[torch.Tensor]:
"""Optionally preprocess (off-loop, with a request-atomicity barrier) then
run a single serial forward.
batched-encode.

With no preprocess pool (``preprocess_concurrency == 0``) raws go straight
to ``forward_batch`` (the backend folds any prep in there). 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.
to the batcher (the backend folds any prep into ``forward_batch``). 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:
if self._batcher is None:
raise RuntimeError("AsyncVisionEncoder.encode() called before load()")
if not raws:
return []
loop = asyncio.get_running_loop()
if self._pool is None:
# No preprocess phase: raw IS the item. No barrier needed — the
# single forward is already all-or-nothing for the request.
items: List[ItemT] = list(raws) # type: ignore[arg-type] # ItemT==RawT
else:
# Request-atomicity barrier: preprocess all images concurrently, wait
# for EVERY one to settle, run the forward only if all succeeded.
# return_exceptions=True makes the gather a true barrier (no short-circuit).
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:
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
)
# No preprocess phase: raw IS the item (cost defaults to 1). No
# barrier needed — the batched forward is all-or-nothing per request.
return await self._batcher.submit(list(raws)) # type: ignore[arg-type]
loop = asyncio.get_running_loop()
# Request-atomicity barrier: preprocess all images concurrently, wait for
# EVERY one to 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)
for result in settled:
if isinstance(result, BaseException):
# Fail the whole request atomically; no item was submitted (no GPU
# work). Surface the first failure, in order.
raise result
# No exception above ⇒ every settled entry is a Preprocessed. Alias the
# list gather() already returned rather than copying it.
preprocessed: List[Preprocessed] = 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)
Comment thread
furionw marked this conversation as resolved.

def shutdown(self) -> None:
"""Run ``backend.close`` on the actor thread, then stop both pools. Safe
before ``load`` and idempotent."""
# Detach both executors before teardown so a repeated cleanup is a no-op,
# including when backend.close() itself raises.
actor = self._actor
"""Stop the actor thread (running ``backend.close`` on it) and the
preprocess pool. Safe before ``load`` and idempotent."""
# Detach both resources before teardown so repeated cleanup is a no-op,
# including if teardown itself raises.
batcher = self._batcher
pool = self._pool
self._actor = None
self._batcher = None
self._pool = None

if actor is not None:
try:
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,
)
actor.shutdown(wait=False)
if batcher is not None:
batcher.shutdown() # runs backend.close() on the actor thread
if pool is not None:
pool.shutdown(wait=False)
Loading
Loading