diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 142affab7..bfa7be0e4 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -20,7 +20,11 @@ import pytest from omegaconf import OmegaConf from pydantic import ValidationError from transformers import PretrainedConfig, Qwen3OmniMoeConfig +from vllm.config import VllmConfig from vllm.engine.arg_utils import EngineArgs +from vllm.exceptions import VLLMValidationError +from vllm.v1.engine.input_processor import InputProcessor +from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState from tests.helpers.mock import patch_hf_snapshot_download from vllm_omni.config.model import OmniModelConfig @@ -32,6 +36,36 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorMode pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("use_v2", [False, True]) +def test_stage_runner_selects_matching_input_processor_validator(monkeypatch, use_v2): + from vllm.sampling_params import SamplingParams + + from vllm_omni.model_executor.models.audex.cfg import AudexCFGLogitsProcessor + + # Upstream's process-wide preference must not override a stage's deploy choice. + monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", not use_v2) + config = VllmConfig() + model_config = object.__new__(OmniModelConfig) + model_config.use_v2_model_runner = use_v2 + model_config.runner_type = "generate" + model_config.logits_processors = [LogitBiasState if use_v2 else AudexCFGLogitsProcessor] + config.model_config = model_config + assert config.use_v2_model_runner is use_v2 + processor = object.__new__(InputProcessor) + processor.vllm_config, processor.model_config = config, model_config + validator = processor._build_logits_processors_params_validator() + validator(SamplingParams(extra_args={"cfg_role": "cond", "cfg_scale": 1.5})) + if not use_v2: + with pytest.raises(VLLMValidationError, match="cfg_role"): + validator(SamplingParams(extra_args={"cfg_role": "invalid"})) + + +@pytest.mark.parametrize("use_v2", [False, True]) +def test_plain_vllm_config_retains_upstream_runner_preference(monkeypatch, use_v2): + monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", use_v2) + assert VllmConfig().use_v2_model_runner is use_v2 + + def test_sync_config_is_omni(): """Ensure create_model_config gives the right type.""" cfg = OmniEngineArgs().create_model_config() diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 85b583b99..d8ab8c2c6 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -10,6 +10,7 @@ from functools import cached_property import torch from aenum import extend_enum from vllm.config import ModelConfig as _OriginalModelConfig +from vllm.config import VllmConfig as _OriginalVllmConfig from vllm.inputs import TokensPrompt as _OriginalTokensPrompt from vllm.model_executor.layers.rotary_embedding import ( MRotaryEmbedding as _OriginalMRotaryEmbedding, @@ -29,6 +30,23 @@ from vllm_omni.request import OmniRequest, OmniStreamingUpdate _PATCH_LOGGER = logging.getLogger("vllm_omni.patch") +# Omni's deploy config selects a runner per stage. vLLM 0.31's config property +# otherwise defaults to v2 independently of the runner selected by our workers, +# making shared consumers such as InputProcessor choose the wrong API. +_runner_property = _OriginalVllmConfig.use_v2_model_runner +if not getattr(_runner_property.fget, "_omni_stage_runner", False): + + def _patched_use_v2_model_runner(self, _upstream=_runner_property.fget): + # Lazy import keeps the package's patch/config initialization acyclic. + from vllm_omni.config.model import OmniModelConfig + + if isinstance(self.model_config, OmniModelConfig): + return self.model_config.use_v2_model_runner + return _upstream(self) + + _patched_use_v2_model_runner._omni_stage_runner = True + _OriginalVllmConfig.use_v2_model_runner = property(_patched_use_v2_model_runner) + # ============================================================================= # Patch ModelConfig.is_mm_prefix_lm to support omni-specific models # ============================================================================= @@ -390,8 +408,8 @@ def _patch_chat_template_registry(): ) if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: - _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = ( - lambda _: CHAT_TEMPLATES_DIR / "template_chatml.jinja" + _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = lambda _: ( + CHAT_TEMPLATES_DIR / "template_chatml.jinja" ) except ImportError: pass