83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
|
"""Unit tests for init_omni_model_state factory dispatch."""
|
||
|
|
|
||
|
|
from types import SimpleNamespace
|
||
|
|
from unittest.mock import MagicMock, patch
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from vllm_omni.worker_v2.model_states import (
|
||
|
|
_OMNI_ARCHITECTURES,
|
||
|
|
init_omni_model_state,
|
||
|
|
)
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||
|
|
|
||
|
|
|
||
|
|
def _make_vllm_config(architectures):
|
||
|
|
model_config = SimpleNamespace(architectures=architectures)
|
||
|
|
return SimpleNamespace(model_config=model_config)
|
||
|
|
|
||
|
|
|
||
|
|
@patch("vllm_omni.worker_v2.model_states.omni_model_state.OmniModelState.__init__", return_value=None)
|
||
|
|
def test_qwen3_omni_dispatches_to_omni_model_state(mock_init):
|
||
|
|
cfg = _make_vllm_config(["Qwen3OmniMoeForConditionalGeneration"])
|
||
|
|
model = MagicMock()
|
||
|
|
device = torch.device("cpu")
|
||
|
|
|
||
|
|
state = init_omni_model_state(cfg, model, None, device)
|
||
|
|
|
||
|
|
from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
|
||
|
|
|
||
|
|
assert isinstance(state, OmniModelState)
|
||
|
|
|
||
|
|
|
||
|
|
@patch("vllm_omni.worker_v2.model_states.omni_model_state.OmniModelState.__init__", return_value=None)
|
||
|
|
def test_mammoth_dispatches_to_omni_model_state(mock_init):
|
||
|
|
cfg = _make_vllm_config(["MammothModa2ForConditionalGeneration"])
|
||
|
|
model = MagicMock()
|
||
|
|
device = torch.device("cpu")
|
||
|
|
|
||
|
|
state = init_omni_model_state(cfg, model, None, device)
|
||
|
|
|
||
|
|
from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
|
||
|
|
|
||
|
|
assert isinstance(state, OmniModelState)
|
||
|
|
|
||
|
|
|
||
|
|
@patch("vllm_omni.worker_v2.model_states._upstream_init_model_state")
|
||
|
|
def test_unknown_arch_delegates_to_upstream(mock_upstream):
|
||
|
|
mock_upstream.return_value = MagicMock()
|
||
|
|
cfg = _make_vllm_config(["LlamaForCausalLM"])
|
||
|
|
model = MagicMock()
|
||
|
|
device = torch.device("cpu")
|
||
|
|
|
||
|
|
state = init_omni_model_state(cfg, model, None, device)
|
||
|
|
|
||
|
|
mock_upstream.assert_called_once_with(cfg, model, None, device)
|
||
|
|
assert state is mock_upstream.return_value
|
||
|
|
|
||
|
|
|
||
|
|
def test_omni_architectures_set_contains_expected():
|
||
|
|
expected = {
|
||
|
|
"Qwen3OmniMoeForConditionalGeneration",
|
||
|
|
"Qwen2_5OmniForConditionalGeneration",
|
||
|
|
"MammothModa2ForConditionalGeneration",
|
||
|
|
"MiMoAudioForConditionalGeneration",
|
||
|
|
"MammothModa2ARForConditionalGeneration",
|
||
|
|
"Qwen3TTSTalkerForConditionalGeneration",
|
||
|
|
"Qwen3TTSCode2Wav",
|
||
|
|
}
|
||
|
|
assert _OMNI_ARCHITECTURES == expected
|
||
|
|
|
||
|
|
|
||
|
|
@patch("vllm_omni.worker_v2.model_states._upstream_init_model_state")
|
||
|
|
def test_none_architectures_delegates_to_upstream(mock_upstream):
|
||
|
|
mock_upstream.return_value = MagicMock()
|
||
|
|
cfg = _make_vllm_config(None)
|
||
|
|
model = MagicMock()
|
||
|
|
device = torch.device("cpu")
|
||
|
|
|
||
|
|
init_omni_model_state(cfg, model, None, device)
|
||
|
|
|
||
|
|
mock_upstream.assert_called_once()
|