diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 07cf85eb07a9..da35b46a394f 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -3,6 +3,7 @@ import asyncio import base64 +import importlib import inspect import logging import math @@ -97,6 +98,8 @@ from .args import Config from .constants import DisaggregationMode, EmbeddingTransferMode from .engine_monitor import VllmEngineMonitor +from .multimodal_utils.async_vision_encoder import AsyncVisionEncoder +from .multimodal_utils.embed_assembler import build_mixed_embeds from .multimodal_utils.hash_utils import compute_mm_uuids_from_images from .multimodal_utils.model import ( ModelFamily, @@ -108,6 +111,7 @@ load_qwen_grid_params, ) from .multimodal_utils.prefill_worker_utils import MultiModalEmbeddingLoader +from .multimodal_utils.vision_encoder_backend import VisionEncoderBackend # Multimodal data dictionary keys IMAGE_URL_KEY: Final = "image_url" @@ -1041,6 +1045,12 @@ def __init__( ) self.embedding_loader = self.init_embedding_loader(config, encode_worker_client) + # Aggregated partial encoder. The attribute is set here so cleanup() is + # always safe, but the encoder is loaded last in __init__ (it starts a + # thread) — see _load_custom_encoder below — so a failure in the rest of + # init does not leak the batcher thread. + self._custom_encoder: Optional[AsyncVisionEncoder] = None + self.use_vllm_tokenizer = use_vllm_tokenizer self.dp_range = get_dp_range_for_worker(self.engine_client.vllm_config) @@ -1092,6 +1102,51 @@ def __init__( # dyn://..rl when --enable-rl / DYN_ENABLE_RL is set. self.rl_route_registry = RLRouteRegistry(self.runtime, logger_=logger) + # Load the custom encoder last: it starts a batcher thread, so deferring + # it until all other (fallible) setup is done means an earlier init + # failure cannot leave that thread orphaned. + self._load_custom_encoder(config) + + def _load_custom_encoder(self, config: Config) -> None: + """Import, instantiate, and load the --custom-encoder-class encoder.""" + custom_encoder_class = getattr(config, "custom_encoder_class", None) + if not custom_encoder_class: + return + # The custom encoder path only ever submits a mixed EmbedsPrompt, so fail + # fast here if prompt-embeds are disabled rather than loading the encoder + # and then rejecting every image request at runtime. + if not config.engine_args.enable_prompt_embeds: + raise ValueError( + "--custom-encoder-class requires --enable-prompt-embeds: the " + "custom encoder submits a mixed EmbedsPrompt, which the engine " + "cannot accept without prompt-embeds enabled." + ) + module_path, _, class_name = custom_encoder_class.rpartition(".") + backend_cls = getattr(importlib.import_module(module_path), class_name) + if not ( + isinstance(backend_cls, type) + and issubclass(backend_cls, VisionEncoderBackend) + ): + raise TypeError( + f"--custom-encoder-class {custom_encoder_class!r} must resolve to a " + f"VisionEncoderBackend subclass, got {backend_cls!r}." + ) + # The author writes the VisionEncoderBackend (L2); Dynamo wraps it in the + # AsyncVisionEncoder glue (L3), which owns the preprocess pool + actor + # thread + micro-batcher. load() runs backend.build() on the actor thread + # (the backend picks its own device — the worker pins it via + # CUDA_VISIBLE_DEVICES) and cleans that thread up on failure. + encoder = AsyncVisionEncoder(backend_cls()) + encoder.load(config.model) + # Assign only after a successful load so a failed load (which already shut + # its own thread down) leaves _custom_encoder None. + self._custom_encoder = encoder + logger.info( + "Loaded CustomEncoder %s from %s", + custom_encoder_class, + config.model, + ) + def _shutdown_on_engine_dead(self, e: EngineDeadError) -> NoReturn: logger.error(f"vLLM EngineDeadError: {e}") logger.warning("Initiating Dynamo Runtime shutdown.") @@ -2322,6 +2377,9 @@ async def list_loras(self, request=None): def cleanup(self): """Clean up resources including temporary directories.""" + if self._custom_encoder is not None: + # Stop the in-process encoder's batcher thread on teardown. + self._custom_encoder.shutdown() for temp_dir in self.temp_dirs: try: temp_dir.cleanup() @@ -2720,6 +2778,7 @@ def _build_prompt_from_request( multi_modal_data: Dict[str, Any] | None, log_prefix: str = "", mm_processor_kwargs: Dict[str, Any] | None = None, + mixed_embeds: tuple[torch.Tensor, list[int], list[bool]] | None = None, ) -> tuple[TokensPrompt | EmbedsPrompt | None, int | None, Dict[str, Any] | None]: """ Build a prompt from request, handling both prompt_embeds and token_ids. @@ -2731,6 +2790,9 @@ def _build_prompt_from_request( log_prefix: Prefix for log messages (e.g., "Prefill " for prefill requests) mm_processor_kwargs: Optional multimodal processor kwargs (e.g. use_audio_in_video) forwarded to the vLLM engine. + mixed_embeds: Optional ``(prompt_embeds, prompt_token_ids, + prompt_is_token_ids)`` assembled by the aggregated CustomEncoder + path. When present, takes the EmbedsPrompt fast path below. Returns: Tuple of (prompt, embedding_sequence_length, error_dict) where: @@ -2739,6 +2801,39 @@ def _build_prompt_from_request( """ embedding_sequence_length = None + # Fast path: mixed token-ids/embeds prompt from the aggregated + # CustomEncoder path, assembled in _generate_token_mode and passed in + # explicitly. The image embeds ride on the EmbedsPrompt itself and the + # request carries no multi_modal_data, so there is nothing to bind + # mm_uuids to here — the normal token path's MM-routing logic below is + # intentionally skipped for this path. + if mixed_embeds is not None: + prompt_embeds, prompt_token_ids, prompt_is_token_ids = mixed_embeds + if not self.config.engine_args.enable_prompt_embeds: + msg = ( + "CustomEncoder requires `--enable-prompt-embeds`; the " + "assembled EmbedsPrompt cannot be submitted without it." + ) + logger.error("Request %s: %s", request_id, msg) + return ( + None, + None, + { + "finish_reason": f"error: {msg}", + "token_ids": [], + }, + ) + seq_len = prompt_embeds.shape[0] + return ( + EmbedsPrompt( + prompt_embeds=prompt_embeds, + prompt_token_ids=prompt_token_ids, + prompt_is_token_ids=prompt_is_token_ids, + ), + seq_len, + None, + ) + if "prompt_embeds" in request and request["prompt_embeds"]: if not self.config.engine_args.enable_prompt_embeds: msg = ( @@ -2781,6 +2876,7 @@ def _build_prompt_from_request( "token_ids": [], }, ) + # Text-only PD + encoder-worker path. # Normal path: use token IDs. # Prefer frontend-forwarded mm_hashes for hash consistency with the # routing layer. Fall back to computing from loaded image data when @@ -3135,6 +3231,99 @@ async def generate(self, request, context): first_token = False yield chunk + async def _assemble_custom_encoder_prompt( + self, + request: Dict[str, Any], + request_id: str, + context, + mm_processor_kwargs: Dict[str, Any] | None = None, + ) -> tuple[ + tuple[torch.Tensor, list[int], list[bool]] | None, + Dict[str, Any] | None, + Dict[str, Any] | None, + ]: + """Run the in-process CustomEncoder and assemble a mixed EmbedsPrompt. + + The CustomEncoder consumes image URLs directly and emits embeds. Returns + ``(mixed_embeds, multi_modal_data, error)``: + - images present: ``(mixed_embeds, None, None)``, + - no image content: ``(None, multi_modal_data, None)`` — text-only request + loaded via the normal aggregated path, + - failure: ``(None, None, error_dict)`` for the caller to yield. + """ + # Internal invariant: callers guard on `self._custom_encoder is not None` + # before reaching here. Use an explicit raise (not assert, which is + # stripped under `python -O`) so a future mis-wire fails loudly. + if self._custom_encoder is None: + raise RuntimeError( + "_assemble_custom_encoder_prompt called without a CustomEncoder" + ) + mm_map = request.get("multi_modal_data") or {} + # CustomEncoder handles images only. Reject any non-image modality + # (video/audio/...) explicitly instead of silently dropping it. + unsupported = sorted(k for k in mm_map if k != IMAGE_URL_KEY and mm_map.get(k)) + if unsupported: + msg = ( + "CustomEncoder supports image inputs only; got " + f"unsupported multimodal data: {unsupported}" + ) + logger.error("Request %s: %s", request_id, msg) + return None, None, {"finish_reason": f"error: {msg}", "token_ids": []} + + image_items = mm_map.get(IMAGE_URL_KEY) or [] + image_urls = [ + item["Url"] + for item in image_items + if isinstance(item, dict) and "Url" in item + ] + if image_items and not image_urls: + # Image data is present but no usable URL could be extracted. + msg = ( + "CustomEncoder received image multimodal data but could not " + f"extract any URLs from {len(image_items)} item(s); each item " + "must be a dict with a 'Url' key" + ) + logger.error("Request %s: %s", request_id, msg) + return None, None, {"finish_reason": f"error: {msg}", "token_ids": []} + + if not image_urls: + # No image content at all → normal aggregated load (text-only). + multi_modal_data = await self._extract_multimodal_data( + request, request_id, context, mm_processor_kwargs=mm_processor_kwargs + ) + return None, multi_modal_data, None + + token_ids: list[int] = request.get("token_ids") or [] + # encode() is user-supplied code (any exception) and + # get_image_placeholder_token_id() can raise (unknown model family), so + # keep encode -> placeholder lookup -> assembly inside one guard: a + # failure becomes a structured request error instead of escaping the + # request coroutine and tearing down the stream. + try: + # encode() is async; AsyncVisionEncoder preprocesses off-thread and the + # ThreadedMicroBatcher coalesces concurrent calls onto one dedicated + # actor thread. The worker holds no lock — batching is the encoder's. + img_tensors: list[torch.Tensor] = await self._custom_encoder.encode( + image_urls + ) + placeholder_id = self._custom_encoder.get_image_placeholder_token_id() + prompt_embeds, mixed_token_ids, is_token_ids = build_mixed_embeds( + token_ids, img_tensors, placeholder_id + ) + except Exception as exc: + msg = f"CustomEncoder failed: {exc}" + logger.exception("Request %s: %s", request_id, msg) + return None, None, {"finish_reason": f"error: {msg}", "token_ids": []} + + logger.debug( + "Request %s: CustomEncoder assembled %d image(s) → seq_len=%d dtype=%s", + request_id, + len(img_tensors), + prompt_embeds.shape[0], + prompt_embeds.dtype, + ) + return (prompt_embeds, mixed_token_ids, is_token_ids), None, None + async def _generate_token_mode(self, request, context, request_id): """Generate tokens using internal protocol format (token-in-token-out).""" # Firstly extract disaggregated params from prefill result if available @@ -3163,6 +3352,9 @@ async def _generate_token_mode(self, request, context, request_id): multi_modal_data: Dict[str, Any] | None = None pre_rendered: Dict[str, Any] | None = None + # Assembled by the aggregated CustomEncoder path (if active) and passed + # explicitly to _build_prompt_from_request below. + mixed_embeds: tuple[torch.Tensor, list[int], list[bool]] | None = None if is_decode_only: # Decode mode: branch on model, not data. if resolve_model_family(self.config.model) is ModelFamily.QWEN_VL: @@ -3209,23 +3401,42 @@ async def _generate_token_mode(self, request, context, request_id): mm_processor_kwargs=mm_processor_kwargs, ) else: - # Fast path: check for pre-processed mm_kwargs via NIXL/SHM from frontend. - # If available, we skip image downloading AND the HF processor. - with _nvtx.annotate("mm_backend:receive_mm_kwargs", color="magenta"): - pre_rendered = await self._try_receive_mm_kwargs(request) - if pre_rendered is not None: - logger.debug( - "[mm-routing] Request %s: received pre-rendered mm_kwargs via NIXL/SHM", - request_id, + if self._custom_encoder is not None and has_mm_data: + # Aggregated CustomEncoder path takes precedence: when an encoder + # is configured it owns the image path, so don't attempt the + # NIXL/SHM pre-render receive (which would otherwise silently + # win). The encoder runs in-process (no NIXL transfer) and + # assembles a mixed EmbedsPrompt, or yields an error; pre_rendered + # stays None. A request with no images is treated as text-only. + ( + mixed_embeds, + multi_modal_data, + assemble_error, + ) = await self._assemble_custom_encoder_prompt( + request, request_id, context, mm_processor_kwargs ) + if assemble_error is not None: + yield assemble_error + return else: - # Aggregated mode: load images normally - multi_modal_data = await self._extract_multimodal_data( - request, - request_id, - context, - mm_processor_kwargs=mm_processor_kwargs, - ) + # Fast path: check for pre-processed mm_kwargs via NIXL/SHM from + # frontend. If available, we skip image downloading AND the HF + # processor. + with _nvtx.annotate("mm_backend:receive_mm_kwargs", color="magenta"): + pre_rendered = await self._try_receive_mm_kwargs(request) + if pre_rendered is not None: + logger.debug( + "[mm-routing] Request %s: received pre-rendered mm_kwargs via NIXL/SHM", + request_id, + ) + else: + # Aggregated mode: load images normally + multi_modal_data = await self._extract_multimodal_data( + request, + request_id, + context, + mm_processor_kwargs=mm_processor_kwargs, + ) # Build prompt from request. `prompt` is either a pre-rendered # MultiModalInput dict (fast path) or a TokensPrompt/EmbedsPrompt from @@ -3254,6 +3465,7 @@ async def _generate_token_mode(self, request, context, request_id): request_id, multi_modal_data, mm_processor_kwargs=mm_processor_kwargs, + mixed_embeds=mixed_embeds, ) if error is not None: yield error diff --git a/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py b/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py new file mode 100644 index 000000000000..f25f772dc73e --- /dev/null +++ b/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Async glue (L3) between the worker's event loop and a ``VisionEncoderBackend``. + +``AsyncVisionEncoder`` is the **Dynamo-owned** layer the worker talks to. It +turns the author's synchronous, thread-affine backend (L2) into an awaitable +``encode(raws) -> list[tensor]`` by: + +- running ``backend.preprocess`` **off the event loop** on a bounded + ``ThreadPoolExecutor`` (CPU-heavy fetch / resize / patchify must not serialize + on the GPU actor thread); +- enforcing **request-level atomicity** (A5): a gather-barrier between preprocess + and submit — ``encode`` waits for *every* image's preprocess to settle and only + submits if **all** succeed; on any failure it submits nothing (zero GPU work) + and raises the request-level error, so a text-only LM never sees a partial + result; +- handing the preprocessed items (with their off-thread-computed 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, Optional + +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 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 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__( + self, + backend: VisionEncoderBackend[RawT, ItemT], + *, + preprocess_concurrency: int = 4, + name: str = "vision-encoder", + ) -> None: + if preprocess_concurrency < 1: + raise ValueError("preprocess_concurrency must be >= 1") + self._backend = backend + self._preprocess_concurrency = preprocess_concurrency + self._name = name + self._batcher: Optional[ThreadedMicroBatcher] = None + self._pool: Optional[ThreadPoolExecutor] = None + + # ---- lifecycle --------------------------------------------------------- + + def load(self, model_id: str) -> None: + """Start the actor thread (running ``backend.build`` on it) and fail fast. + + Re-raises any build error, then ``validate``s the placeholder id so a + misconfigured encoder errors at startup instead of on the first request. + Single-shot: a second ``load()`` raises rather than orphaning the first + batcher's (non-daemon) worker thread and model. + """ + 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 misconfigured buckets / 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: + self._pool = ThreadPoolExecutor( + max_workers=self._preprocess_concurrency, + thread_name_prefix=f"{self._name}-pre", + ) + self._batcher = ThreadedMicroBatcher( + self._backend.forward_batch, + max_batch_cost=self._backend.max_batch_cost, + buckets=self._backend.buckets, + 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() + raise + + def validate(self) -> None: + """Fail-fast check run by ``load`` after ``build``: the author hardcoded a + usable ``image_token_id``.""" + tid = getattr(self._backend, "image_token_id", None) + if not isinstance(tid, int) or isinstance(tid, bool): + raise ValueError( + "VisionEncoderBackend.image_token_id must be a hardcoded int (the " + f"image placeholder token id); got {tid!r}" + ) + + def get_image_placeholder_token_id(self) -> int: + """The token id marking image positions (the backend's hardcoded value).""" + return self._backend.image_token_id + + # ---- request path ------------------------------------------------------ + + async def encode(self, raws: List[RawT]) -> List[torch.Tensor]: + """Preprocess (off-loop, A5 barrier) then batched-encode; all-or-nothing. + + Returns one ``(n_visual_tokens, lm_hidden_dim)`` tensor per raw input, in + order. Raises if any image's preprocess fails (submitting nothing) or if + the batched forward fails. + """ + 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 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] + costs = [p.cost for p in preprocessed] + return await self._batcher.submit(items, costs) + + def shutdown(self) -> None: + """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) 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..0daef21459ec --- /dev/null +++ b/components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py @@ -0,0 +1,601 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Coalesce concurrent async calls into batched calls of a blocking fn on one thread. + +``ThreadedMicroBatcher`` is the generic execution mechanism (L1) behind +``AsyncVisionEncoder`` — it has no model/vision knowledge and stays torch-free. +It owns: + +- a **dedicated worker thread** that runs an optional ``on_start`` (e.g. build + + CUDA-graph capture) and then every ``fn`` call, so anything thread-affine + (CUDA graphs, the current device/stream) is captured and replayed on the same + thread; an optional ``on_stop`` runs on that same thread at teardown; +- a **coalescing micro-batcher**: items from concurrent ``submit()`` calls are + pooled and split into batches whose summed ``cost`` stays within + ``max_batch_cost`` (a compute/token budget, not a raw count). Packing is + **one-dimensional** — by scalar ``cost`` alone; the batcher never inspects item + shape; +- a **graph ladder** (optional ``buckets``): when set, the batcher rounds each + batch's ``sum(cost)`` **up to the nearest rung** and passes it as + ``target_bucket`` so ``fn`` can pad to that rung and replay its captured graph. + ``buckets=None`` ⇒ eager: ``target_bucket=None``. + +The caller speaks in opaque items plus a per-item scalar ``cost`` (int), computed +once off-thread (see ``Preprocessed``); the batcher never interprets the items, so +all model knowledge stays in the caller. Final padding of a batch to a captured +CUDA-graph shape is the ``fn``'s job (it owns +the model), not the batcher's — the batcher only decides *which rung*. + +Coalescing window — **eager drain-on-completion, no timer** (the design default): +whenever the worker is free it pulls everything queued and runs it, then repeats. +A lone item runs on the next free iteration (liveness for free); batch size +auto-scales with load (arrivals during ``fn`` pile up and are scooped next loop). +``max_wait_ms`` is an opt-in, default-off knob for a deliberate accumulate hold. + +Concurrency contract (cross-thread correctness): + +- The event-loop bridge is a ``concurrent.futures.Future`` adapted with + ``asyncio.wrap_future`` — so callers on **any** event loop work, and cancelling + an ``await submit(...)`` retires the request: its not-yet-run items are + tombstoned (never reach ``fn``) and its admission released, and the caller only + returns once the worker is done with every item (``retired``). +- A request is finalised **per item**: admission release, ``completion``, and + ``retired`` fire exactly once, only after *all* of the request's items have been + delivered, failed, or tombstoned — never on the first item of a multi-batch + request. Every request-state transition is under one short lock, so the worker + and a concurrent ``shutdown()`` cannot race on ``remaining``. +- The worker runs under a **supervisor**: any unexpected crash fails every live + request (no hung awaiter) and moves the batcher to ``FAILED``. +- Admission is bounded by an optional ``max_outstanding_cost``; state + admission + are mutated together under the lock (bookkeeping only — never held across + ``fn`` / ``join``). +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import logging +import queue +import threading +import time +from dataclasses import dataclass +from enum import Enum, auto +from typing import Callable, Generic, List, Optional, Sequence, TypeVar + +logger = logging.getLogger(__name__) + +T = TypeVar("T") +R = TypeVar("R") + +# Sentinel pushed onto the queue to stop the worker thread. +_SHUTDOWN = object() +# Sentinel distinguishing "no result" (failed / tombstoned item) from a real +# ``None`` result returned by ``fn``. +_NO_RESULT = object() + + +class BatcherOverloaded(RuntimeError): + """Raised by ``submit()`` when admitting the request would exceed + ``max_outstanding_cost`` (accepted-but-incomplete cost).""" + + +class _State(Enum): + NEW = auto() + RUNNING = auto() + CLOSING = auto() + CLOSED = auto() + 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`` and ``retired`` are thread-safe ``concurrent.futures.Future``s. + 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]]" + retired: "concurrent.futures.Future[None]" + results: List[Optional[R]] + remaining: int + cost_total: int + cancelled: bool = False + 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], target_bucket) -> list[result]`` on a dedicated thread, + coalescing concurrent ``submit()`` calls into cost-bounded batches. + + Args: + fn: Batched work; one result per item, in order. Called as + ``fn(items, target_bucket)`` on the worker thread — ``target_bucket`` + is the ladder rung to pad to (``None`` in eager mode). + 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). With ``buckets`` set, + ``None`` derives the ceiling as ``max(buckets)``. + buckets: Optional sorted graph ladder. When set, the batcher rounds a + batch's ``sum(cost)`` up to the nearest rung and passes it as + ``target_bucket``; ``None``/empty ⇒ eager (``target_bucket=None``). + max_wait_ms: Opt-in coalescing hold after the first item arrives. Default + ``0`` ⇒ eager drain-on-completion (no timer). + on_start: Optional callable run once on the worker thread before serving + (model build / warmup / graph capture); 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. + max_outstanding_cost: Optional admission ceiling on accepted-but-incomplete + cost; ``submit()`` raises ``BatcherOverloaded`` when exceeded. + name: Worker thread name. + join_timeout_s: Seconds ``shutdown()`` waits for an in-flight ``fn``. + """ + + def __init__( + self, + fn: Callable[[List[T], Optional[int]], List[R]], + *, + max_batch_cost: Optional[int] = None, + buckets: Optional[Sequence[int]] = None, + max_wait_ms: float = 0.0, + on_start: Optional[Callable[[], None]] = None, + on_stop: Optional[Callable[[], None]] = None, + max_outstanding_cost: Optional[int] = None, + name: str = "micro-batcher", + join_timeout_s: float = 10.0, + ) -> None: + if max_batch_cost is not None and max_batch_cost < 1: + raise ValueError("max_batch_cost must be >= 1 (or None for pass-through)") + if max_outstanding_cost is not None and max_outstanding_cost < 1: + raise ValueError("max_outstanding_cost must be >= 1") + self._buckets = self._validate_buckets(buckets, max_batch_cost) + # Graph mode needs a bounded ceiling, so derive it from the ladder when the + # author left it None (pass-through is only meaningful in eager mode). + if self._buckets is not None and max_batch_cost is None: + max_batch_cost = self._buckets[-1] + self._fn = fn + self._max_batch_cost = max_batch_cost + self._max_wait_s = max_wait_ms / 1000.0 + self._on_start = on_start + self._on_stop = on_stop + self._max_outstanding_cost = max_outstanding_cost + self._name = name + self._join_timeout_s = join_timeout_s + + self._queue: queue.Queue = queue.Queue() + self._ready = threading.Event() + self._terminated = threading.Event() + self._start_error: Optional[BaseException] = None + self._terminal_error: Optional[BaseException] = None + # Guards _state, _outstanding, _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._outstanding = 0 + self._live: set[_Request] = set() + self._thread: Optional[threading.Thread] = None + + @staticmethod + def _validate_buckets( + buckets: Optional[Sequence[int]], max_batch_cost: Optional[int] + ) -> Optional[tuple]: + """Normalise the ladder to a sorted tuple of positive ints (or None). + + When ``max_batch_cost`` is set, the ladder must cover it so every packed + batch has a rung to round up to: ``max(buckets) >= max_batch_cost``. When + it is ``None``, the ceiling is derived as ``max(buckets)`` by the caller.""" + if not buckets: + return None + rungs = tuple(sorted(int(b) for b in buckets)) + if any(b < 1 for b in rungs): + raise ValueError("buckets must be positive ints") + if max_batch_cost is not None and rungs[-1] < max_batch_cost: + raise ValueError( + f"max(buckets)={rungs[-1]} < max_batch_cost={max_batch_cost}; the " + "ladder must cover the dispatch ceiling (set max_batch_cost to " + "max(buckets) for a graphed encoder)" + ) + return rungs + + def _target_bucket(self, batch_cost: int) -> Optional[int]: + """Round a batch's summed cost up to the nearest ladder rung (None=eager).""" + if self._buckets is None: + return None + for rung in self._buckets: # sorted ascending + if rung >= batch_cost: + return rung + return self._buckets[-1] # unreachable: max(buckets) >= max_batch_cost + + # ---- lifecycle --------------------------------------------------------- + + def start(self) -> None: + """Start the worker thread, run ``on_start`` on it, and re-raise its error. + + Single-shot: a second ``start()`` raises rather than spawning a second + consumer / orphaning the first thread.""" + with self._lock: + if self._thread is not None: + 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) a racing + # shutdown() can never join() an unstarted thread and (b) a failed + # thread creation leaves no stale unstarted thread behind. + thread = threading.Thread(target=self._run, name=self._name, daemon=False) + try: + thread.start() + except BaseException: + self._state = _State.FAILED + raise + self._thread = thread + self._ready.wait() + if self._start_error is not None: + # 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 self._start_error + with self._lock: + if self._state is _State.NEW: + self._state = _State.RUNNING + elif self._state is not _State.RUNNING: + # a concurrent shutdown() closed us during startup + raise RuntimeError("ThreadedMicroBatcher shut down during start()") + + 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. + + Cancellation-safe: cancelling the await tombstones not-yet-run items, + releases admission, and returns only once the worker has retired the + request (so the caller may then release the items' backing memory). + """ + 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(), + retired=concurrent.futures.Future(), + results=[None] * len(items), + remaining=len(items), + cost_total=sum(costs), + ) + 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.RUNNING: + pass + elif self._state is _State.FAILED: + raise RuntimeError( + "ThreadedMicroBatcher.submit() after worker failure" + ) from self._terminal_error + else: # NEW / CLOSING / CLOSED + raise RuntimeError( + "ThreadedMicroBatcher.submit() called after shutdown()" + ) + if ( + self._max_outstanding_cost is not None + and self._outstanding + request.cost_total > self._max_outstanding_cost + ): + raise BatcherOverloaded( + f"submit cost {request.cost_total} would exceed outstanding " + f"budget {self._max_outstanding_cost} " + f"(in flight: {self._outstanding})" + ) + self._outstanding += request.cost_total + self._live.add(request) + for work in works: + self._queue.put(work) + try: + return await asyncio.wrap_future(request.completion) + except asyncio.CancelledError: + # Tombstone: the worker skips not-yet-run items and finalises (releasing + # admission). Wait — through repeated cancellation — until the worker is + # provably done with this request's items before propagating, so the + # caller can safely drop them. + with self._lock: + request.cancelled = True + retirement = asyncio.wrap_future(request.retired) + while not retirement.done(): + try: + await asyncio.shield(retirement) + except asyncio.CancelledError: + continue + raise + + def shutdown(self) -> None: + """Stop the worker, failing not-yet-run items. Idempotent; retries the + join if a slow ``fn`` is still in flight.""" + if self._thread is None: + return + to_fail: List[_Work] = [] + with self._lock: + if self._state not in (_State.CLOSING, _State.CLOSED, _State.FAILED): + self._state = _State.CLOSING + to_fail = self._drain_queue_locked() # fail queued, then signal stop + self._queue.put(_SHUTDOWN) + self._consume_all(to_fail, RuntimeError("ThreadedMicroBatcher shut down")) + self._thread.join(timeout=self._join_timeout_s) + if self._thread.is_alive(): + logger.warning( + "ThreadedMicroBatcher(%s): worker still running after %gs; will " + "reap on a later call", + self._name, + self._join_timeout_s, + ) + return + with self._lock: + if self._state is _State.CLOSING: + self._state = _State.CLOSED + leftover = self._drain_queue_locked() # belt-and-suspenders + self._consume_all(leftover, RuntimeError("ThreadedMicroBatcher shut down")) + + # ---- 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._start_error = exc + self._ready.set() + self._terminated.set() + return + self._ready.set() + 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() + self._terminated.set() + + 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. + + Default (``max_wait_ms=0``): drain only what is immediately available — no + timed hold. With ``max_wait_ms>0``, keep draining within that window.""" + first = self._queue.get() + if first is _SHUTDOWN: + return None + works: List[_Work] = [first] + if self._max_wait_s > 0: + deadline = time.monotonic() + self._max_wait_s + while True: + timeout = deadline - time.monotonic() + if timeout <= 0: + break + try: + item = self._queue.get(timeout=timeout) + except queue.Empty: + break + if item is _SHUTDOWN: + self._queue.put(_SHUTDOWN) # drain this round, stop next loop + break + works.append(item) + return works + # 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 (cancelled / failed / done) items are dropped before batching — + a cancelled or 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 + if batch: + self._run_batch(batch) + + def _run_batch(self, batch: List[_Work]) -> None: + # Re-filter immediately before fn: a request may have been cancelled (or + # failed by a sibling batch) after grouping, so cancellation is synced to + # each fn call rather than only 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] + target_bucket = self._target_bucket(sum(w.cost for w in runnable)) + try: + results = self._fn(items, target_bucket) + 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``.""" + return req.done or req.cancelled or req.error is not None + + def _consume(self, work: _Work, result: object = _NO_RESULT, 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, ``completion`` and + ``retired`` are exactly-once even when items span batches or threads.""" + 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 result is not _NO_RESULT and req.error is None and not req.cancelled: + req.results[work.index] = result + req.remaining -= 1 + if req.remaining == 0: + req.done = True + self._live.discard(req) + self._outstanding -= req.cost_total + 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._outstanding -= req.cost_total + 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) + elif not req.cancelled: + req.completion.set_result(list(req.results)) + # cancelled + no error: leave completion (already cancelled by the waiter) + except concurrent.futures.InvalidStateError: + pass # caller already cancelled the future + try: + req.retired.set_result(None) # caller's cancel path awaits this + except concurrent.futures.InvalidStateError: + pass + + 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).""" + drained: List[_Work] = [] + while True: + try: + item = self._queue.get_nowait() + except queue.Empty: + return drained + if item is _SHUTDOWN: + continue + drained.append(item) + + def _consume_all(self, works: List[_Work], exc: BaseException) -> None: + for work in works: + self._consume(work, error=exc) 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 new file mode 100644 index 000000000000..f9d1ad40e83f --- /dev/null +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for dynamo.vllm.multimodal_utils.async_vision_encoder. + +Pin the glue contract: build / forward / close run on one actor thread; encode +returns one tensor per raw; the A5 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 threading + +import pytest +import torch + +from dynamo.vllm.multimodal_utils.async_vision_encoder import AsyncVisionEncoder +from dynamo.vllm.multimodal_utils.vision_encoder_backend import ( + Preprocessed, + VisionEncoderBackend, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.multimodal, +] + + +class _FakeBackend(VisionEncoderBackend): + """A CPU-only fake backend; records its threads.""" + + max_batch_cost = 8 + buckets = None + image_token_id = 151655 + + def __init__(self, *, fail_on=None): + self.fail_on = set(fail_on or ()) + self.build_thread = None + self.close_thread = None + self.closed = False + self.model_id = None + self.forward_threads: list[int] = [] + + def build(self, model_id): + self.build_thread = threading.get_ident() + self.model_id = model_id + + def preprocess(self, raw): + if raw in self.fail_on: + raise ValueError(f"bad input {raw}") + return Preprocessed(item=raw, cost=1) + + def forward_batch(self, items, target_bucket=None): + self.forward_threads.append(threading.get_ident()) + return [torch.full((2, 4), float(len(str(it)))) for it in items] + + def close(self): + self.close_thread = threading.get_ident() + self.closed = True + + +async def test_encode_returns_one_tensor_per_raw(): + enc = AsyncVisionEncoder(_FakeBackend()) + enc.load("m") + try: + out = await enc.encode(["a", "bb", "ccc"]) + assert len(out) == 3 + assert all(t.shape == (2, 4) for t in out) + finally: + enc.shutdown() + + +async def test_a5_barrier_fails_atomically_with_no_gpu_work(): + be = _FakeBackend(fail_on={"bad"}) + enc = AsyncVisionEncoder(be) + enc.load("m") + try: + with pytest.raises(ValueError, match="bad input"): + await enc.encode(["good", "bad"]) + assert be.forward_threads == [] # nothing was submitted + finally: + enc.shutdown() + + +async def test_build_and_forward_share_one_non_main_thread(): + be = _FakeBackend() + enc = AsyncVisionEncoder(be) + enc.load("m") + try: + await enc.encode(["x"]) + assert be.build_thread is not None + assert set(be.forward_threads) == {be.build_thread} + assert be.build_thread != threading.get_ident() + finally: + enc.shutdown() + + +async def test_load_resolves_placeholder_and_passes_model_id(): + be = _FakeBackend() + enc = AsyncVisionEncoder(be) + enc.load("my-model") + try: + assert enc.get_image_placeholder_token_id() == 151655 + assert be.model_id == "my-model" + finally: + enc.shutdown() + + +async def test_encode_empty_returns_empty(): + enc = AsyncVisionEncoder(_FakeBackend()) + enc.load("m") + try: + assert await enc.encode([]) == [] + finally: + enc.shutdown() + + +async def test_encode_before_load_raises(): + enc = AsyncVisionEncoder(_FakeBackend()) + with pytest.raises(RuntimeError, match="before load"): + await enc.encode(["a"]) + + +def test_load_twice_raises(): + enc = AsyncVisionEncoder(_FakeBackend()) + enc.load("m") + try: + with pytest.raises(RuntimeError, match="called twice"): + enc.load("m") + finally: + enc.shutdown() + + +def test_shutdown_runs_backend_close_on_actor_thread(): + be = _FakeBackend() + enc = AsyncVisionEncoder(be) + enc.load("m") + enc.shutdown() + assert be.closed is True + assert be.close_thread == be.build_thread # close on the actor thread + + +def test_load_fails_fast_on_build_error_and_reaps_thread(): + class _BadBuild(_FakeBackend): + def build(self, model_id): + raise RuntimeError("build failed") + + enc = AsyncVisionEncoder(_BadBuild()) + with pytest.raises(RuntimeError, match="build failed"): + enc.load("m") + assert enc._batcher is not None and not enc._batcher._thread.is_alive() + + +def test_load_fails_fast_on_missing_image_token_id(): + class _NoTokenId(_FakeBackend): + 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._batcher is not None and not enc._batcher._thread.is_alive() + + +def test_shutdown_before_load_is_safe(): + AsyncVisionEncoder(_FakeBackend()).shutdown() # no-op, no raise + + +def test_preprocess_concurrency_must_be_positive(): + with pytest.raises(ValueError, match="preprocess_concurrency"): + AsyncVisionEncoder(_FakeBackend(), preprocess_concurrency=0) + + +def test_load_reaps_pool_if_batcher_ctor_fails(): + """A backend exposing a ladder the batcher rejects must not leak the pool.""" + + class _BadBuckets(_FakeBackend): + max_batch_cost = 8 + buckets = [2, 4] # max(buckets) < max_batch_cost → batcher ctor raises + + enc = AsyncVisionEncoder(_BadBuckets()) + with pytest.raises(ValueError, match="ladder must cover"): + enc.load("m") + assert enc._pool is not None and enc._pool._shutdown is True + assert enc._batcher is None diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_embeds_passthrough.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_embeds_passthrough.py new file mode 100644 index 000000000000..352c719deb26 --- /dev/null +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_embeds_passthrough.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit smoke for the client-supplied-embeddings path (PR6). + +Proves the path works with no new Dynamo runtime code: a client encodes per-image +embeddings as a safetensors ``data:`` URI, the test-only ``EmbedsPassthroughEncoder`` +decodes them (preprocess) and passes them through (forward_batch, identity), and +they splice cleanly via ``build_mixed_embeds``. CPU-only, no model/LM. +""" + +import pytest +import torch + +from dynamo.vllm.multimodal_utils.async_vision_encoder import AsyncVisionEncoder +from dynamo.vllm.multimodal_utils.embed_assembler import build_mixed_embeds +from tests.utils.embeds_passthrough_encoder import ( + EmbedsPassthroughEncoder, + decode_embeds_data_uri, + encode_embeds_data_uri, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.multimodal, +] + + +class _StubNoBuild(EmbedsPassthroughEncoder): + """Skip the real tokenizer load — the embeds stub only decodes + passes through, + and image_token_id is hardcoded (151655 via QwenVisionEncoderBackend).""" + + def build(self, model_id): + pass + + +def test_encode_decode_roundtrip_preserves_shape_dtype_values(): + t = torch.randn(5, 8, dtype=torch.float32) + out = decode_embeds_data_uri(encode_embeds_data_uri(t)) + assert out.shape == (5, 8) and out.dtype == torch.float32 + assert torch.equal(out, t) + + +def test_encode_preserves_bfloat16(): + t = torch.randn(3, 4).to(torch.bfloat16) + out = decode_embeds_data_uri(encode_embeds_data_uri(t)) + assert out.dtype == torch.bfloat16 and out.shape == (3, 4) + + +def test_encode_requires_2d(): + with pytest.raises(ValueError, match="2D"): + encode_embeds_data_uri(torch.randn(8)) + + +def test_decode_rejects_non_data_uri(): + with pytest.raises(ValueError, match="data: URI"): + decode_embeds_data_uri("https://example.com/img.png") + + +def test_decode_rejects_malformed_payload(): + with pytest.raises(ValueError, match="no base64 payload"): + decode_embeds_data_uri("data:application/x-dynamo-embeds;base64,") + + +def test_image_token_id_is_hardcoded(): + # Inherited from QwenVisionEncoderBackend (Qwen <|image_pad|>), no tokenizer. + assert EmbedsPassthroughEncoder.image_token_id == 151655 + + +async def test_passthrough_through_async_encoder(): + """encode([data_uri]) decodes the client embedding and returns it unchanged.""" + enc = AsyncVisionEncoder(_StubNoBuild()) + enc.load("fake-model") + try: + embeds = torch.randn(7, 6) + out = await enc.encode([encode_embeds_data_uri(embeds)]) + assert len(out) == 1 + assert torch.equal(out[0], embeds) # identity passthrough + assert enc.get_image_placeholder_token_id() == 151655 + finally: + enc.shutdown() + + +async def test_bad_embeds_uri_fails_request_atomically(): + """A non-embeds URL fails in preprocess (A5 barrier) — no GPU work, clean error.""" + enc = AsyncVisionEncoder(_StubNoBuild()) + enc.load("fake-model") + try: + with pytest.raises(ValueError, match="data: URI"): + await enc.encode(["http://not-an-embedding"]) + finally: + enc.shutdown() + + +def test_decoded_embeds_splice_into_mixed_prompt(): + """The decoded client embeds land at the placeholder rows via build_mixed_embeds.""" + embeds = torch.randn(3, 4) + decoded = decode_embeds_data_uri(encode_embeds_data_uri(embeds)) + placeholder_id = EmbedsPassthroughEncoder.image_token_id + token_ids = [10, 11, placeholder_id, 12] + prompt_embeds, out_ids, is_token_ids = build_mixed_embeds( + token_ids, [decoded], placeholder_id + ) + assert out_ids == [10, 11, placeholder_id, placeholder_id, placeholder_id, 12] + assert is_token_ids == [True, True, False, False, False, True] + assert torch.equal(prompt_embeds[2:5].to(torch.float32), embeds.to(torch.float32)) 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..13199341ec39 --- /dev/null +++ b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py @@ -0,0 +1,493 @@ +# 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 same-bucket batches, the graph ladder rounds a +batch's cost up to a rung (target_bucket), eager-drain pulls all queued work when +free, errors reach every awaiting caller, and the shutdown lifecycle behaves. + +``fn`` is ``fn(items, target_bucket)``; ``cost`` is a precomputed scalar that rides +on ``submit(items, costs)`` (one-dimensional packing — no bucket_key). +""" + +import asyncio +import threading + +import pytest + +from dynamo.vllm.multimodal_utils.threaded_micro_batcher import ( + BatcherOverloaded, + ThreadedMicroBatcher, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.multimodal, +] + + +def _echo(items, target_bucket=None): + return list(items) + + +class _Recorder: + """fn that records the threads it ran on and the batches / target_buckets it saw.""" + + def __init__(self): + self.threads: list[int] = [] + self.batches: list[list] = [] + self.target_buckets: list = [] + 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, target_bucket=None): + self.threads.append(threading.get_ident()) + self.batches.append(list(items)) + self.target_buckets.append(target_bucket) + return [("r", x) for x in items] + + +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_concurrent_submits_coalesce(): + rec = _Recorder() + b = ThreadedMicroBatcher(rec.fn, max_wait_ms=200.0) + b.start() + try: + results = await asyncio.gather(*(b.submit(["u"]) for _ in range(5))) + assert all(len(r) == 1 for r in results) + assert sum(len(batch) for batch in rec.batches) == 5 + assert max(len(batch) for batch in rec.batches) >= 2 # coalesced + assert len(rec.batches) < 5 + finally: + b.shutdown() + + +async def test_eager_drain_pulls_all_queued_when_free(): + """Default (max_wait_ms=0) 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, target_bucket=None): + 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.""" + rec = _Recorder() + b = ThreadedMicroBatcher(rec.fn, max_wait_ms=200.0, max_batch_cost=5) + b.start() + try: + await b.submit([3, 3, 1], costs=[3, 3, 1]) # 3 | 3,1 → two batches + assert all(sum(batch) <= 5 for batch in rec.batches) + assert sum(len(batch) for batch in rec.batches) == 3 + finally: + b.shutdown() + + +async def test_target_bucket_rounds_packed_cost_up_to_nearest_rung(): + """Graph mode: the batcher rounds a batch's sum(cost) up to the nearest rung + and passes it as target_bucket.""" + rec = _Recorder() + b = ThreadedMicroBatcher( + rec.fn, max_batch_cost=8, buckets=[2, 4, 8], max_wait_ms=200.0 + ) + b.start() + try: + # One coalesced batch of cost 3 → rounds up to rung 4. + await b.submit(["x", "y", "z"], costs=[1, 1, 1]) + assert rec.batches == [["x", "y", "z"]] + assert rec.target_buckets == [4] + finally: + b.shutdown() + + +async def test_eager_mode_passes_none_target_bucket(): + rec = _Recorder() + b = ThreadedMicroBatcher(rec.fn) # buckets=None ⇒ eager + b.start() + try: + await b.submit(["a"]) + assert rec.target_buckets == [None] + finally: + b.shutdown() + + +def test_buckets_below_max_batch_cost_rejected(): + with pytest.raises(ValueError, match="ladder must cover"): + ThreadedMicroBatcher(_echo, max_batch_cost=8, buckets=[2, 4]) + + +def test_buckets_derive_max_batch_cost_when_none(): + """Graph mode with no explicit budget derives the ceiling from the ladder.""" + b = ThreadedMicroBatcher(_echo, buckets=[2, 4, 8]) # max_batch_cost=None + assert b._max_batch_cost == 8 + + +async def test_max_batch_cost_none_is_passthrough(): + """Default (max_batch_cost=None): no cap and no per-item ceiling — the whole + drained same-bucket group runs as ONE fn call regardless of summed cost.""" + rec = _Recorder() + b = ThreadedMicroBatcher(rec.fn, max_wait_ms=200.0) # max_batch_cost=None + b.start() + try: + # Big per-item costs that would be split (or rejected) under any finite cap. + out = await b.submit([1, 2, 3, 4], costs=[1000, 1000, 1000, 1000]) + assert len(out) == 4 + assert rec.batches == [[1, 2, 3, 4]] # one un-split batch + assert rec.target_buckets == [None] # eager (no ladder) + finally: + b.shutdown() + + +async def test_error_reaches_every_caller(): + def boom(items, target_bucket=None): + raise ValueError("boom") + + b = ThreadedMicroBatcher(boom, max_wait_ms=50.0) + 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_wrong_result_count_raises(): + b = ThreadedMicroBatcher(lambda items, target_bucket=None: [], max_wait_ms=10.0) + 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, target_bucket=None): + 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_cancelled_submit_is_retired_and_releases_admission(): + """Cancelling an await retires the request and frees its admission cost.""" + entered = threading.Event() + release = threading.Event() + + def blocking(items, target_bucket=None): + entered.set() + release.wait(timeout=5.0) + return [("r", x) for x in items] + + b = ThreadedMicroBatcher(blocking, max_outstanding_cost=10) + b.start() + try: + first = asyncio.ensure_future(b.submit(["a"])) # occupies the worker + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + + second = asyncio.ensure_future(b.submit(["b"])) # queued behind first + await asyncio.sleep(0.05) + assert b._outstanding == 2 # both admitted (cost 1 each) + + second.cancel() + release.set() # let the worker finish "a", then collect+retire cancelled "b" + with pytest.raises(asyncio.CancelledError): + await second + assert len(await first) == 1 + assert b._outstanding == 0 # cancelled request's admission released + finally: + b.shutdown() + + +async def test_max_outstanding_cost_rejects_when_full(): + """submit() raises BatcherOverloaded once accepted-but-incomplete cost is full.""" + entered = threading.Event() + release = threading.Event() + + def blocking(items, target_bucket=None): + entered.set() + release.wait(timeout=5.0) + return [("r", x) for x in items] + + b = ThreadedMicroBatcher(blocking, max_outstanding_cost=1) + b.start() + try: + first = asyncio.ensure_future(b.submit(["a"])) # cost 1 fills the budget + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + + with pytest.raises(BatcherOverloaded): + await b.submit(["b"]) + + release.set() + assert len(await first) == 1 + finally: + b.shutdown() + + +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_and_releases(): + """A multi-item request split across batches where the FIRST batch raises + fails the whole request exactly once, releases all of its admission, and the + later sibling item is tombstoned — it never reaches fn.""" + seen: list = [] + + def fn(items, target_bucket=None): + 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_wait_ms=200.0, max_batch_cost=1, max_outstanding_cost=10 + ) + b.start() + try: + with pytest.raises(ValueError, match="boom"): + await b.submit(["bad", "good"], costs=[1, 1]) + assert b._outstanding == 0 + 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, target_bucket=None): + 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_wait_ms=50.0, 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_worker_thread_is_not_daemon(): + b = ThreadedMicroBatcher(_echo) + b.start() + try: + assert b._thread.daemon is False + finally: + b.shutdown() diff --git a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py index 1f66f4880aa1..a2c041543ea5 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py @@ -662,6 +662,9 @@ def _make_decode_handler( handler.input_param_manager = MagicMock() handler.input_param_manager.get_extra_params.return_value = {} handler._deferred_aborts = {} + # Real BaseWorkerHandler.__init__ (patched out above) sets this; the + # aggregated branch in _generate_token_mode reads it, so mirror the default. + handler._custom_encoder = None return handler diff --git a/examples/custom_encoder/README.md b/examples/custom_encoder/README.md new file mode 100644 index 000000000000..4b51c34e1dc7 --- /dev/null +++ b/examples/custom_encoder/README.md @@ -0,0 +1,96 @@ + + +# Custom vision encoders for the aggregated `dynamo.vllm` worker + +Run a **text-only LLM** in vLLM and plug in your **own** vision encoder (bespoke +ViT/projector weights, your own CUDA-graph capture, your own batched forward). +Dynamo intercepts image inputs, runs them through your encoder, and splices the +resulting image embeds into a mixed `EmbedsPrompt` fed to the text LM — no +separate encode worker, no NIXL transfer. + +You implement one class, `VisionEncoderBackend` +(`dynamo.vllm.multimodal_utils.vision_encoder_backend`), and point the worker at +it: + +```bash +python -m dynamo.vllm --model \ + --custom-encoder-class my_pkg.encoders.MyEncoder \ + --enable-multimodal --enable-prompt-embeds +``` + +The author contract (see the class docstring for the full surface): + +- `build(model_id, device)` — load weights/tokenizer on the actor thread. +- `preprocess(raw) -> Preprocessed{item, cost, bucket_key}` — off-thread, CUDA-free. +- `forward_batch(items, target_bucket=None) -> list[torch.Tensor]` — actor thread; + one `(n_visual_tokens, lm_hidden_dim)` tensor per image, in input order. +- `get_image_placeholder_token_id()` — Qwen-family encoders get this free by + subclassing `QwenVisionEncoderBackend` (`qwen_vision_encoder.py`). + +Worked examples in this directory: `hitchhikers_vision_encoder.py` (a fake encoder +that returns a fixed phrase's embeddings — a semantic smoke that answers "42") and +`qwen3vl_vit_encoder.py` (the real Qwen3-VL vision tower with a CUDA-graph bucket +ladder). + +## Sending pre-computed embeddings (no ViT in Dynamo) + +If you have already computed per-image embeddings on the client (e.g. CLIP/ViT +outputs) and only need a small **projector** (`embed_dim → lm_hidden_dim`), you do +**not** need Dynamo to run a vision tower. Send each embedding **inline** as a +safetensors `data:` URI on a normal `image_url` content part, and write a backend +whose `preprocess` decodes it and whose `forward_batch` projects it. No new Dynamo +request format is required today (a first-class `image_embeds` request field is +planned as a follow-up). + +**Client — encode each image's embeddings** (`(n_tokens, embed_dim)` tensor): + +```python +import base64, torch +from safetensors.torch import save as st_save + +def encode_embeds_data_uri(embeds: torch.Tensor) -> str: # embeds: (n_tokens, dim) + blob = st_save({"embeds": embeds.contiguous().cpu()}) + return "data:application/x-dynamo-embeds;base64," + base64.b64encode(blob).decode() + +# ... then send it as a normal image content part: +content = [ + {"type": "text", "text": "describe the product"}, + {"type": "image_url", "image_url": {"url": encode_embeds_data_uri(my_embeds)}}, +] +``` + +**Server — your projector backend** (decode + project; Dynamo splices the result): + +```python +import torch +from dynamo.vllm.multimodal_utils.vision_encoder_backend import Preprocessed +from examples.custom_encoder.qwen_vision_encoder import QwenVisionEncoderBackend +# decode_embeds_data_uri: inverse of the client snippet above (base64 -> safetensors load) + +class MyProjector(QwenVisionEncoderBackend): + image_token_id = 151655 # hardcode your model's <|image_pad|> + + def build(self, model_id): # pick your own device + self.device = "cuda" + self.proj = load_my_projector().to(self.device) # nn.Linear/MLP: dim -> hidden + + def preprocess(self, image_url): # off-thread, CUDA-free + return Preprocessed(item=decode_embeds_data_uri(image_url)) + + @torch.inference_mode() + def forward_batch(self, items, target_bucket=None): # actor thread, CPU out + sizes = [t.shape[0] for t in items] + x = torch.cat(items, dim=0).to(self.device) + y = self.proj(x) # -> (sum_tokens, hidden) + return [p.detach().cpu() for p in y.split(sizes, dim=0)] +``` + +That's it — the embeddings ride the existing `image_url` channel, your projector +maps them to the LM hidden dim, and `build_mixed_embeds` splices them at the +placeholder positions before the prompt reaches vLLM (`--enable-prompt-embeds`). +A reference stub used by the test suite lives at +`tests/utils/embeds_passthrough_encoder.py` (it passes embeddings through +unchanged — your real backend projects them). diff --git a/examples/custom_encoder/__init__.py b/examples/custom_encoder/__init__.py new file mode 100644 index 000000000000..e5725ea5a481 --- /dev/null +++ b/examples/custom_encoder/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/examples/custom_encoder/hitchhikers_vision_encoder.py b/examples/custom_encoder/hitchhikers_vision_encoder.py new file mode 100644 index 000000000000..ea7035d046c5 --- /dev/null +++ b/examples/custom_encoder/hitchhikers_vision_encoder.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example ``VisionEncoderBackend`` that fakes an image as a known text phrase. + +Instead of a real vision encoder, ``forward_batch()`` returns the LM's +``embed_tokens`` embeddings of a fixed phrase (default: *"the Ultimate Question +of Life, the Universe, and Everything"*). Splicing those embeddings in at the +image placeholder makes the assembled prompt read as one coherent sentence, so +the mixed-embeds path can be checked for **semantic** correctness, not just +shape: + + "Based on The Hitchhiker's Guide to the Galaxy, The Answer to" + + # → embeds of " the Ultimate Question of Life, ..." + + " is?" + → the model answers "42". + +The image URL is ignored — any URL yields the same phrase embeddings. This is an +**eager** backend: ``buckets`` stays ``None`` (no CUDA graphs), ``cost`` is 1 per +image, and ``preprocess`` is a pass-through (no fetch / decode), so it isolates +the splice + contract + LM path from any real vision compute. + +Subclasses ``QwenVisionEncoderBackend``: the default model is a Qwen-family LM +(Qwen2.5), so the base loads the tokenizer and resolves the ``<|image_pad|>`` +placeholder id; this class only loads the ``embed_tokens`` weight and implements +``preprocess`` + the synchronous ``forward_batch``. + +Usage (via agg_custom.sh): + DYN_ENCODER_CLASS=examples.custom_encoder.hitchhikers_vision_encoder.HitchhikersVisionEncoder + DYN_MODEL=Qwen/Qwen2.5-1.5B-Instruct + DYN_CUSTOM_PHRASE=" the Ultimate Question of Life, the Universe, and Everything" + ./agg_custom.sh +""" + +from __future__ import annotations + +import json +import logging +import os +from pathlib import Path +from typing import List, Optional + +import torch +from safetensors import safe_open +from transformers.utils import cached_file + +from dynamo.vllm.multimodal_utils.vision_encoder_backend import Preprocessed +from examples.custom_encoder.qwen_vision_encoder import QwenVisionEncoderBackend + +logger = logging.getLogger(__name__) + +# The Answer to the Ultimate Question of Life, the Universe, and Everything is 42. +_PHRASE = os.environ.get( + "DYN_CUSTOM_PHRASE", + " the Ultimate Question of Life, the Universe, and Everything", +) + + +def _load_embed_tokens_weight(model_id: str) -> torch.Tensor: + """Load only ``embed_tokens.weight`` from a HF checkpoint (lazy safetensors read). + + Works for both local directories and HF hub model IDs (resolved through the + HF cache), and for sharded and single-file checkpoints. + """ + try: + index_path = cached_file(model_id, "model.safetensors.index.json") + model_dir = Path(index_path).parent + weight_map = json.loads(Path(index_path).read_text())["weight_map"] + embed_key = next( + (k for k in weight_map if k.endswith("embed_tokens.weight")), None + ) + if embed_key is None: + raise FileNotFoundError( + f"No embed_tokens.weight key in safetensors index for {model_id}" + ) + shard_path = model_dir / weight_map[embed_key] + except (OSError, StopIteration): + # Fallback: single-file safetensors model. + shard_path = Path(cached_file(model_id, "model.safetensors")) + embed_key = None # scanned below + + with safe_open(str(shard_path), framework="pt", device="cpu") as f: + if embed_key is None: + embed_key = next( + (k for k in f.keys() if k.endswith("embed_tokens.weight")), None + ) + if embed_key is None: + raise FileNotFoundError(f"embed_tokens.weight not found in {shard_path}") + return f.get_tensor(embed_key) + + +class HitchhikersVisionEncoder(QwenVisionEncoderBackend): + """Backend that returns the LM embeddings of a fixed phrase for any image URL. + + A test/example backend, not a production vision encoder: it loads the LM's + ``embed_tokens`` weight and returns the embeddings of ``DYN_CUSTOM_PHRASE`` + so the spliced prompt reads as a coherent sentence. + """ + + # Eager: no graph ladder. Count-based budget (cost == 1 per image). + buckets = None + max_batch_cost = 8 + + def build(self, model_id: str) -> None: + """Load the tokenizer (via the Qwen base) and the LM ``embed_tokens`` weight.""" + super().build(model_id) # loads self.tokenizer + self._embed_weight = _load_embed_tokens_weight(model_id) + logger.info( + "[HitchhikersVisionEncoder] ready: embed_weight=%s dtype=%s phrase=%r", + tuple(self._embed_weight.shape), + self._embed_weight.dtype, + _PHRASE, + ) + + def preprocess(self, image_url: str) -> Preprocessed[str]: + """Pass-through: the URL is ignored, so there is nothing to fetch/decode. + + cost=1 (count-based).""" + return Preprocessed(item=image_url, cost=1) + + def forward_batch( + self, items: List[str], target_bucket: Optional[int] = None + ) -> List[torch.Tensor]: + """Return the ``embed_tokens`` embeddings of the phrase for each item. + + Synchronous batched forward — the AsyncVisionEncoder runs it on the + dedicated actor thread, one batch at a time. Eager (``target_bucket`` is + always ``None`` here since ``buckets`` is ``None``).""" + # Explicit check (not assert): asserts are stripped under `python -O`, + # which would turn a missing build() into an opaque None-index crash. + if ( + getattr(self, "_embed_weight", None) is None + or getattr(self, "tokenizer", None) is None + ): + raise RuntimeError( + "HitchhikersVisionEncoder.forward_batch() called before " + "build(); load the encoder first." + ) + ids = self.tokenizer.encode(_PHRASE, add_special_tokens=False) + phrase_embeds = self._embed_weight[torch.tensor(ids, dtype=torch.long)] + logger.debug( + "[HitchhikersVisionEncoder] phrase tokens=%d → shape=%s", + len(ids), + tuple(phrase_embeds.shape), + ) + return [phrase_embeds.clone() for _ in items] diff --git a/examples/custom_encoder/launch/agg_custom.sh b/examples/custom_encoder/launch/agg_custom.sh new file mode 100755 index 000000000000..3bb81dad2283 --- /dev/null +++ b/examples/custom_encoder/launch/agg_custom.sh @@ -0,0 +1,132 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Aggregated serving with a CustomEncoder. +# +# Architecture: single aggregated worker +# - Frontend: dynamo.frontend (OpenAI-compatible HTTP gateway) +# - Worker: dynamo.vllm with a CustomEncoder loaded in-process +# +# The CustomEncoder is called for each multimodal request: +# 1. encoder.encode(image_urls) → list[(n_visual_tokens, lm_hidden_dim)] +# 2. Dynamo builds a mixed token-ids/embeds EmbedsPrompt: vLLM embeds the text +# itself and substitutes the encoder's image embeds at the placeholder span. +# 3. vLLM engine runs transformer layers on the EmbedsPrompt. +# +# No separate encode worker, no NIXL inter-process transfer. +# +# The default model is a text-only LM (Qwen2.5-1.5B-Instruct) — the standard +# topology for this feature (custom encoder + stock LM), served with a minimal +# custom chat template that emits the image placeholder. The default encoder +# (HitchhikersVisionEncoder) fakes an image as the embeddings of a fixed phrase +# so the path can be checked for semantic correctness; replace it with a real +# CustomEncoder subclass for production. +# +# Usage: +# ./agg_custom.sh [--model ] [--encoder-class ] +# [--gpu ] +# +# Defaults: +# --model: Qwen/Qwen2.5-1.5B-Instruct +# --encoder-class: examples.custom_encoder.hitchhikers_vision_encoder.HitchhikersVisionEncoder +# --gpu: 0 + +set -e +trap 'echo "Cleaning up..."; kill 0' EXIT + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +source "$SCRIPT_DIR/../../common/gpu_utils.sh" +source "$SCRIPT_DIR/../../common/launch_utils.sh" + +# ── Defaults ────────────────────────────────────────────────────────────────── +MODEL="${DYN_MODEL:-Qwen/Qwen2.5-1.5B-Instruct}" +ENCODER_CLASS="${DYN_ENCODER_CLASS:-examples.custom_encoder.hitchhikers_vision_encoder.HitchhikersVisionEncoder}" +# Precedence: explicit DYN_WORKER_GPU/--gpu > harness-set CUDA_VISIBLE_DEVICES +# (e.g. the pytest profile runner) > device 0 (portable default; on a generic +# host or single-GPU container, GPU 0 is the natural choice). +WORKER_GPU="${DYN_WORKER_GPU:-${CUDA_VISIBLE_DEVICES:-0}}" +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +MAX_MODEL_LEN="${DYN_MAX_MODEL_LEN:-16384}" +# A text-only LM's own chat template can't render image content parts, so default +# to the bundled minimal template that emits the <|image_pad|> placeholder. +CUSTOM_JINJA_TEMPLATE="${DYN_CUSTOM_JINJA_TEMPLATE:-$SCRIPT_DIR/../templates/qwen_vl.jinja}" +EXTRA_ARGS=() + +# ── Argument parsing ────────────────────────────────────────────────────────── +while [[ $# -gt 0 ]]; do + case $1 in + --model) + MODEL=$2; shift 2 ;; + --encoder-class) + ENCODER_CLASS=$2; shift 2 ;; + --gpu) + WORKER_GPU=$2; shift 2 ;; + -h|--help) + cat <<'EOF' +Usage: agg_custom.sh [OPTIONS] + +Aggregated serving with a CustomEncoder (no separate encode worker). + +Options: + --model LLM checkpoint (default: Qwen/Qwen2.5-1.5B-Instruct) + --encoder-class Dotted module.ClassName for CustomEncoder subclass + --gpu GPU index for the worker (default: 0) + -h, --help Show this help + +Environment variables: + DYN_MODEL LLM model checkpoint + DYN_ENCODER_CLASS Dotted class path for the CustomEncoder subclass + DYN_WORKER_GPU GPU index (default: 0) + DYN_CUSTOM_JINJA_TEMPLATE Path to a .jinja chat template (defaults to the + bundled templates/qwen_vl.jinja) + DYN_CUSTOM_PHRASE Phrase the HitchhikersEncoder embeds as the "image" +EOF + exit 0 ;; + *) + EXTRA_ARGS+=("$1"); shift ;; + esac +done + +# VRAM sizing: honors the profiler/test-harness override +# (_PROFILE_OVERRIDE_VLLM_KV_CACHE_BYTES) when set, else falls back to a sane +# local default. Required for VRAM-safe parallel test scheduling. +GPU_MEM_ARGS=$(build_vllm_gpu_mem_args) +[[ -z "$GPU_MEM_ARGS" ]] && GPU_MEM_ARGS="--gpu-memory-utilization 0.8" + +print_launch_banner --no-curl "CustomEncoder — Aggregated Serving" "$MODEL" "$HTTP_PORT" \ + "Worker GPU: $WORKER_GPU" \ + "Encoder: $ENCODER_CLASS" \ + "Jinja tmpl: ${CUSTOM_JINJA_TEMPLATE:-(model default)}" \ + "NOTE: HitchhikersVisionEncoder fakes an image as a fixed-phrase embedding;" \ + " replace with a real CustomEncoder subclass for production use." + +export DYN_REQUEST_PLANE=tcp +export DYN_TCP_MAX_MESSAGE_SIZE=209715200 +export DYN_HTTP_BODY_LIMIT_MB=200 + +# ── Frontend ────────────────────────────────────────────────────────────────── +echo "[1/2] Starting frontend (port $HTTP_PORT)..." +python -m dynamo.frontend & + +# ── Aggregated worker ───────────────────────────────────────────────────────── +echo "[2/2] Starting aggregated worker (model=$MODEL, GPU=$WORKER_GPU)..." +JINJA_ARG=() +[[ -n "$CUSTOM_JINJA_TEMPLATE" ]] && JINJA_ARG=(--custom-jinja-template "$CUSTOM_JINJA_TEMPLATE") +CUDA_VISIBLE_DEVICES=$WORKER_GPU \ +DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \ +python -m dynamo.vllm \ + --model "$MODEL" \ + --custom-encoder-class "$ENCODER_CLASS" \ + --enable-multimodal \ + --enable-prompt-embeds \ + --max-model-len "$MAX_MODEL_LEN" \ + $GPU_MEM_ARGS \ + "${JINJA_ARG[@]}" \ + "${EXTRA_ARGS[@]}" & + +echo "==================================================================" +echo "All components started. Waiting for initialization (~30-60s)..." +echo "==================================================================" + +wait_any_exit diff --git a/examples/custom_encoder/qwen3vl_vit_encoder.py b/examples/custom_encoder/qwen3vl_vit_encoder.py new file mode 100644 index 000000000000..86b1fe83f6bb --- /dev/null +++ b/examples/custom_encoder/qwen3vl_vit_encoder.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Real Qwen3-VL vision-tower ``VisionEncoderBackend`` (in-process, CUDA-graphed). + +Loads the **actual** Qwen3-VL vision tower (ViT patch-embed + transformer blocks + +spatial merger) and runs it on the ``AsyncVisionEncoder`` dedicated actor thread. +It demonstrates the **bucket-ladder** graph scheme (design theme D2, the vLLM +``EncoderCudaGraphManager`` pattern) end to end: + +- ``buckets`` exposes a sorted ladder of **merged-visual-token** rungs. +- The ``ThreadedMicroBatcher`` packs images by ``cost`` (merged tokens) and rounds + the packed ``sum(cost)`` **up to the nearest rung**, passing it as + ``target_bucket``. +- ``forward_batch`` **pads** the packed batch up to ``target_bucket`` (appends + dummy images), replays the graph captured for that rung's shape, and slices the + real images' embeds back out. + +Padding the input to a rung quantises the forward's shape to the (bounded) ladder, +so ``torch.compile(mode="reduce-overhead")`` captures **one CUDA graph per rung** +(not one per arbitrary batch size — that would be SGLang's unbounded per-exact-S +scheme, which the design rejects). Capture + replay both happen on the actor +thread, the affinity the batcher guarantees. + +Stable per-image shape: every image is resized to one fixed square, so each image +is exactly ``TOKENS_PER_SIDE**2`` merged tokens — every rung is a whole number of +images and padding is always whole dummy images. ``forward_batch`` copies its +output to CPU, which also detaches it from the reused CUDA-graph output buffer. + +Limitation (intentional for this harness): Qwen3-VL's vision tower also emits +``deepstack_features`` that the LM injects at specific layers. The contract +returns one embed tensor per image (the merged ``pooler_output``); deepstack +features are not plumbed through the mixed-embeds splice path, so image grounding +is approximate. This exercises the in-process encoder + batcher + CUDA-graph +bucket-ladder path, **not** Qwen3-VL accuracy. + +Usage (via agg_custom.sh): + DYN_MODEL=Qwen/Qwen3-VL-2B-Instruct + DYN_ENCODER_CLASS=examples.custom_encoder.qwen3vl_vit_encoder.Qwen3VLViTEncoder + DYN_WORKER_GPU=2 ./agg_custom.sh + +Env knobs: + DYN_VIT_COMPILE 1 (default) → torch.compile + bucket ladder; 0 → eager (no graphs) + DYN_VIT_TOKENS_PER_SIDE merged visual tokens per image side (default 16 → 256/img) + DYN_VIT_MAX_IMAGES max images per forward → top rung (default 16) +""" + +from __future__ import annotations + +import base64 +import io +import logging +import os +from typing import Any, Dict, List, Optional + +import requests +import torch +from PIL import Image +from transformers import AutoModelForImageTextToText, AutoProcessor + +from dynamo.vllm.multimodal_utils.vision_encoder_backend import Preprocessed +from examples.custom_encoder.qwen_vision_encoder import QwenVisionEncoderBackend + +logger = logging.getLogger(__name__) + +_COMPILE = os.environ.get("DYN_VIT_COMPILE", "1") == "1" +_TOKENS_PER_SIDE = int(os.environ.get("DYN_VIT_TOKENS_PER_SIDE", "16")) +_MAX_IMAGES = int(os.environ.get("DYN_VIT_MAX_IMAGES", "16")) + + +def _load_image(image_url: str) -> Image.Image: + """Load a PIL RGB image from a data: URI, http(s) URL, or local path.""" + if image_url.startswith("data:"): + b64 = image_url.split(",", 1)[1] + return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB") + if image_url.startswith(("http://", "https://")): + resp = requests.get(image_url, timeout=30) + resp.raise_for_status() + return Image.open(io.BytesIO(resp.content)).convert("RGB") + return Image.open(image_url).convert("RGB") + + +def _power_of_two_ladder(tokens_per_img: int, max_images: int) -> List[int]: + """Merged-token rungs at 1, 2, 4, ... images up to ``max_images``. + + Power-of-two image counts (not every count) so the ladder — and thus the + captured-graph count — stays bounded; sub-rung batches pad up to the next.""" + mults: List[int] = [] + m = 1 + while m < max_images: + mults.append(m) + m *= 2 + mults.append(max_images) + return sorted({tokens_per_img * x for x in mults}) + + +class Qwen3VLViTEncoder(QwenVisionEncoderBackend): + """In-process Qwen3-VL vision-tower backend with a CUDA-graphed bucket ladder.""" + + def __init__(self) -> None: + self._tokens_per_img = _TOKENS_PER_SIDE**2 + if _COMPILE: + # Graphed: expose the ladder; the top rung is the dispatch ceiling. + self.buckets = _power_of_two_ladder(self._tokens_per_img, _MAX_IMAGES) + self.max_batch_cost = self.buckets[-1] + else: + # Eager: no ladder; pack up to max_images worth of tokens, no padding. + self.buckets = None + self.max_batch_cost = self._tokens_per_img * _MAX_IMAGES + + def build(self, model_id: str) -> None: + """Load tokenizer (Qwen base) + processor + ViT; compile and warm up so one + CUDA graph per rung is captured on this (the actor) thread.""" + super().build(model_id) # self.tokenizer + # The worker pins the GPU via CUDA_VISIBLE_DEVICES, so the current device + # ("cuda") is correct — the backend picks its own device (no device arg). + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.processor = AutoProcessor.from_pretrained(model_id) + + model = AutoModelForImageTextToText.from_pretrained( + model_id, dtype=torch.bfloat16 + ) + # Qwen3VLForConditionalGeneration.model is Qwen3VLModel(.visual, .language_model). + inner = getattr(model, "model", model) + visual = getattr(inner, "visual", None) or getattr(model, "visual") + self.visual = visual.to(self.device).eval() + del model # drop the LM half; the vLLM worker owns the LM + torch.cuda.empty_cache() + + vc = self.visual.config + self.merge = int(vc.spatial_merge_size) + patch = int(vc.patch_size) + # Fixed square so grid_thw (→ cost) is constant: every image is + # exactly _tokens_per_img merged tokens. + self.side = patch * self.merge * _TOKENS_PER_SIDE + self._fixed_hw = (self.side, self.side) + + # One dummy (gray) image's processed tensors, reused for padding to a rung. + self._dummy = self._process(Image.new("RGB", self._fixed_hw, (127, 127, 127))) + + self._eager_visual = self.visual + self._compiled = False + if _COMPILE: + self.visual = torch.compile(self.visual, mode="reduce-overhead") + self._compiled = True + + self._warmup() + + def _warmup(self) -> None: + """Forward once at **each rung** so torch.compile captures that rung's CUDA + graph here; fall back to eager if compile/capture fails.""" + try: + rungs = self.buckets if self.buckets else [self._tokens_per_img] + for _ in range(2 if self._compiled else 1): + for rung in rungs: + # One real image padded up to the rung (the forward_batch path). + self.forward_batch( + [self._dummy], target_bucket=rung if self._compiled else None + ) + torch.cuda.synchronize() + logger.info( + "[Qwen3VLViTEncoder] ready: side=%d merge=%d compile=%s " + "tokens/img=%d buckets=%s max_batch_cost=%d", + self.side, + self.merge, + self._compiled, + self._tokens_per_img, + list(self.buckets) if self.buckets else None, + self.max_batch_cost, + ) + except Exception as exc: # noqa: BLE001 — compile/capture is best-effort + if self._compiled: + logger.warning( + "[Qwen3VLViTEncoder] compile/warmup failed (%s); using eager", exc + ) + self.visual = self._eager_visual + self._compiled = False + self.buckets = None + self.forward_batch([self._dummy]) + torch.cuda.synchronize() + else: + raise + + # ---- preprocess (off the actor thread) --------------------------------- + + def preprocess(self, image_url: str) -> Preprocessed[Dict[str, Any]]: + """Off-thread: fetch + resize to the fixed square + HF patchify. + + ``cost`` = merged visual tokens for this image (the scalar Dynamo packs + by). Every image is the same fixed square, so the ViT shapes are uniform — + the author owns any shape/padding concerns inside ``forward_batch``.""" + img = _load_image(image_url).resize(self._fixed_hw) + item = self._process(img) + t, h, w = item["grid_thw"][0].tolist() + cost = (t * h * w) // (self.merge**2) + return Preprocessed(item=item, cost=cost) + + def _process(self, img: Image.Image) -> Dict[str, Any]: + out = self.processor.image_processor(images=[img], return_tensors="pt") + return {"pixel_values": out["pixel_values"], "grid_thw": out["image_grid_thw"]} + + # ---- forward (on the actor thread) ------------------------------------- + + @torch.inference_mode() + def forward_batch( + self, items: List[Dict[str, Any]], target_bucket: Optional[int] = None + ) -> List[torch.Tensor]: + """Pad the packed batch up to ``target_bucket``, replay that rung's graph, + and slice the real images' embeds back out (one tensor per item, in order).""" + if os.environ.get("DYN_VIT_LOG_BATCH") == "1": + logger.info( + "[Qwen3VLViTEncoder] forward_batch images=%d target_bucket=%s", + len(items), + target_bucket, + ) + n_real = len(items) + pix_parts = [it["pixel_values"] for it in items] + grid_parts = [it["grid_thw"] for it in items] + + # Pad up to the rung (whole dummy images) so the forward shape is one of the + # captured rungs. target_bucket is None in eager mode → no padding. + if target_bucket is not None: + real_tokens = sum( + (t * h * w) // (self.merge**2) + for t, h, w in torch.cat(grid_parts, dim=0).tolist() + ) + n_dummy = (target_bucket - real_tokens) // self._tokens_per_img + for _ in range(max(0, n_dummy)): + pix_parts.append(self._dummy["pixel_values"]) + grid_parts.append(self._dummy["grid_thw"]) + + pix = torch.cat(pix_parts, dim=0).to(self.device, dtype=torch.bfloat16) + grid = torch.cat(grid_parts, dim=0).to(self.device) + embeds = self.visual(pix, grid).pooler_output # (total_merged_tokens, hidden) + sizes = [(t * h * w) // (self.merge**2) for t, h, w in grid.tolist()] + parts = torch.split(embeds, sizes, dim=0) + # Copy to CPU: detaches from any reused CUDA-graph output buffer and matches + # the assembler's CPU prompt_embeds layout (it preserves dtype). Drop the + # padding dummies — return only the real images, in input order. + return [parts[i].detach().to("cpu", copy=True) for i in range(n_real)] + + def close(self) -> None: + """Drop the ViT + compiled graphs on the actor thread.""" + self.visual = None + self._eager_visual = None + torch.cuda.empty_cache() diff --git a/examples/custom_encoder/qwen_vision_encoder.py b/examples/custom_encoder/qwen_vision_encoder.py new file mode 100644 index 000000000000..68df737c4ab8 --- /dev/null +++ b/examples/custom_encoder/qwen_vision_encoder.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reusable example base for Qwen-family ``VisionEncoderBackend`` authors. + +Hardcodes the Qwen ``<|image_pad|>`` placeholder id and loads the model tokenizer +(handy for subclasses that tokenize text). A concrete Qwen-family encoder +subclasses this and implements only ``preprocess`` + ``forward_batch``. + + class MyQwenEncoder(QwenVisionEncoderBackend): + def build(self, model_id): + super().build(model_id) # loads self.tokenizer + # ... load ViT + projector (pick the device yourself) ... + def preprocess(self, raw): + ... # off-thread, returns Preprocessed + def forward_batch(self, items, target_bucket=None): + ... # actor thread, batched forward (CPU out) +""" + +from __future__ import annotations + +from transformers import AutoTokenizer + +from dynamo.vllm.multimodal_utils.vision_encoder_backend import VisionEncoderBackend + + +class QwenVisionEncoderBackend(VisionEncoderBackend): + """``VisionEncoderBackend`` base for Qwen-family models (Qwen2-VL / Qwen3-VL / + Qwen3.5). + + Hardcodes ``image_token_id`` to Qwen3-VL's ``<|image_pad|>`` (151655) — override + it for other versions (e.g. 248056 for Qwen3.5). ``build`` loads the model + tokenizer; ``preprocess`` and ``forward_batch`` stay abstract, so this class + cannot be instantiated directly — subclass it and implement them. + """ + + # Qwen3-VL <|image_pad|>; override for other Qwen versions. + image_token_id = 151655 + + def build(self, model_id: str) -> None: + """Load the model tokenizer. Subclasses extend this (call super) to also + load their encoder weights (picking the device themselves).""" + self.tokenizer = AutoTokenizer.from_pretrained(model_id) diff --git a/examples/custom_encoder/templates/qwen_vl.jinja b/examples/custom_encoder/templates/qwen_vl.jinja new file mode 100644 index 000000000000..927638ec2dfc --- /dev/null +++ b/examples/custom_encoder/templates/qwen_vl.jinja @@ -0,0 +1,26 @@ +{#- Minimal chat template for a text-only LM that hosts a custom image encoder. -#} +{#- Renders text content normally and emits one bare <|image_pad|> placeholder -#} +{#- per image content part. The CustomEncoder replaces each placeholder token -#} +{#- with the encoder's image embeddings before the prompt reaches the engine. -#} +{#- The frontend rewrites each {"type":"image_url",...} part to a bare -#} +{#- {"type":"image"} before this template runs, so match on type == 'image'. -#} +{#- The assembler maps one placeholder token to one image tensor, so -#} +{#- consecutive images need no separator. -#} +{%- for message in messages %} + {{- '<|im_start|>' + message.role + '\n' }} + {%- if message.content is string %} + {{- message.content }} + {%- else %} + {%- for content in message.content %} + {%- if content.type == 'image' %} + {{- '<|image_pad|>' }} + {%- elif content.type == 'text' %} + {{- content.text }} + {%- endif %} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} +{%- endif %} diff --git a/tests/report_pytest_markers.py b/tests/report_pytest_markers.py index 4d47cde47af7..ef14f9314cc1 100755 --- a/tests/report_pytest_markers.py +++ b/tests/report_pytest_markers.py @@ -114,6 +114,10 @@ "transformers.models", "transformers.models.qwen2_vl", "transformers.models.qwen2_vl.image_processing_qwen2_vl", + # safetensors — used by the client-embeddings stub backend + # (tests/utils/embeds_passthrough_encoder.py); not in the pre-commit env. + "safetensors", + "safetensors.torch", "pandas", "matplotlib", "matplotlib.pyplot", diff --git a/tests/serve/multimodal_profiles/vllm.py b/tests/serve/multimodal_profiles/vllm.py index 3765f390c5cc..e914082f406d 100644 --- a/tests/serve/multimodal_profiles/vllm.py +++ b/tests/serve/multimodal_profiles/vllm.py @@ -1,13 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os + import pytest +from dynamo.common.utils.paths import WORKSPACE_DIR from tests.utils.multimodal import ( MmCase, MultimodalModelProfile, TopologyConfig, make_audio_payload, + make_custom_encoder_payload, make_image_payload, make_image_payload_b64, make_image_payload_cached_tokens, @@ -53,6 +57,10 @@ "epd": "disagg_multimodal_epd.sh", "epd_video": "disagg_multimodal_epd.sh", "p_d": "disagg_multimodal_p_d.sh", + # CustomEncoder: a custom in-process vision encoder on a text-only LM + # (no separate encode worker, no NIXL). Lives in examples/custom_encoder, + # not examples/backends/vllm — the TopologyConfig sets `directory` to match. + "agg_custom": "agg_custom.sh", } VLLM_MULTIMODAL_PROFILES: list[MultimodalModelProfile] = [ @@ -505,4 +513,38 @@ ), }, ), + # CustomEncoder coverage. NOTE: Qwen2.5-1.5B-Instruct is a TEXT-ONLY LM — + # it sits in the multimodal profiles because the in-process CustomEncoder + # *plugin* (a custom vision encoder) gives it the image->embeds serving + # path; the multimodality comes from the encoder, not the model. The + # `agg_custom` topology launches examples/custom_encoder/launch/agg_custom.sh + # (hence the `directory` override) with the example HitchhikersVisionEncoder, + # which fakes an image as a fixed phrase so the spliced prompt answers "42". + MultimodalModelProfile( + name="Qwen/Qwen2.5-1.5B-Instruct", + short_name="custom-encoder", + topologies={ + "agg_custom": TopologyConfig( + marks=[pytest.mark.post_merge], + timeout_s=300, + directory=os.path.join(WORKSPACE_DIR, "examples/custom_encoder"), + env={ + # The single-GPU test container remaps its host GPU to + # device 0, so pin the worker there (agg_custom.sh also + # honors an externally-set CUDA_VISIBLE_DEVICES). + "DYN_WORKER_GPU": "0", + "DYN_ENCODER_CLASS": ( + "examples.custom_encoder.hitchhikers_vision_encoder." + "HitchhikersVisionEncoder" + ), + "DYN_CUSTOM_JINJA_TEMPLATE": os.path.join( + WORKSPACE_DIR, + "examples/custom_encoder/templates/qwen_vl.jinja", + ), + "PYTHONPATH": str(WORKSPACE_DIR), + }, + tests=[MmCase(payload=make_custom_encoder_payload())], + ), + }, + ), ] diff --git a/tests/utils/embeds_passthrough_encoder.py b/tests/utils/embeds_passthrough_encoder.py new file mode 100644 index 000000000000..1fbef287614e --- /dev/null +++ b/tests/utils/embeds_passthrough_encoder.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Test-only stub backend for the client-supplied-embeddings path. + +A client that already has per-image embeddings (e.g. CLIP/ViT outputs) can drive +the existing in-process custom-encoder path **with no new Dynamo code**: it sends +each embedding inline as a safetensors ``data:`` URI on a normal ``image_url`` +content part, and provides a ``VisionEncoderBackend`` whose ``preprocess`` decodes +the embedding and whose ``forward_batch`` **projects** it to the LM hidden dim. +Dynamo splices the result in via ``build_mixed_embeds`` (unchanged). + +This module is **not a shipped example** — the projector is the client's code. It +exists only to prove + test that path end to end: ``EmbedsPassthroughEncoder`` +decodes the client's embedding and passes it through unchanged (identity — it +assumes the client already sent hidden-dim embeds). See +``examples/custom_encoder/README.md`` for the client-facing contract. +""" + +from __future__ import annotations + +import base64 +from typing import List, Optional + +import torch +from safetensors.torch import load as _safetensors_load +from safetensors.torch import save as _safetensors_save + +from dynamo.vllm.multimodal_utils.vision_encoder_backend import Preprocessed +from examples.custom_encoder.qwen_vision_encoder import QwenVisionEncoderBackend + +# Media type marking a Dynamo per-image embeddings payload (safetensors bytes, +# base64-encoded). This is the client↔server wire convention for PR6; PR7 will add +# a first-class ``image_embeds`` request field that carries the same bytes. +EMBEDS_DATA_URI_PREFIX = "data:application/x-dynamo-embeds;base64," +_EMBEDS_KEY = "embeds" + + +def encode_embeds_data_uri(embeds: torch.Tensor) -> str: + """Serialize a 2D ``(n_tokens, hidden)`` tensor as a safetensors base64 ``data:`` URI. + + The client-side contract: compute your per-image embeddings, encode them with + this, and send the returned string as the ``image_url`` of an image content + part. safetensors carries dtype + shape and never executes code on load. + """ + if embeds.dim() != 2: + raise ValueError(f"embeds must be 2D (n_tokens, hidden), got {embeds.dim()}D") + blob = _safetensors_save({_EMBEDS_KEY: embeds.contiguous().cpu()}) + return EMBEDS_DATA_URI_PREFIX + base64.b64encode(blob).decode("ascii") + + +def decode_embeds_data_uri(url: str) -> torch.Tensor: + """Inverse of :func:`encode_embeds_data_uri`: ``data:`` URI → CPU ``(n_tokens, hidden)``. + + Raises ``ValueError`` on a non-``data:`` URL (e.g. an http URL) or a malformed + payload, so a bad input fails only that image (the A5 barrier turns it into a + clean request-level error). + """ + if not isinstance(url, str) or not url.startswith("data:"): + raise ValueError( + f"expected an embeds data: URI starting with 'data:', got {url[:48]!r}" + ) + _, _, b64 = url.partition(",") + if not b64: + raise ValueError("malformed embeds data: URI (no base64 payload after ',')") + try: + blob = base64.b64decode(b64, validate=True) + tensors = _safetensors_load(blob) + except Exception as exc: # noqa: BLE001 — bad input fails just this image + raise ValueError(f"failed to decode embeds data: URI: {exc}") from exc + if _EMBEDS_KEY not in tensors: + raise ValueError(f"embeds payload missing '{_EMBEDS_KEY}' tensor key") + return tensors[_EMBEDS_KEY] + + +class EmbedsPassthroughEncoder(QwenVisionEncoderBackend): + """Decode client embeddings and pass them through unchanged (identity). + + Test-only stub. A real client backend would **project** (``embed_dim`` → + ``lm_hidden_dim``) inside ``forward_batch``; this stub assumes the client + already sent hidden-dim embeds, so the LM reads them directly through the + mixed-embeds splice. Eager + pass-through (no batch cap, no graph ladder). + """ + + buckets = None + max_batch_cost = None # pass-through: the whole drained batch in one forward + + def preprocess(self, image_url: str) -> Preprocessed: + """Off-thread: decode the client's safetensors ``data:`` URI to a CPU tensor.""" + embeds = decode_embeds_data_uri(image_url) + if embeds.dim() != 2: + raise ValueError( + f"decoded embeds must be 2D (n_tokens, hidden), got {embeds.dim()}D" + ) + return Preprocessed(item=embeds) + + def forward_batch( + self, items: List[torch.Tensor], target_bucket: Optional[int] = None + ) -> List[torch.Tensor]: + """Identity: return one tensor per item, unchanged (a real client projects here).""" + return list(items) diff --git a/tests/utils/multimodal.py b/tests/utils/multimodal.py index 7147739fae45..802ee676d39d 100644 --- a/tests/utils/multimodal.py +++ b/tests/utils/multimodal.py @@ -231,6 +231,32 @@ def make_audio_payload(expected_response: list[str]) -> ChatPayload: ) +def make_custom_encoder_payload() -> ChatPayload: + """Semantic check for the aggregated CustomEncoder path. + + The example HitchhikersVisionEncoder splices the embeddings of "the + Ultimate Question of Life, the Universe, and Everything" at the image + placeholder, so the assembled prompt must answer "42". The served image + content is irrelevant (the encoder ignores the URL). ``expected_log`` + asserts the encoder loaded in-process at worker startup. + """ + return chat_payload( + [ + { + "type": "text", + "text": "Based on The Hitchhiker's Guide to the Galaxy, The Answer to", + }, + {"type": "image_url", "image_url": {"url": MULTIMODAL_IMG_URL}}, + {"type": "text", "text": " is?"}, + ], + repeat_count=1, + expected_response=["42"], + expected_log=["Loaded CustomEncoder"], + max_tokens=32, + temperature=0.0, + ) + + # --------------------------------------------------------------------------- # Config dataclasses # ---------------------------------------------------------------------------