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>
72 lines
1.9 KiB
Python
72 lines
1.9 KiB
Python
"""Unit tests for OmniModelStatePlugin interface correctness."""
|
|
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.worker_v2.model_states.plugin import OmniModelStatePlugin
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _ConcretePlugin(OmniModelStatePlugin):
|
|
"""Minimal concrete plugin for testing the ABC."""
|
|
|
|
def __init__(self):
|
|
self.added: list[tuple[int, Any]] = []
|
|
self.removed: list[int] = []
|
|
|
|
def on_add_request(self, req_index, new_req_data):
|
|
self.added.append((req_index, new_req_data))
|
|
|
|
def on_remove_request(self, req_index):
|
|
self.removed.append(req_index)
|
|
|
|
def postprocess(self, text_hidden, multimodal_outputs, input_batch, req_states):
|
|
return text_hidden, multimodal_outputs
|
|
|
|
|
|
def test_concrete_plugin_instantiates():
|
|
plugin = _ConcretePlugin()
|
|
assert isinstance(plugin, OmniModelStatePlugin)
|
|
|
|
|
|
def test_on_add_request_called():
|
|
plugin = _ConcretePlugin()
|
|
plugin.on_add_request(0, SimpleNamespace(req_id="r1"))
|
|
assert len(plugin.added) == 1
|
|
assert plugin.added[0][0] == 0
|
|
|
|
|
|
def test_on_remove_request_called():
|
|
plugin = _ConcretePlugin()
|
|
plugin.on_remove_request(3)
|
|
assert plugin.removed == [3]
|
|
|
|
|
|
def test_prepare_extra_inputs_default_empty():
|
|
plugin = _ConcretePlugin()
|
|
result = plugin.prepare_extra_inputs(SimpleNamespace(), SimpleNamespace())
|
|
assert result == {}
|
|
|
|
|
|
def test_dummy_run_default_noop():
|
|
plugin = _ConcretePlugin()
|
|
plugin.dummy_run(None, 128)
|
|
|
|
|
|
def test_postprocess_returns_inputs():
|
|
plugin = _ConcretePlugin()
|
|
hidden = torch.randn(4, 8)
|
|
mm = {"audio": torch.randn(4, 2)}
|
|
out_h, out_mm = plugin.postprocess(hidden, mm, None, None)
|
|
assert out_h is hidden
|
|
assert out_mm is mm
|
|
|
|
|
|
def test_abc_cannot_instantiate_directly():
|
|
with pytest.raises(TypeError):
|
|
OmniModelStatePlugin()
|