From 039278305eb2d03075d2d9a82f72d3a5f10a3761 Mon Sep 17 00:00:00 2001 From: Qi Wang Date: Mon, 13 Jul 2026 19:45:32 -0700 Subject: [PATCH] feat(multimodal): batch custom vision encoder requests Add a supervised, thread-affine micro-batcher and route AsyncVisionEncoder forwards through it. Preserve preprocess atomicity while making lifecycle, failures, and teardown deterministic. Signed-off-by: Qi Wang --- components/src/dynamo/vllm/handlers.py | 8 +- .../multimodal_utils/async_vision_encoder.py | 172 ++--- .../threaded_micro_batcher.py | 480 ++++++++++++ .../test_vllm_async_vision_encoder.py | 67 +- .../test_vllm_threaded_micro_batcher.py | 683 ++++++++++++++++++ docs/features/multimodal/README.md | 1 + .../multimodal/custom-vision-encoder.md | 160 ++++ docs/features/multimodal/multimodal-vllm.md | 7 + docs/index.yml | 2 + 9 files changed, 1461 insertions(+), 119 deletions(-) create mode 100644 components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py create mode 100644 components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py create mode 100644 docs/features/multimodal/custom-vision-encoder.md diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index c96098ca90e1..041a226536d0 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -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) @@ -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 ) diff --git a/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py b/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py index a6e7a142a79a..db483a655b43 100644 --- a/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py +++ b/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py @@ -1,37 +1,40 @@ # 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, @@ -39,22 +42,26 @@ 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__( @@ -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 = ( @@ -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() @@ -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) 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) diff --git a/components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py b/components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py new file mode 100644 index 000000000000..92272efd5bc4 --- /dev/null +++ b/components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run a blocking ``fn(list[item]) -> list[result]`` on one dedicated worker +thread, coalescing items from concurrent async ``submit()`` calls into batches. + +``ThreadedMicroBatcher`` is the generic execution mechanism behind +``AsyncVisionEncoder`` — no model/vision knowledge, torch-free. It owns: + +- a **dedicated worker thread** that runs an optional ``on_start`` (e.g. build + + CUDA-graph capture), then every ``fn`` call, then an optional ``on_stop`` at + teardown — so anything thread-affine (CUDA graphs, the device/stream) is + captured and replayed on the same thread; +- **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). + +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 +caller. A request is finalised exactly once — its ``completion`` future resolves +only after *all* its items are delivered or failed — with every request-state +transition under one short lock, so the worker and a concurrent ``shutdown()`` +never race. A worker crash fails every live request (no hung awaiter) via a +supervisor. + +Cancellation tombstones the request so work that has not passed the final +pre-``fn`` dispatch check is dropped. A committed ``fn`` remains non-preemptible. +TODO: admission has no backpressure (``submit()`` always accepts). +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import logging +import queue +import threading +from dataclasses import dataclass +from enum import Enum, auto +from typing import Callable, Generic, List, Optional, TypeVar + +logger = logging.getLogger(__name__) + +T = TypeVar("T") +R = TypeVar("R") + +# Sentinel pushed onto the queue to stop the worker thread. +_SHUTDOWN = object() + + +class _State(Enum): + NEW = auto() + RUNNING = auto() + CLOSED = auto() # shutdown begun — no further submits (worker may still be exiting) + FAILED = auto() + + +@dataclass(eq=False) # identity-hashable: tracked in a set, compared by identity +class _Request(Generic[R]): + """One ``submit()`` call: its future resolves to one result per item, in order. + + ``completion`` is a thread-safe ``concurrent.futures.Future``. All mutation + (``remaining`` / ``done`` / ``error`` / ``results``) happens under the batcher + lock, so the worker and a concurrent ``shutdown()`` never race. + """ + + completion: "concurrent.futures.Future[List[R]]" + results: List[Optional[R]] + remaining: int + error: Optional[BaseException] = None + done: bool = False + + +@dataclass +class _Work(Generic[T]): + """A single item plus where its result belongs in the owning request.""" + + item: T + cost: int + request: _Request + index: int + + +class ThreadedMicroBatcher(Generic[T, R]): + """Run ``fn(list[item]) -> list[result]`` on a dedicated thread, coalescing + concurrent ``submit()`` calls into cost-bounded batches. + + Args: + fn: Batched work; one result per item, in order. Runs on the worker thread. + 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). + 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 + the serving loop ends), iff ``on_start`` succeeded. Its failure is + logged, never raised. + name: Worker thread name. + join_timeout_s: Seconds ``shutdown()`` waits for an in-flight ``fn``. + """ + + def __init__( + self, + fn: Callable[[List[T]], List[R]], + *, + max_batch_cost: 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)") + self._fn = fn + self._max_batch_cost = max_batch_cost + self._on_start = on_start + self._on_stop = on_stop + self._name = name + self._join_timeout_s = join_timeout_s + + self._queue: queue.Queue = queue.Queue() + # Completed by the worker once on_start settles: result(None) on success, + # set_exception(exc) on failure. start() blocks on it and re-raises — + # one primitive in place of a ready-Event + a start-error field. + self._started: "concurrent.futures.Future[None]" = concurrent.futures.Future() + self._terminal_error: Optional[BaseException] = None + # Guards _state, _live, every _Request transition, and the queue-commit so + # a state change and its work enqueue happen atomically. Bookkeeping only + # — never held across fn / join. + self._lock = threading.Lock() + self._state = _State.NEW + self._live: set[_Request] = set() + self._thread: Optional[threading.Thread] = None + + # ---- lifecycle --------------------------------------------------------- + + def start(self) -> None: + """Start the worker thread, run ``on_start`` on it, and re-raise its error. + + Single-shot: once ``start()`` has been called — even if the thread failed + to spawn — a second call raises rather than spawning a second consumer / + orphaning the first thread.""" + with self._lock: + # Reject a second start() on two axes, both read under the lock: + # `_thread` is already set — a first start() spawned the worker, even + # if it is still blocked waiting for on_start while `_state` is NEW + # (covers concurrent double-start); OR `_state` moved past NEW — a + # failed spawn set FAILED with `_thread` still None (covers a + # failed-spawn retry). Either alone leaves a hole; together they don't. + if self._thread is not None or self._state is not _State.NEW: + raise RuntimeError("ThreadedMicroBatcher.start() called twice") + # Non-daemon: a clean stop is via shutdown(); a daemon worker could be + # torn down mid-fn at interpreter exit. Pin daemon=False explicitly so + # it never inherits a daemon creator thread. Start under the lock and + # publish _thread only after a successful start(), so a racing + # shutdown() can never join() an unstarted thread. + thread = threading.Thread(target=self._run, name=self._name, daemon=False) + try: + thread.start() + except BaseException: + self._state = _State.FAILED + raise + self._thread = thread + try: + self._started.result() # blocks until on_start settles; re-raises it + except BaseException: + # on_start ran on the thread, which then exited — mark closed so a + # later submit() raises instead of queueing to a dead consumer. + self.shutdown() + raise + with self._lock: + # Only this (winning) start() ever sets RUNNING, so the state is NEW + # unless a concurrent shutdown()/crash moved it to CLOSED/FAILED. + if self._state is not _State.NEW: + raise RuntimeError("ThreadedMicroBatcher shut down during start()") + self._state = _State.RUNNING + + async def submit( + self, + items: List[T], + costs: Optional[List[int]] = 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. + + Cancelling the await tombstones this request. Work that has not passed the + final pre-``fn`` dispatch check is dropped; a committed ``fn`` still finishes. + """ + if self._thread is None: + raise RuntimeError("ThreadedMicroBatcher.submit() called before start()") + if not items: + return [] + if costs is None: + costs = [1] * len(items) + elif len(costs) != len(items): + raise ValueError(f"costs has {len(costs)} 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}") + if self._max_batch_cost is not None and c > self._max_batch_cost: + raise ValueError( + f"item cost {c} exceeds max_batch_cost {self._max_batch_cost}; " + "it has no batch it can fit" + ) + 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)) + ] + # State check + admission + queue-commit under one lock so a concurrent + # shutdown() cannot strand the request and capacity cannot leak. + with self._lock: + if self._state is _State.FAILED: + raise RuntimeError( + "ThreadedMicroBatcher.submit() after worker failure" + ) from self._terminal_error + if self._state is not _State.RUNNING: # NEW / CLOSED + raise RuntimeError( + "ThreadedMicroBatcher.submit() called after shutdown()" + ) + self._live.add(request) + for work in works: + self._queue.put(work) + try: + return await asyncio.wrap_future(request.completion) + except asyncio.CancelledError: + # asyncio cancels the bridged completion future, but request state is + # separate. Tombstone it so queued work cannot enter a shared batch and + # make an unrelated live request fail. Keep `done` false: every work is + # still consumed exactly once, which releases `_live` on the final item. + with self._lock: + if not request.done and request.error is None: + request.error = asyncio.CancelledError() + raise + + def shutdown(self) -> None: + """Stop the worker, failing not-yet-run items. Idempotent and best-effort: + a slow in-flight ``fn`` keeps running and the worker exits on its own once + it returns and reads the stop signal.""" + to_fail: List[_Work] = [] + with self._lock: + # Snapshot _thread under the lock: start() publishes it under the same + # lock, so reading it here can't race with a concurrent spawn. + thread = self._thread + if thread is None: + return + if self._state not in (_State.CLOSED, _State.FAILED): + self._state = _State.CLOSED + to_fail = self._drain_queue_locked() # fail queued, then signal stop + self._queue.put(_SHUTDOWN) + # Admission is now closed and the queue drained under the lock, so no new + # work can arrive; failing the drained set covers everything not already + # in-flight on the worker. No post-join re-drain is needed. + for work in to_fail: + self._consume(work, error=RuntimeError("ThreadedMicroBatcher shut down")) + thread.join(timeout=self._join_timeout_s) + if thread.is_alive(): + logger.warning( + "ThreadedMicroBatcher(%s): worker still finishing an in-flight fn " + "after %gs; it will exit on its own.", + self._name, + self._join_timeout_s, + ) + + # ---- worker thread ----------------------------------------------------- + + def _run(self) -> None: + try: + if self._on_start is not None: + self._on_start() # build / warmup / CUDA-graph capture HERE + except BaseException as exc: # noqa: BLE001 — surface to start() + self._started.set_exception(exc) + return + self._started.set_result(None) + try: + while True: + works = self._collect() + if works is None: + return + self._dispatch(works) + except BaseException as exc: # noqa: BLE001 — supervisor: never hang awaiters + logger.exception( + "ThreadedMicroBatcher(%s): worker crashed; failing live requests", + self._name, + ) + with self._lock: + self._state = _State.FAILED + self._terminal_error = exc + live = list(self._live) + self._drain_queue_locked() # clear queue; those reqs are in `live` + for req in live: + self._abort(req, exc) + finally: + # on_stop runs on the actor thread (so CUDA teardown is same-thread), + # only after on_start succeeded (this finally is unreachable otherwise). + self._run_on_stop() + + def _run_on_stop(self) -> None: + if self._on_stop is None: + return + try: + self._on_stop() + except BaseException: # noqa: BLE001 — teardown best-effort, never raise + logger.exception( + "ThreadedMicroBatcher(%s): on_stop raised during teardown", + self._name, + ) + + def _collect(self) -> Optional[List[_Work]]: + """Block for one item, then eager-drain everything else already queued. + + No timed hold: pull only what is immediately available, then run.""" + first = self._queue.get() + if first is _SHUTDOWN: + return None + works: List[_Work] = [first] + # Eager drain: pull everything immediately available, then run. + while True: + try: + item = self._queue.get_nowait() + except queue.Empty: + break + if item is _SHUTDOWN: + self._queue.put(_SHUTDOWN) # drain this round, stop next loop + break + works.append(item) + return works + + def _dispatch(self, works: List[_Work]) -> None: + """Split live items by cost budget, run ``fn`` (one-dimensional packing). + + Tombstoned (failed / done) items are dropped before batching — + an already-failed request never reaches ``fn``.""" + live: List[_Work] = [] + for work in works: + if self._is_tombstoned(work.request): + self._consume(work) # account the dropped item + else: + 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 + for work in live: + if batch and batch_cost + work.cost > self._max_batch_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) + + def _run_batch(self, batch: List[_Work]) -> None: + # Re-filter immediately before fn: a request may have been failed by a + # sibling batch after grouping, so a failed request's remaining items are + # dropped at each fn call rather than only at the per-dispatch snapshot. + runnable: List[_Work] = [] + for work in batch: + if self._is_tombstoned(work.request): + self._consume(work) + else: + runnable.append(work) + if not runnable: + return + # Once shutdown/failure has begun, do not START new fn calls: fail these + # collected-but-not-yet-run items with the shutdown error. The fn already + # in flight when shutdown() was called still finishes (it is past this + # check); this just bounds work after teardown intent to that one batch. + with self._lock: + stopping = self._state is not _State.RUNNING + if stopping: + for work in runnable: + self._consume( + work, error=RuntimeError("ThreadedMicroBatcher shut down") + ) + return + items = [w.item for w in runnable] + try: + results = self._fn(items) + except ( + BaseException + ) as exc: # noqa: BLE001 — a bad batch must not hang awaiters + for work in runnable: + self._consume(work, error=exc) + return + if len(results) != len(items): + err = RuntimeError( + f"batch fn returned {len(results)} results for {len(items)} items; " + "it must return one result per item" + ) + for work in runnable: + self._consume(work, error=err) + return + for work, result in zip(runnable, results): + self._consume(work, result=result) + + def _is_tombstoned(self, req: _Request) -> bool: + """True once the request must not send further items to ``fn``.""" + with self._lock: + return req.done or req.error is not None + + def _consume(self, work: _Work, result: object = None, error=None) -> None: + """Account one item of a request (delivered / failed / tombstoned). + + Decrements ``remaining`` under the lock and finalises the request only + when the last item is consumed — so admission release and ``completion`` + are exactly-once even when items span batches or threads. A bare + ``_consume(work)`` only ever accounts an already-tombstoned item — such a + request has ``error`` set, so the ``result`` write below is skipped and + the defaulted ``None`` is never stored.""" + req = work.request + finalize = False + with self._lock: + if req.done: + return + if error is not None and req.error is None: + req.error = error + elif req.error is None: + req.results[work.index] = result + req.remaining -= 1 + if req.remaining == 0: + req.done = True + self._live.discard(req) + finalize = True + if finalize: + self._complete(req) + + def _abort(self, req: _Request, exc: BaseException) -> None: + """Force-finalise a live request (worker crash): items may be lost, so do + not wait for ``remaining``.""" + with self._lock: + if req.done: + return + req.done = True + if req.error is None: + req.error = exc + self._live.discard(req) + self._complete(req) + + def _complete(self, req: _Request) -> None: + """Settle a finalised request's futures (outside the lock; idempotent).""" + try: + if req.error is not None: + req.completion.set_exception(req.error) + else: + req.completion.set_result(list(req.results)) + except concurrent.futures.InvalidStateError: + pass # caller already cancelled / abandoned the future + + def _drain_queue_locked(self) -> List[_Work]: + """Pop all queued works (caller holds the lock); returns them to consume + outside the lock (``_consume`` re-takes the lock). + + Both callers own the stop signal: ``shutdown()`` drains, then enqueues a + fresh ``_SHUTDOWN``; the crash supervisor drains as the sole worker exits. + So any ``_SHUTDOWN`` seen here is stale — drop it and keep draining rather + than re-queueing it.""" + drained: List[_Work] = [] + while True: + try: + item = self._queue.get_nowait() + except queue.Empty: + return drained + if item is not _SHUTDOWN: + drained.append(item) diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py index 41f6022f0bf6..794c96384f09 100644 --- a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py @@ -1,16 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for the serial dynamo.vllm.multimodal_utils.async_vision_encoder. +"""Unit tests for dynamo.vllm.multimodal_utils.async_vision_encoder. -Pin the serial-glue contract: build / forward / close run on one actor thread; -encode returns one tensor per raw; the preprocess barrier fails a request -atomically (no GPU work) if any image's preprocess fails; concurrent encodes are -**not** coalesced (one forward_batch call each — the batched version's job, added -later); load fails fast on a build error or a missing/invalid image_token_id. +Pin the glue contract: build / forward / close run on one actor thread; encode +returns one tensor per raw; the preprocess barrier fails a request atomically +(no GPU work) if any image's preprocess fails; load fails fast on a build error or +a missing/invalid hardcoded image_token_id and reaps its thread. """ -import asyncio import threading import pytest @@ -32,8 +30,9 @@ class _FakeBackend(VisionEncoderBackend): - """A CPU-only fake backend; records its threads and forward-call boundaries.""" + """A CPU-only fake backend; records its threads.""" + max_batch_cost = 8 image_token_id = 151655 # This fake overrides preprocess, so opt into the off-loop pool (the default # is 0 ⇒ no pool); the passthrough path is covered by its own test below. @@ -47,7 +46,6 @@ def __init__(self, *, fail_on=None): self.close_calls = 0 self.model_id = None self.forward_threads: list[int] = [] - self.forward_calls: list[list] = [] # one entry per forward_batch call def build(self, model_id): self.build_thread = threading.get_ident() @@ -60,7 +58,6 @@ def preprocess(self, raw): def forward_batch(self, items, target_bucket=None): self.forward_threads.append(threading.get_ident()) - self.forward_calls.append(list(items)) return [torch.full((2, 4), float(len(str(it)))) for it in items] def close(self): @@ -87,7 +84,7 @@ async def test_preprocess_barrier_fails_atomically_with_no_gpu_work(): try: with pytest.raises(ValueError, match="bad input"): await enc.encode(["good", "bad"]) - assert be.forward_calls == [] # nothing ran + assert be.forward_threads == [] # nothing was submitted finally: enc.shutdown() @@ -105,20 +102,6 @@ async def test_build_and_forward_share_one_non_main_thread(): enc.shutdown() -async def test_concurrent_encodes_are_not_coalesced(): - """Serial glue: each encode runs its own forward_batch — no cross-request - batching (that is the batched version's job, added in a later PR).""" - be = _FakeBackend() - enc = AsyncVisionEncoder(be) - enc.load("m") - try: - await asyncio.gather(enc.encode(["a"]), enc.encode(["b"]), enc.encode(["c"])) - assert len(be.forward_calls) == 3 - assert all(len(call) == 1 for call in be.forward_calls) - finally: - enc.shutdown() - - async def test_load_resolves_placeholder_and_passes_model_id(): be = _FakeBackend() enc = AsyncVisionEncoder(be) @@ -161,7 +144,7 @@ def test_shutdown_runs_backend_close_on_actor_thread(): enc.load("m") enc.shutdown() assert be.closed is True - assert be.close_thread == be.build_thread + assert be.close_thread == be.build_thread # close on the actor thread def test_shutdown_is_idempotent(): @@ -171,7 +154,7 @@ def test_shutdown_is_idempotent(): enc.shutdown() enc.shutdown() assert be.close_calls == 1 - assert enc._actor is None + assert enc._batcher is None assert enc._pool is None @@ -183,18 +166,18 @@ def build(self, model_id): enc = AsyncVisionEncoder(_BadBuild()) with pytest.raises(RuntimeError, match="build failed"): enc.load("m") - assert enc._actor is None + assert enc._batcher is None assert enc._pool is None def test_load_fails_fast_on_missing_image_token_id(): class _NoTokenId(_FakeBackend): - image_token_id = None + image_token_id = None # author forgot to hardcode it enc = AsyncVisionEncoder(_NoTokenId()) with pytest.raises(ValueError, match="image_token_id"): enc.load("m") - assert enc._actor is None + assert enc._batcher is None assert enc._pool is None @@ -220,8 +203,8 @@ def forward_batch(self, items, target_bucket=None): async def test_passthrough_skips_preprocess_when_no_pool(): - """With no pool (preprocess_concurrency=0) the preprocess phase is skipped - entirely and raws go straight to forward_batch — no barrier, no pool thread.""" + """With no pool (preprocess_concurrency=0) preprocess is skipped and raws go + straight to the batcher → forward_batch, no barrier, no pool thread.""" be = _PassthroughBackend() # Instance-level boom (invisible to the class-override check) proves the # phase is never entered, not merely run through an identity hook. @@ -234,7 +217,11 @@ async def test_passthrough_skips_preprocess_when_no_pool(): assert enc._pool is None # no pool created out = await enc.encode(["a", "bb"]) assert len(out) == 2 - assert be.forward_calls == [["a", "bb"]] # raws passed straight through + assert all(t.shape == (2, 4) for t in out) + # Passthrough must hand the raws themselves to forward_batch, unchanged + # and in order. A single request may span multiple micro-batches, so + # flatten the recorded calls rather than requiring one forward call. + assert [item for batch in be.forward_calls for item in batch] == ["a", "bb"] finally: enc.shutdown() @@ -269,3 +256,17 @@ def test_driver_override_to_zero_with_overriding_backend_raises(): def test_preprocess_concurrency_rejects_negative(): with pytest.raises(ValueError, match="preprocess_concurrency"): AsyncVisionEncoder(_FakeBackend(), preprocess_concurrency=-1) + + +def test_load_bad_batch_cost_fails_without_spawning_pool(): + """A backend whose max_batch_cost the batcher rejects fails load() cleanly: the + batcher ctor raises before the preprocess pool is ever spawned (nothing to reap).""" + + class _BadBudget(_FakeBackend): + max_batch_cost = 0 # ThreadedMicroBatcher requires >= 1 → ctor raises + + enc = AsyncVisionEncoder(_BadBudget()) + with pytest.raises(ValueError, match="max_batch_cost"): + enc.load("m") + assert enc._pool is None + assert enc._batcher is None diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py new file mode 100644 index 000000000000..daa3d0623aa0 --- /dev/null +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py @@ -0,0 +1,683 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for dynamo.vllm.multimodal_utils.threaded_micro_batcher. + +Pin the execution contract: on_start + every fn call (+ on_stop) run on one +dedicated thread (so CUDA-graph capture/replay share a thread), concurrent submits +coalesce into cost-bounded batches up to max_batch_cost (or pass-through when +None), eager-drain pulls all queued work when free, errors reach every awaiting +caller, and the shutdown lifecycle behaves. + +``fn`` is ``fn(items)``; ``cost`` is a precomputed scalar that rides on +``submit(items, costs)`` (one-dimensional packing — no bucket_key, no ladder). +""" + +import asyncio +import threading +from typing import Callable + +import pytest + +from dynamo.vllm.multimodal_utils.threaded_micro_batcher import ThreadedMicroBatcher + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.multimodal, + pytest.mark.timeout(30), +] + + +def _echo(items): + return list(items) + + +class _Recorder: + """fn that records the threads it ran on and the batches it received.""" + + def __init__(self): + self.threads: list[int] = [] + self.start_thread: int | None = None + self.stop_thread: int | None = None + + def on_start(self): + self.start_thread = threading.get_ident() + + def on_stop(self): + self.stop_thread = threading.get_ident() + + def fn(self, items): + self.threads.append(threading.get_ident()) + return [("r", x) for x in items] + + +async def _wait_until( + predicate: Callable[[], bool], + message: str, + timeout_s: float = 5.0, +) -> None: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + while not predicate(): + if loop.time() >= deadline: + raise AssertionError(message) + await asyncio.sleep(0.01) + + +class _Gate: + """Park the worker inside a sentinel "gate" batch until released, so a test can + enqueue several submits that then drain together as ONE batch — deterministic + coalescing without relying on a timed hold.""" + + GATE = "__gate__" + + def __init__(self): + self.entered = threading.Event() + self.release = threading.Event() + self.batches: list[list] = [] # excludes the gate batch + + def fn(self, items): + if list(items) == [self.GATE]: + self.entered.set() + self.release.wait(timeout=5.0) + return [("r", x) for x in items] + self.batches.append(list(items)) + return [("r", x) for x in items] + + async def park(self, b): + """Submit the gate item; return its future once the worker is parked in fn.""" + fut = asyncio.ensure_future(b.submit([self.GATE])) + await _wait_until(self.entered.is_set, "worker never entered the gate batch") + return fut + + +async def test_submit_returns_one_result_per_item(): + rec = _Recorder() + b = ThreadedMicroBatcher(rec.fn, on_start=rec.on_start) + b.start() + try: + out = await b.submit(["a", "b", "c"]) + assert out == [("r", "a"), ("r", "b"), ("r", "c")] + finally: + b.shutdown() + + +async def test_on_start_fn_and_on_stop_share_one_non_main_thread(): + rec = _Recorder() + b = ThreadedMicroBatcher(rec.fn, on_start=rec.on_start, on_stop=rec.on_stop) + b.start() + await asyncio.gather(b.submit(["x"]), b.submit(["y"])) + b.shutdown() + assert rec.start_thread is not None + assert rec.start_thread != threading.get_ident() + assert set(rec.threads) == {rec.start_thread} + # on_stop runs on the same actor thread (so CUDA teardown is same-thread). + assert rec.stop_thread == rec.start_thread + + +def test_on_stop_not_run_if_on_start_failed(): + ran = {"stop": False} + + def bad_start(): + raise RuntimeError("start failed") + + def on_stop(): + ran["stop"] = True + + b = ThreadedMicroBatcher(_echo, on_start=bad_start, on_stop=on_stop) + with pytest.raises(RuntimeError, match="start failed"): + b.start() + assert ran["stop"] is False + + +async def test_eager_drain_pulls_all_queued_when_free(): + """Eager-drain: items queued while the worker is busy are all pulled into ONE + batch on the next free iteration (no timer).""" + entered = threading.Event() + release = threading.Event() + batches: list[list] = [] + + def fn(items): + batches.append(list(items)) + if "block" in items: + entered.set() + release.wait(timeout=5.0) + return [("r", x) for x in items] + + b = ThreadedMicroBatcher(fn) # default: eager-drain + b.start() + first = asyncio.ensure_future(b.submit(["block"])) + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + # Queue three while the worker is blocked in fn("block"). + rest = [asyncio.ensure_future(b.submit([x])) for x in ("a", "b", "c")] + await asyncio.sleep(0.05) + release.set() + await asyncio.gather(first, *rest) + b.shutdown() + # The post-release _collect drains a, b, c in one batch. + assert ["a", "b", "c"] in batches + + +async def test_cost_budget_caps_each_batch(): + """costs ride on submit; with budget 5, batches never exceed summed cost 5. + + Park the worker so [3, 3, 1] drain in one collect, deterministically + exercising the cost-split (3 | 3,1).""" + g = _Gate() + b = ThreadedMicroBatcher(g.fn, max_batch_cost=5) + b.start() + try: + gate = await g.park(b) + real = asyncio.ensure_future(b.submit([3, 3, 1], costs=[3, 3, 1])) + await asyncio.sleep(0.05) # let all three enqueue while the worker is parked + g.release.set() + await asyncio.gather(gate, real) + assert all(sum(batch) <= 5 for batch in g.batches) + assert sum(len(batch) for batch in g.batches) == 3 + assert [3] in g.batches and [3, 1] in g.batches # split actually happened + finally: + b.shutdown() + + +async def test_max_batch_cost_none_is_passthrough(): + """Default (max_batch_cost=None): no cap and no per-item ceiling — the whole + drained set runs as ONE fn call regardless of summed cost.""" + g = _Gate() + b = ThreadedMicroBatcher(g.fn) # max_batch_cost=None + b.start() + try: + gate = await g.park(b) + # Big per-item costs that would be split (or rejected) under any finite cap. + real = asyncio.ensure_future( + b.submit([1, 2, 3, 4], costs=[1000, 1000, 1000, 1000]) + ) + await asyncio.sleep(0.05) # let all four enqueue while the worker is parked + g.release.set() + out, _ = await asyncio.gather(real, gate) + assert len(out) == 4 + assert g.batches == [[1, 2, 3, 4]] # one un-split batch + finally: + b.shutdown() + + +async def test_error_reaches_every_caller(): + def boom(items): + raise ValueError("boom") + + b = ThreadedMicroBatcher(boom) + b.start() + try: + results = await asyncio.gather( + *(b.submit(["u"]) for _ in range(3)), return_exceptions=True + ) + assert all(isinstance(r, ValueError) and str(r) == "boom" for r in results) + finally: + b.shutdown() + + +async def test_cancelled_queued_request_cannot_poison_live_request(): + """A cancelled request is dropped before dispatch, so its bad item cannot + fail an unrelated request that was queued alongside it.""" + gate = _Gate() + + def fn(items): + results = gate.fn(items) + if any(item.startswith("bad") for item in items): + raise ValueError("bad poisoned batch") + return results + + b = ThreadedMicroBatcher(fn) + b.start() + try: + parked = await gate.park(b) + bad = asyncio.ensure_future(b.submit(["bad-1", "bad-2"])) + good = asyncio.ensure_future(b.submit(["good"])) + await _wait_until( + lambda: b._queue.qsize() == 3, + "requests were not queued behind the gate", + ) + + bad.cancel() + with pytest.raises(asyncio.CancelledError): + await bad + + gate.release.set() + await parked + assert await good == [("r", "good")] + assert gate.batches == [["good"]] + assert not b._live + finally: + gate.release.set() + b.shutdown() + + +async def test_cancellation_after_partial_dispatch_drops_later_siblings(): + """Cancellation after one item completes lets the claimed item finish, drops + later siblings, releases the request, and leaves the actor usable.""" + blocked = threading.Event() + release = threading.Event() + seen: list[str] = [] + + def fn(items): + seen.extend(items) + if items == ["blocked"]: + blocked.set() + release.wait(timeout=5.0) + return [("r", item) for item in items] + + b = ThreadedMicroBatcher(fn, max_batch_cost=1) + b.start() + try: + task = asyncio.ensure_future(b.submit(["first", "blocked", "later"])) + await _wait_until(blocked.is_set, "second item never entered fn") + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + release.set() + assert await b.submit(["probe"]) == [("r", "probe")] + assert seen == ["first", "blocked", "probe"] + assert not b._live + finally: + release.set() + b.shutdown() + + +async def test_cancellation_after_batch_claim_is_best_effort(): + """Once a shared batch is claimed, cancellation cannot remove its bad item; + the live peer may fail, but the actor remains usable afterward.""" + park_entered = threading.Event() + park_release = threading.Event() + batch_claimed = threading.Event() + batch_release = threading.Event() + seen: list[list[str]] = [] + + def fn(items): + if items == ["park"]: + park_entered.set() + park_release.wait(timeout=5.0) + return [("r", "park")] + seen.append(list(items)) + if "bad" in items: + batch_claimed.set() + batch_release.wait(timeout=5.0) + raise ValueError("claimed bad poisoned batch") + return [("r", item) for item in items] + + b = ThreadedMicroBatcher(fn) + b.start() + try: + parked = asyncio.ensure_future(b.submit(["park"])) + await _wait_until(park_entered.is_set, "worker never entered park batch") + bad = asyncio.ensure_future(b.submit(["bad"])) + good = asyncio.ensure_future(b.submit(["good"])) + await _wait_until( + lambda: b._queue.qsize() == 2, + "shared batch was not queued behind park", + ) + + park_release.set() + await parked + await _wait_until(batch_claimed.is_set, "shared batch was not claimed") + bad.cancel() + with pytest.raises(asyncio.CancelledError): + await bad + + batch_release.set() + with pytest.raises(ValueError, match="claimed bad poisoned batch"): + await good + assert seen == [["bad", "good"]] + assert await b.submit(["probe"]) == [("r", "probe")] + assert b._thread.is_alive() + assert not b._live + finally: + park_release.set() + batch_release.set() + b.shutdown() + + +async def test_cancellation_racing_shutdown_cleans_request(): + """Shutdown consumes every item of a canceled queued request without a hang, + double finalization, or live-request leak.""" + entered = threading.Event() + release = threading.Event() + + def fn(items): + if items == ["block"]: + entered.set() + release.wait(timeout=5.0) + return [("r", item) for item in items] + + b = ThreadedMicroBatcher(fn, join_timeout_s=0.1) + b.start() + try: + in_flight = asyncio.ensure_future(b.submit(["block"])) + await _wait_until(entered.is_set, "blocker never entered fn") + queued = asyncio.ensure_future(b.submit(["cancel-1", "cancel-2"])) + await _wait_until( + lambda: b._queue.qsize() == 2, + "canceled request was not queued behind blocker", + ) + + queued.cancel() + with pytest.raises(asyncio.CancelledError): + await queued + b.shutdown() + + release.set() + assert await in_flight == [("r", "block")] + b.shutdown() + assert not b._thread.is_alive() + assert not b._live + finally: + release.set() + b.shutdown() + + +async def test_cancellation_after_finalization_keeps_actor_usable(): + """Cancellation after done=True but before future settlement leaves the actor + alive when completion observes the already-canceled bridged future.""" + completion_entered = threading.Event() + completion_release = threading.Event() + first_completion = True + + b = ThreadedMicroBatcher(_echo) + original_complete = b._complete + + def gated_complete(request): + nonlocal first_completion + if first_completion: + first_completion = False + completion_entered.set() + completion_release.wait(timeout=5.0) + original_complete(request) + + b._complete = gated_complete + b.start() + try: + task = asyncio.ensure_future(b.submit(["done"])) + await _wait_until( + completion_entered.is_set, + "request did not reach finalization gate", + ) + assert not b._live + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + completion_release.set() + assert await b.submit(["probe"]) == ["probe"] + assert b._thread.is_alive() + assert not b._live + finally: + completion_release.set() + b.shutdown() + + +async def test_wrong_result_count_raises(): + b = ThreadedMicroBatcher(lambda items: []) + b.start() + try: + with pytest.raises(RuntimeError, match="one result per item"): + await b.submit(["a", "b"]) + finally: + b.shutdown() + + +async def test_costs_length_mismatch_raises(): + b = ThreadedMicroBatcher(_echo) + b.start() + try: + with pytest.raises(ValueError, match="costs has"): + await b.submit(["a", "b"], costs=[1]) + finally: + b.shutdown() + + +async def test_submit_before_start_raises(): + b = ThreadedMicroBatcher(_echo) + with pytest.raises(RuntimeError, match="before start"): + await b.submit(["a"]) + + +async def test_submit_after_shutdown_raises(): + b = ThreadedMicroBatcher(_echo) + b.start() + b.shutdown() + with pytest.raises(RuntimeError, match="after shutdown"): + await b.submit(["a"]) + + +def test_start_error_propagates_and_reaps(): + def bad_start(): + raise RuntimeError("start failed") + + b = ThreadedMicroBatcher(_echo, on_start=bad_start) + with pytest.raises(RuntimeError, match="start failed"): + b.start() + assert not b._thread.is_alive() + + +async def test_shutdown_fails_queued_items(): + entered = threading.Event() + release = threading.Event() + + def blocking(items): + entered.set() + release.wait(timeout=5.0) + return [("r", x) for x in items] + + b = ThreadedMicroBatcher(blocking, join_timeout_s=0.2) + b.start() + in_flight = asyncio.ensure_future(b.submit(["a"])) + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + + queued = [ + asyncio.ensure_future(b.submit(["b"])), + asyncio.ensure_future(b.submit(["c"])), + ] + await asyncio.sleep(0.05) + b.shutdown() # fails b, c; a is in flight + for q in queued: + with pytest.raises(RuntimeError, match="shut down"): + await q + release.set() + assert len(await in_flight) == 1 + b.shutdown() + assert not b._thread.is_alive() + + +def test_shutdown_stops_thread(): + b = ThreadedMicroBatcher(_echo) + b.start() + assert b._thread.is_alive() + b.shutdown() + assert not b._thread.is_alive() + + +async def test_worker_supervisor_fails_awaiters_on_crash(): + """An unexpected worker crash fails live awaiters and moves to a failed state + (later submits raise) instead of hanging.""" + b = ThreadedMicroBatcher(_echo) + b.start() + + def explode(_works): + raise RuntimeError("worker boom") + + b._dispatch = explode # force a crash inside the serve loop + try: + with pytest.raises(RuntimeError, match="worker boom"): + await b.submit(["a"]) + with pytest.raises(RuntimeError): # FAILED state rejects new work + await b.submit(["b"]) + finally: + b.shutdown() + + +async def test_oversized_cost_is_rejected(): + """A per-item cost above the batch budget has no batch it can fit → rejected.""" + b = ThreadedMicroBatcher(_echo, max_batch_cost=5) + b.start() + try: + with pytest.raises(ValueError, match="exceeds max_batch_cost"): + await b.submit([6], costs=[6]) + finally: + b.shutdown() + + +async def test_nonpositive_cost_is_rejected(): + b = ThreadedMicroBatcher(_echo) + b.start() + try: + with pytest.raises(ValueError, match="positive int"): + await b.submit([1], costs=[0]) + finally: + b.shutdown() + + +async def test_partial_batch_failure_fails_request_once(): + """A multi-item request split across batches where the FIRST batch raises + fails the whole request exactly once, and the later sibling item is + tombstoned — it never reaches fn.""" + seen: list = [] + + def fn(items): + seen.extend(items) + if "bad" in items: + raise ValueError("boom") + return [("r", x) for x in items] + + # max_batch_cost=1 → "bad" and "good" are separate (cost-1) batches; "bad" + # runs (and fails) first, tombstoning the request before "good" runs. + b = ThreadedMicroBatcher(fn, max_batch_cost=1) + b.start() + try: + with pytest.raises(ValueError, match="boom"): + await b.submit(["bad", "good"], costs=[1, 1]) + assert "good" not in seen # tombstoned sibling never reached fn + finally: + b.shutdown() + + +async def test_no_fn_after_shutdown_for_collected_items(): + """Items pulled off the queue by _collect but not yet run must not reach fn + once shutdown begins; they fail with the shutdown error.""" + entered = threading.Event() + release = threading.Event() + seen: list = [] + + def fn(items): + seen.extend(items) + if "a" in items: + entered.set() + release.wait(timeout=5.0) + return [("r", x) for x in items] + + # max_batch_cost=1 → one batch per item; all three collected together, the + # "a" batch blocks in fn while shutdown() is called. + b = ThreadedMicroBatcher(fn, max_batch_cost=1) + b.start() + task = asyncio.ensure_future(b.submit(["a", "b", "c"], costs=[1, 1, 1])) + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + b.shutdown() # b, c are collected but not yet run + release.set() + with pytest.raises(RuntimeError, match="shut down"): + await task + assert seen == ["a"] # b and c never reached fn + b.shutdown() + + +def test_double_start_raises(): + b = ThreadedMicroBatcher(_echo) + b.start() + try: + with pytest.raises(RuntimeError, match="twice"): + b.start() + finally: + b.shutdown() + + +def test_concurrent_start_starts_one_worker(): + """Two threads race start(); the winner parks inside on_start while the loser + must already be rejected — one worker only. + + Deterministic regression guard for the double-start race: the winner blocks in + on_start, so `_state` is still NEW when the loser races in. A state-only guard + would let the loser through *here* (state == NEW), spawn a second worker, and + run on_start twice; the fix also keys on `_thread` (published under the lock), + so the loser is rejected while the winner is still parked. Blocking on_start + forces the window every run (no reliance on scheduler timing). + """ + entered = threading.Event() + release = threading.Event() + ran = {"n": 0} + ran_lock = threading.Lock() + + def on_start(): + with ran_lock: + ran["n"] += 1 + entered.set() + release.wait(timeout=5.0) + + b = ThreadedMicroBatcher(_echo, on_start=on_start) + outcomes: list[str] = [] + out_lock = threading.Lock() + outcome_recorded = threading.Event() + gate = threading.Barrier(2) + + def do_start(): + gate.wait() # both threads reach start() together + try: + b.start() + tag = "ok" + except RuntimeError: + tag = "rejected" + with out_lock: + outcomes.append(tag) + outcome_recorded.set() + + t1 = threading.Thread(target=do_start) + t2 = threading.Thread(target=do_start) + t1.start() + t2.start() + try: + assert entered.wait(timeout=5.0), "winner never reached on_start" + # Winner is parked in on_start (state still NEW, _thread published). The + # loser must already be rejected — the window a state-only guard misses. + assert outcome_recorded.wait(timeout=5.0), "loser did not finish start()" + assert outcomes == [ + "rejected" + ], f"loser not rejected while winner parked: {outcomes}" + finally: + release.set() + t1.join() + t2.join() + assert sorted(outcomes) == ["ok", "rejected"], outcomes + assert ran["n"] == 1 # one worker only — on_start never ran twice + b.shutdown() + + +def test_worker_thread_is_not_daemon(): + b = ThreadedMicroBatcher(_echo) + b.start() + try: + assert b._thread.daemon is False + finally: + b.shutdown() diff --git a/docs/features/multimodal/README.md b/docs/features/multimodal/README.md index f19bbcb306cf..cd137f4ae698 100644 --- a/docs/features/multimodal/README.md +++ b/docs/features/multimodal/README.md @@ -36,6 +36,7 @@ Dynamo provides support for improving latency and throughput for vision-and-lang |---------|-------------| | **[Embedding Cache](embedding-cache.md)** | CPU-side LRU cache that skips re-encoding repeated images | | **[Encoder Disaggregation](encoder-disaggregation.md)** | Separate vision encoder worker for independent scaling | +| **[Custom Vision Encoders](custom-vision-encoder.md)** | In-process author-provided vision towers with cross-request batching | | **[Multimodal KV Routing](multimodal-kv-routing.md)** | MM-aware KV cache routing for optimal worker selection | ## Support Matrix diff --git a/docs/features/multimodal/custom-vision-encoder.md b/docs/features/multimodal/custom-vision-encoder.md new file mode 100644 index 000000000000..7e537043f5ee --- /dev/null +++ b/docs/features/multimodal/custom-vision-encoder.md @@ -0,0 +1,160 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: Custom Vision Encoders +subtitle: Run a bespoke vision tower in an aggregated Dynamo vLLM worker +--- + +A custom vision encoder lets an aggregated `dynamo.vllm` worker use an +author-provided vision tower or projector instead of vLLM's built-in multimodal +encoder. The encoder runs in the worker process, produces one embedding tensor +per image, and Dynamo inserts those embeddings at the image-placeholder positions +of a mixed prompt. + +Use this path when the language model can consume external prompt embeddings but +the vision encoder is private, experimental, or otherwise unavailable in vLLM. +It is not encoder disaggregation: the encoder and language model share one worker +process and GPU, and no embedding transfer occurs. + +## Current Scope + +| Capability | Support | +| --- | --- | +| Backend | Legacy Python `dynamo.vllm` worker | +| Topology | Aggregated only | +| Input | `image_url` content | +| Video and audio | Not supported | +| Language-model input | Mixed token IDs and prompt embeddings | +| Cross-request batching | Yes | +| CUDA graph bucket selection | Reserved for future support; current dispatch is eager | + +The custom-encoder path requires `--enable-multimodal` and +`--enable-prompt-embeds`. It is incompatible with frontend decoding, the vLLM +tokenizer mode, legacy multimodal worker roles, and non-aggregated disaggregation +modes. Dynamo validates these combinations during startup. + +## Run the Included Path + +From the repository root, launch the aggregated worker: + +```bash +bash examples/custom_encoder/launch/agg_custom.sh --gpu 0 +``` + +The launcher defaults to `Qwen/Qwen2.5-1.5B-Instruct` and the +`HitchhikersVisionEncoder`. That encoder intentionally ignores the image and +substitutes embeddings for a fixed phrase so the complete prompt-embedding path +can be checked semantically. It is a test backend, not a production vision +encoder. + +Select your own backend with a dotted Python class path: + +```bash +DYN_MODEL=my-org/my-language-model \ +DYN_ENCODER_CLASS=my_package.encoders.MyVisionEncoder \ +bash examples/custom_encoder/launch/agg_custom.sh --gpu 0 +``` + +The launcher supplies the required multimodal and prompt-embedding flags. If the +language model's chat template does not render an image-placeholder token, also +provide `DYN_CUSTOM_JINJA_TEMPLATE` or `--custom-jinja-template`. + + +The backend owns any media retrieval performed by `preprocess()`. Apply Dynamo's +[media URL policy](README.md#security-url-validation), finite network timeouts, +response-size limits, and image decode limits rather than fetching arbitrary +request URLs directly. + + +## Implement `VisionEncoderBackend` + +Subclass +`dynamo.vllm.multimodal_utils.vision_encoder_backend.VisionEncoderBackend` and +implement the following contract: + +| Member | Execution context | Responsibility | +| --- | --- | --- | +| `image_token_id` | Configuration | Hardcoded placeholder token ID used by the model or custom template | +| `build(model_id)` | Encoder actor thread, once | Load the encoder, choose its device, and initialize thread-affine resources | +| `preprocess(raw)` | Optional CPU thread pool | Fetch, decode, resize, or patchify one image and return `Preprocessed(item, cost)` | +| `forward_batch(items, target_bucket=None)` | Encoder actor thread | Run one synchronous batched forward and return one CPU tensor per item, in order | +| `close()` | Encoder actor thread, once | Release thread-affine resources | + +`forward_batch()` returns one tensor shaped +`(number_of_visual_tokens, language_model_hidden_size)` for every input item. It +must synchronize any device work and copy the results to CPU before returning; +the request coroutine consumes them from a different thread. + +Preprocessing is disabled by default. To enable it, override `preprocess()` and +set `preprocess_concurrency` to a positive value. The method must be synchronous, +thread-safe, deterministic, and CUDA-free because multiple pool threads may call +it concurrently. With `preprocess_concurrency = 0`, Dynamo skips `preprocess()` +and passes the raw image URL directly to `forward_batch()`. + +The reusable Qwen-family base and the semantic test backend are under +[`examples/custom_encoder`](https://github.com/ai-dynamo/dynamo/tree/main/examples/custom_encoder). + +## Cross-Request Batching + +Each request calls the async encoder with its images. A dedicated actor thread +collects items from all concurrent requests, invokes `forward_batch()` once for +the physical batch, and returns each result to the request and position that +submitted it. + +The batcher does not add a timer. A lone image runs as soon as the actor is free; +images that accumulate while the actor is busy are coalesced on its next pass. +Batching therefore helps when requests overlap. It does not change a serial, +single-request workload into a larger batch. + +### Choose a Batch Cost + +`Preprocessed.cost` is the amount one image contributes to +`max_batch_cost`. The batcher does not inspect image tensors or shapes. + +| Processed image regime | Recommended cost | +| --- | --- | +| Every item has the same bounded shape | `1`; the limit acts as a maximum image count | +| Native or variable resolution | Number of visual patches or tokens after preprocessing | +| Backend-specific memory relationship | A documented positive unit that remains proportional to the limiting resource | + +For variable-resolution inputs, a count-only limit can combine several maximum +resolution images and exhaust GPU memory. Compute the processed grid in +`preprocess()`, use its patch or visual-token count as `cost`, and set a finite +`max_batch_cost` that the encoder can serve alongside the language model. + +With `max_batch_cost = None`, the batcher passes every item already queued when +the actor becomes free to one `forward_batch()` call and ignores per-item costs. +Use this pass-through mode only when the backend performs its own safe sizing. +A finite limit rejects an individual item whose cost cannot fit any batch. + +The encoder and vLLM share GPU memory. Leave enough memory outside vLLM's +`gpu_memory_utilization` for the encoder weights and the peak activation memory +at `max_batch_cost`, and exercise that maximum legal batch during startup or +deployment validation. + +### Failure and Cancellation Behavior + +Validate malformed, oversized, or unsupported images in `preprocess()`. Dynamo +waits for every image in one request to finish preprocessing and submits no GPU +work if any of them fails. + +Once submitted, items from different requests may share one `forward_batch()`. +An exception from that call fails every live request represented in that physical +batch, so input-dependent failures should not be deferred to the GPU forward. + +If an awaiting request is canceled before its work passes the final dispatch +check, Dynamo tombstones its items and excludes them from later shared batches. A +synchronous `forward_batch()` already committed for execution cannot be +preempted; it finishes, and the canceled request's result is discarded. + +## Operational Checklist + +- Confirm the chat template emits exactly one placeholder span for every image. +- Confirm each returned tensor has the language model's hidden size and the row + count expected at its placeholder span. +- Use distinct images in correctness tests so reordered results cannot pass. +- Test the largest permitted image and maximum batch cost with the language model + resident on the same GPU. +- Exercise concurrent requests; a serial smoke test does not prove coalescing. +- Keep blocking media operations out of the actor thread by enabling the + preprocessing pool. diff --git a/docs/features/multimodal/multimodal-vllm.md b/docs/features/multimodal/multimodal-vllm.md index ea10b3daca82..ee4c158196ea 100644 --- a/docs/features/multimodal/multimodal-vllm.md +++ b/docs/features/multimodal/multimodal-vllm.md @@ -43,6 +43,13 @@ The main multimodal vLLM launchers in this repo are: | E/PD (Encode + PD) | CUDA | `disagg_multimodal_e_pd.sh` | No | Separate encoder and embedding-cache workflows | | E/P/D (Full Disaggregation) | CUDA | `disagg_multimodal_epd.sh` | No | Separate encode, prefill, and decode workers | +### Custom Vision Encoders + +The legacy aggregated vLLM worker can load an author-provided vision tower in +process, batch images across concurrent requests, and splice the resulting +embeddings into the language-model prompt. See [Custom Vision +Encoders](custom-vision-encoder.md) for the backend contract, launch instructions, +batch sizing guidance, and current limitations. ## Image/Video Serving diff --git a/docs/index.yml b/docs/index.yml index b9775eafd517..4efa270b039c 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -330,6 +330,8 @@ navigation: path: features/multimodal/embedding-cache.md - page: Encoder Disaggregation path: features/multimodal/encoder-disaggregation.md + - page: Custom Vision Encoders + path: features/multimodal/custom-vision-encoder.md - page: Multimodal KV Routing path: features/multimodal/multimodal-kv-routing.md - page: SGLang Multimodal