vllm-omni/tests/engine/test_arg_utils.py

"""
Tests for Omni config utils. For stability, these tests should largely be
invariant to the specific attributes of vLLM config except in cases where we
explicitly patch values that differ from vLLM.
"""

import argparse
import inspect
import json
import os
import shutil
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
from omegaconf import OmegaConf
from pydantic import ValidationError
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
from vllm.engine.arg_utils import EngineArgs

from vllm_omni.config.model import OmniModelConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_init_utils import build_engine_args_dict

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


def test_sync_config_is_omni():
    """Ensure create_model_config gives the right type."""
    cfg = OmniEngineArgs().create_model_config()
    assert isinstance(cfg, OmniModelConfig)


def test_default_stage_id_is_concrete_int():
    """Ensure `stage_id` stays safe for downstream arithmetic/indexing."""
    engine_args = OmniEngineArgs()

    assert engine_args.stage_id == 0
    assert isinstance(engine_args.stage_id, int)
    assert engine_args.log_stats is False

    cfg = engine_args.create_model_config()
    assert cfg.stage_id == 0


def test_multimodal_kwarg_overrides(mocker):
    """Ensure that overrides in the multimodal config are preserved."""
    sig = inspect.signature(OmniEngineArgs)
    default_mm_cache = sig.parameters["mm_processor_cache_gb"].default
    override_val = default_mm_cache + 1

    fake_model_config = SimpleNamespace(
        multimodal_config=SimpleNamespace(mm_processor_cache_gb=override_val),
    )

    def _fake_parent_create_model_config(self):
        assert self.mm_processor_cache_gb == override_val
        return fake_model_config

    mocker.patch.object(EngineArgs, "create_model_config", _fake_parent_create_model_config)
    mocker.patch.object(OmniModelConfig, "from_vllm_model_config", side_effect=lambda model_config, **_: model_config)

    cfg = OmniEngineArgs(
        model="Qwen/Qwen2-VL-2B-Instruct",
        mm_processor_cache_gb=override_val,
    ).create_model_config()

    assert cfg.multimodal_config is not None
    assert cfg.multimodal_config.mm_processor_cache_gb == override_val


def test_from_vllm_config_validates_invalid_omni_kwargs():
    """Ensure omni-specific field validation catches invalid keys."""
    model_config = EngineArgs().create_model_config()
    with pytest.raises(ValueError, match="Unexpected omni kwarg"):
        OmniModelConfig.from_vllm_model_config(model_config, foo="bar")


def test_from_vllm_config_validates_bad_omni_kwarg_types():
    """Ensure omni-specific field validation catches type errors."""
    model_config = EngineArgs().create_model_config()
    with pytest.raises(ValidationError):
        OmniModelConfig.from_vllm_model_config(model_config, stage_id="not_an_int")


def test_default_all_values_are_initialized():
    """Ensure omni-specific field initializes all fields"""
    model_config = EngineArgs().create_model_config()
    cfg = OmniModelConfig.from_vllm_model_config(model_config)

    # Test a primitive
    assert cfg.model_stage == "thinker"
    # Test a field initialized with a default factory
    assert cfg.stage_connector_config == {
        "name": "SharedMemoryConnector",
        "extra": {},
    }

    # Ensure that hf_config is initialized on model_config in the vLLM by ModelConfig's
    # __post_init__, and that the hf_config is copied over to the OmniModelConfig;
    # we explicitly set this since the field sets init=False
    assert isinstance(model_config.hf_config, PretrainedConfig)
    assert cfg.hf_config is model_config.hf_config

    # Ensure that we can convert it to a string; this will convert
    # all attributes, so should raise if we have attributes that are
    # not initialized correctly, e.g., due to default factories
    str(cfg)


