68 lines
2.9 KiB
Python
68 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.v1.worker.gpu.sample.output import SamplerOutput
|
|
|
|
from vllm_omni.worker_v2.omni_sampler import OmniSampler, OmniSamplingOutput, sample_with_output
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _output():
|
|
counts, rejected = torch.tensor([1, 0]), torch.tensor([0, 0])
|
|
return SamplerOutput(torch.tensor([[7], [0]]), None, None, counts, rejected), counts, rejected
|
|
|
|
|
|
@pytest.mark.parametrize("wrapped", [False, True])
|
|
def test_standard_sampling_keeps_counts_grammar_and_lifecycle(wrapped):
|
|
base = Mock()
|
|
base.req_states = object()
|
|
sampler = OmniSampler(base) if wrapped else base
|
|
expected = _output()
|
|
standard = Mock(return_value=expected)
|
|
hidden, batch, grammar = torch.zeros(2, 4), object(), object()
|
|
result = sample_with_output(sampler, standard, hidden, batch, base.req_states, grammar)
|
|
assert result.sampler_output is expected[0]
|
|
assert result.num_sampled is expected[1] and result.num_rejected is expected[2]
|
|
assert result.multimodal_outputs is None
|
|
standard.assert_called_once_with(hidden, batch, grammar)
|
|
sampler.add_request(3, "request")
|
|
sampler.apply_staged_writes()
|
|
base.add_request.assert_called_once_with(3, "request")
|
|
base.apply_staged_writes.assert_called_once_with()
|
|
assert sampler.req_states is base.req_states
|
|
sampler(hidden, batch)
|
|
base.assert_called_once_with(hidden, batch)
|
|
|
|
|
|
@pytest.mark.parametrize("payload", [{}, {"codes.audio": torch.tensor([[1, 2]])}])
|
|
def test_audio_sampler_uses_registered_adapter_without_standard_resampling(payload):
|
|
expected = OmniSamplingOutput(*_output(), multimodal_outputs=payload)
|
|
|
|
class AudioSampler(OmniSampler):
|
|
def sample_step(self, hidden, batch, states, grammar, standard):
|
|
assert states is self.req_states
|
|
if grammar is not None:
|
|
raise ValueError("unsupported grammar")
|
|
return expected
|
|
|
|
sampler = AudioSampler(SimpleNamespace(req_states=object()))
|
|
standard = Mock(side_effect=AssertionError("must not resample"))
|
|
result = sample_with_output(sampler, standard, torch.zeros(2, 4), object(), sampler.req_states, None)
|
|
assert result is expected and result.multimodal_outputs is payload
|
|
with pytest.raises(ValueError, match="unsupported grammar"):
|
|
sample_with_output(sampler, standard, torch.zeros(2, 4), object(), sampler.req_states, object())
|
|
standard.assert_not_called()
|
|
|
|
|
|
def test_adapter_preserves_stock_sampler_staged_write_fast_path():
|
|
from vllm.v1.worker.gpu.sample.sampler import Sampler
|
|
|
|
assert OmniSampler(object.__new__(Sampler)).omni_static_staged_writes
|
|
assert not OmniSampler(SimpleNamespace()).omni_static_staged_writes
|
|
assert OmniSampler(SimpleNamespace(omni_static_staged_writes=True)).omni_static_staged_writes
|