Taichang ZhouandClaude Opus 4.8
9678ae58cf
Squashed migration rebased onto current main. Collapses the 46-commit dev/migrate-MR-v2 history (10 merge commits + superseded pre-rebase work) into a single commit, so the diff reflects only the genuine MR-V2 feature (worker_v2/ runners + model_states, omni schedulers, stage engine pipeline, pooling_output serialization bridge). Content already merged to main separately (step_audio2, tts_adapters, npu/_310p, etc.) drops out. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
84 lines
2.6 KiB
Python
84 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()
|