From fb94b87d4b0ec77f1aa04d57da3d32cc9dbc1ef4 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Thu, 10 Sep 2026 00:09:09 +0800 Subject: [PATCH] refactor(inputs): build the Omni renderer as a subclass instead of a proxy OmniRenderer wrapped the upstream renderer, forwarded unknown attributes via __getattr__, and re-bound eight BaseRenderer methods (render_cmpl, render_chat, process_for_engine, _process_enc_dec and their async twins) onto the wrapper so that they would reach Omni's _process_singleton. That allow-list hid any concrete-renderer override of those methods and would silently skip the Omni extras for any future BaseRenderer entry point not on the list. Make OmniRenderer a mixin and create the renderer as a real subclass of the upstream class instead: - omni_renderer_cls(base_cls) builds (and caches) Omni = (OmniRenderer, base_cls); only _process_singleton / _process_singleton_async are extended, everything else is the upstream renderer unchanged. - build_omni_renderer(vllm_config) resolves the class the way upstream renderer_from_config does (cached_tokenizer_from_config, tokenizer_args_from_config, RENDERER_REGISTRY.load_renderer_cls) and instantiates the Omni subclass; renderer_cls/tokenizer can be passed to skip the resolution. - build_stage0_input_processor passes the renderer to InputProcessor up front instead of replacing it afterwards; the skip-tokenizer branch keeps _build_token_only_renderer and never resolves a tokenizer. Rebinding an existing instance's __class__ was rejected: CPython refuses it for HfRenderer ("object layout differs"), so the subclass is constructed. Tests: tests/inputs/test_preprocess.py builds the renderer with omni_renderer_cls and adds checks that the subclass is created once, that a concrete renderer's render_cmpl override is no longer suppressed, and that build_omni_renderer resolves like upstream or honours an explicit class; tests/engine/test_async_omni_engine_stage_init.py asserts the stage-0 processor receives the built subclass instance and that the skip-tokenizer branch never calls build_omni_renderer. Two pre-existing mypy complaints in test_preprocess.py are fixed with annotations. Validation: ruff, pre-commit (incl. mypy) on the four files; the modules import and omni_renderer_cls(HfRenderer) builds; tests/inputs (32), tests/engine/test_async_omni_engine_stage_init.py (55) and tests/entrypoints (548) pass on CPU. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MpXe2LeRrUQGjZfKFaocMZ Signed-off-by: tzhouam --- .../test_async_omni_engine_stage_init.py | 36 ++++-- tests/inputs/test_preprocess.py | 72 ++++++++++-- vllm_omni/engine/stage_init_utils.py | 22 ++-- vllm_omni/inputs/preprocess.py | 107 +++++++++++++----- 4 files changed, 184 insertions(+), 53 deletions(-) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 79317805f..27677bd54 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1593,22 +1593,34 @@ def test_model_path_resolver_is_generic_and_model_owned(tmp_path): assert "model_path_resolver" not in engine_args -def test_build_stage0_input_processor_wraps_active_renderer(monkeypatch): - import vllm_omni.engine.stage_init_utils as init_mod - from vllm_omni.inputs.preprocess import OmniRenderer +def test_build_stage0_input_processor_uses_omni_renderer_subclass(monkeypatch): + from vllm.renderers import BaseRenderer - original = object() + import vllm_omni.engine.stage_init_utils as init_mod + from vllm_omni.inputs.preprocess import OmniRenderer, omni_renderer_cls + + class _Base(BaseRenderer): + def __init__(self, config, tokenizer): + self.config, self.tokenizer = config, tokenizer + + def render_messages(self, messages, params): # pragma: no cover - abstract stub + raise NotImplementedError + + config = types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {})) + built = omni_renderer_cls(_Base)(config, "tok") + seen = {} class DummyInputProcessor: def __init__(self, vllm_config, renderer=None): - self.renderer = renderer or original + seen["renderer"] = renderer + self.renderer = renderer monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) - processor = build_stage0_input_processor( - types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {})) - ) + monkeypatch.setattr(init_mod, "build_omni_renderer", lambda cfg: built if cfg is config else None) + processor = build_stage0_input_processor(config) + assert seen["renderer"] is built assert isinstance(processor.renderer, OmniRenderer) - assert processor.renderer._renderer is original + assert isinstance(processor.renderer, _Base) assert not hasattr(processor, "input_preprocessor") @@ -1622,13 +1634,17 @@ def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped(mo assert renderer is original self.renderer = renderer + def _must_not_resolve(_cfg): + raise AssertionError("tokenizer must not be resolved when skip_tokenizer_init=True") + monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) monkeypatch.setattr(init_mod, "_build_token_only_renderer", lambda _: original) + monkeypatch.setattr(init_mod, "build_omni_renderer", _must_not_resolve) config = types.SimpleNamespace( model_config=types.SimpleNamespace(skip_tokenizer_init=True, try_get_generation_config=lambda: {}) ) processor = build_stage0_input_processor(config) - assert processor.renderer._renderer is original + assert processor.renderer is original def test_inject_kv_stage_info_infers_sender_tp_topology(): diff --git a/tests/inputs/test_preprocess.py b/tests/inputs/test_preprocess.py index 1b5e057fb..3fb48a7cc 100644 --- a/tests/inputs/test_preprocess.py +++ b/tests/inputs/test_preprocess.py @@ -14,7 +14,8 @@ from vllm.renderers import BaseRenderer from vllm.renderers.params import TokenizeParams from vllm.v1.engine.input_processor import InputProcessor -from vllm_omni.inputs.preprocess import OmniRenderer +from vllm_omni.inputs import preprocess as preprocess_mod +from vllm_omni.inputs.preprocess import OmniRenderer, build_omni_renderer, omni_renderer_cls pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -35,7 +36,7 @@ class _Renderer(BaseRenderer): @pytest.fixture def renderer(): - return OmniRenderer(_Renderer()) + return omni_renderer_cls(_Renderer)() @pytest.mark.parametrize("kwargs", [{}, {"target_h": 512, "target_w": 768}]) @@ -55,12 +56,12 @@ def test_process_inputs_routes_no_media_processor_kwargs(renderer, kwargs, token processor._validate_params = Mock() processor._validate_lora = Mock() processor._validate_model_inputs = Mock() - prompt = {"prompt_token_ids": [1, 2, 3]} if tokenized else {"prompt": "hello"} + prompt: dict[str, object] = {"prompt_token_ids": [1, 2, 3]} if tokenized else {"prompt": "hello"} prompt.update(mm_processor_kwargs=kwargs, cache_salt="salt") request = processor.process_inputs("image-request", prompt, PoolingParams(), ("embed",)) assert request.prompt_token_ids == [1, 2, 3, 99] assert request.cache_salt == "salt" - renderer._renderer._process_multimodal.assert_called_once_with( + renderer._process_multimodal.assert_called_once_with( [1, 2, 3], {}, mm_processor_kwargs=kwargs, @@ -72,7 +73,11 @@ def test_process_inputs_routes_no_media_processor_kwargs(renderer, kwargs, token @pytest.mark.parametrize("async_mode", [False, True]) @pytest.mark.parametrize("kind", ["text", "tokens", "media", "kwargs", "embeds"]) def test_renderer_preserves_routes_extras_and_cache_policy(renderer, async_mode, kind): - prompt = {"prompt": "hello", "additional_information": {"speaker": 1}, "model_intermediate_buffer": {"x": 2}} + prompt: dict[str, object] = { + "prompt": "hello", + "additional_information": {"speaker": 1}, + "model_intermediate_buffer": {"x": 2}, + } if kind == "tokens": prompt["prompt_token_ids"] = [1, 2, 3] elif kind == "media": @@ -90,10 +95,63 @@ def test_renderer_preserves_routes_extras_and_cache_policy(renderer, async_mode, assert result["prompt"] == "hello" if kind in ("media", "kwargs"): assert result["prompt_token_ids"] == [1, 2, 3, 99] - assert renderer._renderer._process_multimodal.call_args.kwargs["skip_mm_cache"] is True + assert renderer._process_multimodal.call_args.kwargs["skip_mm_cache"] is True else: - renderer._renderer._process_multimodal.assert_not_called() + renderer._process_multimodal.assert_not_called() if kind == "embeds": torch.testing.assert_close(result["prompt_embeds"], prompt["prompt_embeds"]) else: assert result["prompt_token_ids"] == [1, 2, 3] + + +def test_omni_renderer_is_a_real_subclass_created_once(): + cls = omni_renderer_cls(_Renderer) + assert issubclass(cls, OmniRenderer) and issubclass(cls, _Renderer) + assert omni_renderer_cls(_Renderer) is cls + assert omni_renderer_cls(cls) is cls + assert cls.__name__ == "Omni_Renderer" + instance = cls() + assert isinstance(instance, _Renderer) + assert not hasattr(instance, "_renderer") + + +def test_concrete_renderer_overrides_are_not_suppressed(): + """The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides.""" + + class _Custom(_Renderer): + def render_cmpl(self, prompts, *args, **kwargs): + return ["custom"] + + assert omni_renderer_cls(_Custom)().render_cmpl([{"prompt": "x"}]) == ["custom"] + + +class _CtorRenderer(BaseRenderer): + def __init__(self, config, tokenizer): + self.config = config + self.tokenizer = tokenizer + + def render_messages(self, messages, params): + raise NotImplementedError + + +def test_build_omni_renderer_with_explicit_class_skips_registry(monkeypatch): + monkeypatch.setattr(preprocess_mod, "cached_tokenizer_from_config", Mock(side_effect=AssertionError("resolved"))) + config = SimpleNamespace(model_config=SimpleNamespace()) + built = build_omni_renderer(config, renderer_cls=_CtorRenderer, tokenizer=None) + assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer) + assert built.config is config and built.tokenizer is None + + +def test_build_omni_renderer_resolves_like_upstream(monkeypatch): + config = SimpleNamespace(model_config=SimpleNamespace()) + tokenizer = object() + monkeypatch.setattr( + preprocess_mod, "cached_tokenizer_from_config", lambda mc: tokenizer if mc is config.model_config else None + ) + monkeypatch.setattr(preprocess_mod, "tokenizer_args_from_config", lambda mc: ("fake-mode", None)) + load_cls = Mock(return_value=_CtorRenderer) + monkeypatch.setattr(preprocess_mod.RENDERER_REGISTRY, "load_renderer_cls", load_cls) + built = build_omni_renderer(config) + load_cls.assert_called_once_with("fake-mode") + assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer) + assert built.tokenizer is tokenizer diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 9bdcd647f..de97f7aa3 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -54,7 +54,7 @@ from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams -from vllm_omni.inputs.preprocess import OmniRenderer +from vllm_omni.inputs.preprocess import build_omni_renderer, omni_renderer_cls from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig @@ -1549,22 +1549,24 @@ class _TokenOnlyRenderer(BaseRenderer): def _build_token_only_renderer(stage_vllm_config: Any) -> BaseRenderer: - return _TokenOnlyRenderer(stage_vllm_config, tokenizer=None) + return omni_renderer_cls(_TokenOnlyRenderer)(stage_vllm_config, tokenizer=None) def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor: - """Build the shared stage-0 input processor.""" + """Build the shared stage-0 input processor. + + The renderer is the Omni subclass of the upstream renderer class that + ``renderer_from_config`` would have picked (or of the token-only renderer + when the stage skips tokenizer initialization), so upstream's + ``InputProcessor`` runs unmodified on top of it. + """ patch_generation_config_if_needed(stage_vllm_config.model_config) if bool(getattr(stage_vllm_config.model_config, "skip_tokenizer_init", False)): - input_processor = InputProcessor( - vllm_config=stage_vllm_config, - renderer=_build_token_only_renderer(stage_vllm_config), - ) + renderer = _build_token_only_renderer(stage_vllm_config) else: - input_processor = InputProcessor(vllm_config=stage_vllm_config) - input_processor.renderer = OmniRenderer(input_processor.renderer) - return input_processor + renderer = build_omni_renderer(stage_vllm_config) + return InputProcessor(vllm_config=stage_vllm_config, renderer=renderer) def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str: diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 28fe29236..dc4105b4d 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -1,34 +1,46 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni renderer: the upstream renderer class with Omni engine-input routing. -from typing import Any +``OmniRenderer`` is a mixin. ``omni_renderer_cls`` puts it in front of a +concrete ``BaseRenderer`` subclass, so the Omni instance *is* an upstream +renderer (tokenizer, executors, multimodal cache and every public entry point +come from upstream unchanged) and only ``_process_singleton`` / +``_process_singleton_async`` are extended. ``build_omni_renderer`` resolves +the concrete class the same way upstream ``renderer_from_config`` does. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast from vllm.renderers import BaseRenderer +from vllm.renderers.registry import RENDERER_REGISTRY +from vllm.tokenizers.registry import cached_tokenizer_from_config, tokenizer_args_from_config + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.tokenizers import TokenizerLike class OmniRenderer: - """Delegate tokenization/chat rendering while extending engine-input routing. + """Mixin that routes Omni prompts through the upstream renderer. - Keep the original renderer's tokenizer, executors, and multimodal caches. - The upstream pipeline methods run on this adapter so both sync and async - entrypoints reach the Omni singleton processing below. + It has no base class of its own; ``omni_renderer_cls`` prepends it to a + concrete ``BaseRenderer`` subclass, so ``self`` and ``super()`` are that + renderer at runtime (hence the casts below). + + Two extensions over upstream ``_process_singleton``: + + * a token prompt that carries ``mm_processor_kwargs`` but no + ``multi_modal_data`` still goes through the multimodal processor (AR + image-generation models such as GLM-Image read their target size from + those kwargs); + * Omni-only prompt keys (``prompt``, ``cache_salt``, + ``additional_information``, ``model_intermediate_buffer``) are copied + onto the engine input. """ - render_cmpl = BaseRenderer.render_cmpl - render_cmpl_async = BaseRenderer.render_cmpl_async - render_chat = BaseRenderer.render_chat - render_chat_async = BaseRenderer.render_chat_async - process_for_engine = BaseRenderer.process_for_engine - process_for_engine_async = BaseRenderer.process_for_engine_async - _process_enc_dec = BaseRenderer._process_enc_dec - _process_enc_dec_async = BaseRenderer._process_enc_dec_async - - def __init__(self, renderer: BaseRenderer) -> None: - self._renderer = renderer - - def __getattr__(self, name: str) -> Any: - return getattr(self._renderer, name) - @staticmethod def _with_omni_extras(inputs, prompt): for key in ("prompt", "cache_salt", "additional_information", "model_intermediate_buffer"): @@ -36,9 +48,14 @@ class OmniRenderer: inputs[key] = prompt[key] return inputs + @staticmethod + def _routes_no_media_kwargs(prompt) -> bool: + return "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data") + def _process_singleton(self, prompt, *, skip_mm_cache: bool = False): - if "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data"): - inputs = self._renderer._process_multimodal( + renderer = cast(BaseRenderer, self) + if self._routes_no_media_kwargs(prompt): + inputs = renderer._process_multimodal( prompt["prompt_token_ids"], {}, mm_processor_kwargs=prompt["mm_processor_kwargs"], @@ -46,12 +63,13 @@ class OmniRenderer: skip_mm_cache=skip_mm_cache, ) else: - inputs = self._renderer._process_singleton(prompt, skip_mm_cache=skip_mm_cache) + inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache) return self._with_omni_extras(inputs, prompt) async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False): - if "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data"): - inputs = await self._renderer._process_multimodal_async( + renderer = cast(BaseRenderer, self) + if self._routes_no_media_kwargs(prompt): + inputs = await renderer._process_multimodal_async( prompt["prompt_token_ids"], {}, mm_processor_kwargs=prompt["mm_processor_kwargs"], @@ -59,5 +77,42 @@ class OmniRenderer: skip_mm_cache=skip_mm_cache, ) else: - inputs = await self._renderer._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache) + inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache) return self._with_omni_extras(inputs, prompt) + + +_OMNI_RENDERER_CLASSES: dict[type[BaseRenderer], type[BaseRenderer]] = {} + + +def omni_renderer_cls(base_cls: type[BaseRenderer]) -> type[BaseRenderer]: + """Return the Omni subclass of ``base_cls`` (created once per base class).""" + if issubclass(base_cls, OmniRenderer): + return base_cls + cls = _OMNI_RENDERER_CLASSES.get(base_cls) + if cls is None: + name = f"Omni{base_cls.__name__}" + cls = type(name, (OmniRenderer, base_cls), {"__module__": __name__, "__qualname__": name}) + _OMNI_RENDERER_CLASSES[base_cls] = cls + return cls + + +def build_omni_renderer( + vllm_config: VllmConfig, + *, + renderer_cls: type[BaseRenderer] | None = None, + tokenizer: TokenizerLike | None = None, +) -> BaseRenderer: + """Build the Omni renderer for ``vllm_config``. + + With ``renderer_cls`` unset this mirrors upstream ``renderer_from_config``: + the tokenizer and renderer mode come from the model config and the class + from ``RENDERER_REGISTRY``. Pass ``renderer_cls`` (and optionally + ``tokenizer``) to skip that resolution, e.g. for tokenizer-less stages. + """ + model_config = vllm_config.model_config + if renderer_cls is None: + if tokenizer is None: + tokenizer = cached_tokenizer_from_config(model_config) + renderer_mode, *_ = tokenizer_args_from_config(model_config) + renderer_cls = RENDERER_REGISTRY.load_renderer_cls(renderer_mode) + return omni_renderer_cls(renderer_cls)(vllm_config, tokenizer)