def test_qwen3_tts_codec_frame_rate_patching():
    """Ensure the patch for qwen3 tts is applied correctly when creating the omni config."""
    # Create a vLLM ModelConfig
    vllm_config = EngineArgs().create_model_config()

    # Create a mock talking config with a dummy value for position_id_per_seconds
    mock_talker_config = SimpleNamespace()
    mock_talker_config.position_id_per_seconds = 12.3
    vllm_config.hf_config.talker_config = mock_talker_config

    # Ensure creating the config for a Qwen3TTSTalkerForConditionalGenerationARVLLM
    # model calls the patch func to apply position_id_per_seconds from the talker
    # config to the config's codec_frame_rate_hz
    omni_config = OmniModelConfig.from_vllm_model_config(
        vllm_config,
        model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM",
    )

    # Verify codec_frame_rate_hz was patched
    assert omni_config.codec_frame_rate_hz == 12.3


def test_qwen3_tts_startup_task_type_is_validated():
    vllm_config = EngineArgs().create_model_config()

    config = OmniModelConfig.from_vllm_model_config(
        vllm_config,
        model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM",
        task_type="Base",
    )
    assert config.task_type == "Base"

    with pytest.raises(ValueError, match="Qwen3-TTS --task-type must be one of"):
        OmniModelConfig.from_vllm_model_config(
            vllm_config,
            model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM",
            task_type="fl2va",
        )


def test_from_cli_args_picks_up_stage_configs_path():
    """from_cli_args should pick up stage_configs_path from namespace."""
    ns = argparse.Namespace(
        model="facebook/opt-125m",
        stage_configs_path="/some/path.yaml",
        custom_pipeline_args=None,
    )

    args = OmniEngineArgs.from_cli_args(ns)
    assert args.stage_configs_path == "/some/path.yaml"
    assert args.custom_pipeline_args is None


def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
    """Ensure Code2Wav mirrors stage max_model_len into nested HF overrides.

    Qwen3-TTS Code2Wav is a pure decoder stage whose runtime max_model_len can
    legitimately exceed the base checkpoint's default text max length. Recent
    vLLM validates these values during ModelConfig creation, so we inject
    ``talker_config.max_position_embeddings`` before delegating to vLLM.
    """
    captured: dict[str, object] = {}
    baseline_config = Mock()

    def fake_create_model_config(self):
        captured["hf_overrides"] = self.hf_overrides
        return baseline_config

    monkeypatch.setattr(EngineArgs, "create_model_config", fake_create_model_config)
    monkeypatch.setattr(
        OmniModelConfig,
        "from_vllm_model_config",
        classmethod(lambda cls, model_config, **omni_kwargs: model_config),
    )

    OmniEngineArgs(
        model_arch="Qwen3TTSCode2Wav",
        max_model_len=65536,
    ).create_model_config()

    assert captured["hf_overrides"] == {
        "architectures": ["Qwen3TTSCode2Wav"],
        "talker_config": {
            "max_position_embeddings": 65536,
        },
    }


def test_patch_missing_local_hf_config(tmp_path):
    """Native model bundles may only contain a checkpoints/ directory."""
    (tmp_path / "checkpoints").mkdir()
    args = object.__new__(OmniEngineArgs)
    args.model = str(tmp_path)
    args.model_arch = "IndexTTS25TalkerForConditionalGeneration"
    args.hf_config_path = None

    try:
        args._patch_empty_hf_config("indextts2_5")

        assert args.hf_config_path is not None
        config_path = args.hf_config_path + "/config.json"
        with open(config_path, encoding="utf-8") as config_file:
            config = json.load(config_file)
        assert config == {
            "model_type": "indextts2_5",
            "architectures": ["IndexTTS25TalkerForConditionalGeneration"],
        }
    finally:
        if hasattr(args, "_temp_config_dir"):
            shutil.rmtree(args._temp_config_dir, ignore_errors=True)


