71 lines
1.9 KiB
Python
71 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()
|