@pytest.mark.parametrize(
    "config_entry",
    ["malformed", "directory", "broken_symlink", "permission_error"],
)
def test_non_missing_local_hf_config_error_reaches_parent_loader(tmp_path, monkeypatch, config_entry):
    """Non-missing config errors must reach vLLM's normal loader."""
    config_path = tmp_path / "config.json"
    if config_entry == "malformed":
        config_path.write_text("{not valid json", encoding="utf-8")
        loader_error = json.JSONDecodeError("invalid config", "{not valid json", 1)
    elif config_entry == "directory":
        config_path.mkdir()
        loader_error = IsADirectoryError(str(config_path))
    elif config_entry == "broken_symlink":
        config_path.symlink_to(tmp_path / "missing-target.json")
        loader_error = FileNotFoundError(str(config_path))
    else:
        config_path.write_text("{}", encoding="utf-8")
        loader_error = PermissionError(str(config_path))

    parent_inputs = {}

    def controlled_get_config_dict(cls, model, **kwargs):
        if model == str(tmp_path):
            raise loader_error
        with open(os.path.join(model, "config.json"), encoding="utf-8") as config_file:
            return json.load(config_file), {}

    def fake_parent_create_model_config(self):
        parent_inputs["called"] = True
        parent_inputs["model"] = self.model
        parent_inputs["hf_config_path"] = self.hf_config_path
        config_root = self.hf_config_path or self.model
        parent_inputs["config_root"] = config_root
        return PretrainedConfig.get_config_dict(config_root)[0]

    monkeypatch.setattr(PretrainedConfig, "get_config_dict", classmethod(controlled_get_config_dict))
    monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda self: True)
    monkeypatch.setattr(EngineArgs, "create_model_config", fake_parent_create_model_config)
    monkeypatch.setattr(
        OmniModelConfig,
        "from_vllm_model_config",
        classmethod(lambda cls, model_config, **omni_kwargs: model_config),
    )

    args = OmniEngineArgs(
        model=str(tmp_path),
        model_arch="IndexTTS25TalkerForConditionalGeneration",
    )

    with pytest.raises(type(loader_error)) as exc_info:
        args.create_model_config()

    assert str(exc_info.value) == str(loader_error)
    assert parent_inputs["called"] is True
    assert parent_inputs["model"] == str(tmp_path)
    assert parent_inputs["config_root"] == str(tmp_path)
    assert parent_inputs["hf_config_path"] is None
    assert args.model == str(tmp_path)
    assert args.hf_config_path is None
    assert not hasattr(args, "_temp_config_dir")


def test_patch_valid_hf_config_without_model_type_preserves_keys(tmp_path):
    """A valid partial config should be copied before required keys are added."""
    original_config = {
        "hidden_size": 1024,
        "custom_key": {"nested": True},
    }
    (tmp_path / "config.json").write_text(json.dumps(original_config), encoding="utf-8")
    args = object.__new__(OmniEngineArgs)
    args.model = str(tmp_path)
    args.model_arch = "IndexTTS25TalkerForConditionalGeneration"
    args.hf_config_path = None

    try:
        args._patch_empty_hf_config("indextts2_5")

        assert args.hf_config_path is not None
        with open(os.path.join(args.hf_config_path, "config.json"), encoding="utf-8") as config_file:
            patched_config = json.load(config_file)
        assert patched_config["hidden_size"] == original_config["hidden_size"]
        assert patched_config["custom_key"] == original_config["custom_key"]
        assert patched_config["model_type"] == "indextts2_5"
        assert patched_config["architectures"] == ["IndexTTS25TalkerForConditionalGeneration"]
    finally:
        if hasattr(args, "_temp_config_dir"):
            shutil.rmtree(args._temp_config_dir, ignore_errors=True)


def test_remote_hf_config_error_reaches_parent_loader(monkeypatch):
    """Remote resolution failures must stay on vLLM's normal error path."""
    loader_error = OSError("remote config resolution failed")
    parent_inputs = {}

    def controlled_get_config_dict(cls, model, **kwargs):
        raise loader_error

    def fake_parent_create_model_config(self):
        parent_inputs["called"] = True
        parent_inputs["model"] = self.model
        parent_inputs["hf_config_path"] = self.hf_config_path
        return PretrainedConfig.get_config_dict(self.model)[0]

    monkeypatch.setattr(PretrainedConfig, "get_config_dict", classmethod(controlled_get_config_dict))
    monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda self: True)
    monkeypatch.setattr(EngineArgs, "create_model_config", fake_parent_create_model_config)

    args = OmniEngineArgs(
        model="remote/model",
        model_arch="IndexTTS25TalkerForConditionalGeneration",
    )

    with pytest.raises(OSError) as exc_info:
        args.create_model_config()

    assert str(exc_info.value) == str(loader_error)
    assert parent_inputs["called"] is True
    assert parent_inputs["model"] == "remote/model"
    assert parent_inputs["hf_config_path"] is None
    assert args.model == "remote/model"
    assert args.hf_config_path is None
    assert not hasattr(args, "_temp_config_dir")


def test_stage_specific_text_config_override():
    """Stage swap must refresh hf_text_config, dependent attrs, and model_arch_config."""
    vllm_config = EngineArgs().create_model_config()
    vllm_config.disable_sliding_window = True
    thinker_mac = vllm_config.model_arch_config

    talker_num_heads = max(2, thinker_mac.total_num_attention_heads // 2)
    talker_num_kv_heads = max(1, talker_num_heads // 8)
    talker_head_dim = 128
    stage_text_config = SimpleNamespace(
        sliding_window=4096,
        attention_chunk_size=2048,
        max_position_embeddings=4096,
        num_attention_heads=talker_num_heads,
        num_key_value_heads=talker_num_kv_heads,
        head_dim=talker_head_dim,
        hidden_size=talker_num_heads * talker_head_dim,
        vocab_size=thinker_mac.vocab_size,
        num_hidden_layers=4,
    )

    vllm_config.hf_text_config = SimpleNamespace()
    vllm_config.hf_config.thinker_config = SimpleNamespace(get_text_config=lambda: stage_text_config)

    omni_config = OmniModelConfig.from_vllm_model_config(
        vllm_config,
        hf_config_name="thinker_config",
    )

    assert omni_config.hf_text_config is stage_text_config
    assert omni_config.attention_chunk_size == 2048
    assert omni_config.max_model_len == 4096
    assert omni_config.hf_text_config.sliding_window is None

    stage_mac = omni_config.model_arch_config
    assert stage_mac is not thinker_mac
    assert stage_mac.total_num_attention_heads == talker_num_heads
    assert stage_mac.total_num_kv_heads == talker_num_kv_heads
    assert stage_mac.head_size == talker_head_dim

    parallel_config = SimpleNamespace(
        tensor_parallel_size=1,
        pipeline_parallel_size=1,
        decode_context_parallel_size=1,
    )
    assert omni_config.get_num_attention_heads(parallel_config) == talker_num_heads
    assert omni_config.get_num_kv_heads(parallel_config) == talker_num_kv_heads
    assert omni_config.get_head_size() == talker_head_dim


def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size():
    """A stage without an embedding-size override keeps the default input width."""
    vllm_config = EngineArgs().create_model_config()
    vllm_config.hf_config = Qwen3OmniMoeConfig(enable_audio_output=True)

    omni_config = OmniModelConfig.from_vllm_model_config(
        vllm_config,
        model_stage="talker",
        hf_config_name="talker_config",
    )

    assert (
        getattr(
            omni_config.hf_config.talker_config,
            "embedding_size",
            None,
        )
        is None
    )
    assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size


def test_stage_configs_path_field():
    """OmniEngineArgs with stage_configs_path should construct without error."""
    args = OmniEngineArgs(stage_configs_path="/some/path.yaml")
    assert args.stage_configs_path == "/some/path.yaml"


def test_strip_single_engine_args():
    """_strip_single_engine_args should remove EngineArgs fields but keep omni fields."""
    kwargs = {
        # Parent EngineArgs fields — stripped unless explicitly allowlisted
        "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
        "tensor_parallel_size": 4,
        "gpu_memory_utilization": 0.9,
        "model": "some/model",
        # Parent field that should be kept (allowlisted)
        "worker_extension_cls": "some.Extension",
        # OmniEngineArgs-only / non-engine fields — should pass through
        "stage_configs_path": "/path/to/yaml",
        "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
        "mode": "text-to-image",
        "lora_path": "/some/lora",
    }

    filtered = AsyncOmniEngine._strip_single_engine_args(kwargs)

    # Stripped — parent EngineArgs fields
    assert "compilation_config" not in filtered
    assert filtered["tensor_parallel_size"] == 4
    assert "gpu_memory_utilization" not in filtered
    assert "model" not in filtered

    # Stripped — orchestrator-level OmniEngineArgs field
    assert "stage_configs_path" not in filtered

    # Kept
    assert filtered["worker_extension_cls"] == "some.Extension"
    assert filtered["custom_pipeline_args"] == {"pipeline_class": "my.Pipeline"}
    assert filtered["mode"] == "text-to-image"
    assert filtered["lora_path"] == "/some/lora"


def test_strip_single_engine_args_model_does_not_trigger_warning(mocker):
    """model is always in kwargs (callers set it via from_cli_args/asdict),
    so it should not cause the override warning by itself or appear in it."""
    mock_warn = mocker.patch("vllm_omni.engine.async_omni_engine.logger.warning")

    # Typical caller kwargs: model is always present, no other parent
    # EngineArgs fields are explicitly overridden.
    AsyncOmniEngine._strip_single_engine_args(
        {
            "model": "some/model",
            "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
        }
    )
    mock_warn.assert_not_called()

    # When there *are* genuinely surprising overrides alongside model,
    # the warning should mention them but not model. Keep-listed fields such as
    # tensor_parallel_size are intentionally passed through and should not warn.
    AsyncOmniEngine._strip_single_engine_args(
        {
            "model": "some/model",
            "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
            "tensor_parallel_size": 4,
            "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
        }
    )
    mock_warn.assert_called_once()
    warned_args = mock_warn.call_args[0][-1]  # the formatted arg list
    assert "compilation_config" in warned_args
    assert "tensor_parallel_size" not in warned_args
    assert "model" not in warned_args


def test_new_format_deploy_yaml_does_not_strip_interleave_cli_overrides(tmp_path):
    """New-format deploy YAMLs must not strip interleave_mm_strings / media_io_kwargs.

    Pytest historically passed ``--stage-configs-path`` for deploy YAMLs; stripping
    those CLI flags dropped MiniCPM interleaved AV packing and cut Daily-Omni
    accuracy from ~78% to ~70%.
    """
    from vllm_omni.entrypoints.utils import is_new_format_deploy_config

    yaml_path = tmp_path / "minicpm_deploy.yaml"
    yaml_path.write_text(
        """
pipeline: minicpmo_4_5
async_chunk: true
stages:
  - stage_id: 0
    devices: "0"
""".strip(),
        encoding="utf-8",
    )
    path = str(yaml_path)
    assert is_new_format_deploy_config(path)

    kwargs = {
        "stage_configs_path": path,
        "trust_remote_code": True,
        "interleave_mm_strings": True,
        "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
        "gpu_memory_utilization": 0.55,
    }
    # Mirror AsyncOmniEngine._resolve_stage_configs strip decision.
    if kwargs.get("stage_configs_path") is not None and not is_new_format_deploy_config(kwargs["stage_configs_path"]):
        base = AsyncOmniEngine._strip_single_engine_args(dict(kwargs))
    else:
        base = dict(kwargs)
    assert base.get("interleave_mm_strings") is True
    assert base.get("media_io_kwargs") == {"video": {"fps": 1, "num_frames": 128}}
    # Parent EngineArgs fields also remain available for deploy merge.
    assert base.get("gpu_memory_utilization") == 0.55

    legacy = tmp_path / "legacy.yaml"
    legacy.write_text("stage_args:\n  - stage_id: 0\n", encoding="utf-8")
    assert not is_new_format_deploy_config(str(legacy))
    legacy_kwargs = dict(kwargs, stage_configs_path=str(legacy))
    stripped = AsyncOmniEngine._strip_single_engine_args(legacy_kwargs)
    assert "interleave_mm_strings" not in stripped
    assert "media_io_kwargs" not in stripped
    assert "gpu_memory_utilization" not in stripped


# For https://github.com/vllm-project/vllm-omni/issues/3293
def test_tensor_parallel_size_none_is_handled():
    """Ensure the tensor parallel size of None isn't forwarded."""
    engine_args = OmegaConf.create({"stage_id": 0, "engine_args": {"tensor_parallel_size": None}})
    args = build_engine_args_dict(
        engine_args,
        model="snu-aidas/Dynin-Omni",
    )
    assert isinstance(args, dict)
    assert "tensor_parallel_size" not in args