89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Msgpack roundtrip tests for OmniEngineCoreOutputs with multimodal_output.
|
|
|
|
Validates that tensor-only payloads survive msgspec encode/decode and that
|
|
non-tensor values fail decoding (enforcing the wire invariant).
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
|
|
|
|
from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _roundtrip(outputs: OmniEngineCoreOutputs) -> OmniEngineCoreOutputs:
|
|
encoder = MsgpackEncoder()
|
|
encoded = encoder.encode(outputs)
|
|
decoder = MsgpackDecoder(OmniEngineCoreOutputs)
|
|
return decoder.decode(encoded)
|
|
|
|
|
|
def test_tensor_only_roundtrip():
|
|
"""Tensor-only multimodal_output survives msgpack roundtrip."""
|
|
audio = torch.randn(16000)
|
|
sr = torch.tensor(24000)
|
|
eco = OmniEngineCoreOutput(
|
|
request_id="req-1",
|
|
new_token_ids=[1, 2, 3],
|
|
finish_reason=None,
|
|
multimodal_output={"audio": audio, "sr": sr},
|
|
)
|
|
decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
|
|
assert len(decoded.outputs) == 1
|
|
out = decoded.outputs[0]
|
|
assert out.request_id == "req-1"
|
|
assert isinstance(out.multimodal_output, dict)
|
|
assert torch.allclose(out.multimodal_output["audio"], audio)
|
|
assert out.multimodal_output["sr"].item() == 24000
|
|
|
|
|
|
def test_empty_multimodal_roundtrip():
|
|
"""None multimodal_output roundtrips correctly."""
|
|
eco = OmniEngineCoreOutput(
|
|
request_id="req-2",
|
|
new_token_ids=[4],
|
|
finish_reason=None,
|
|
multimodal_output=None,
|
|
)
|
|
decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
|
|
assert decoded.outputs[0].multimodal_output is None
|
|
|
|
|
|
def test_multiple_tensor_keys_roundtrip():
|
|
"""Multiple tensor keys survive roundtrip."""
|
|
hidden = torch.randn(10, 4096)
|
|
codes = torch.randint(0, 1024, (1, 100))
|
|
eco = OmniEngineCoreOutput(
|
|
request_id="req-3",
|
|
new_token_ids=[],
|
|
finish_reason=None,
|
|
multimodal_output={"hidden": hidden, "codes": codes},
|
|
)
|
|
decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
|
|
out = decoded.outputs[0]
|
|
assert torch.allclose(out.multimodal_output["hidden"], hidden)
|
|
assert torch.equal(out.multimodal_output["codes"], codes)
|
|
|
|
|
|
def test_non_tensor_value_fails_decode():
|
|
"""Non-tensor values in multimodal_output cause decode failure.
|
|
|
|
This validates that the wire field must be tensor-only and that
|
|
the _ensure_tensor_values boundary in model runners is necessary.
|
|
"""
|
|
eco = OmniEngineCoreOutput(
|
|
request_id="req-bad",
|
|
new_token_ids=[],
|
|
finish_reason=None,
|
|
multimodal_output={"audio": torch.randn(10), "sr": 24000}, # scalar!
|
|
)
|
|
encoder = MsgpackEncoder()
|
|
encoded = encoder.encode(OmniEngineCoreOutputs(outputs=[eco]))
|
|
decoder = MsgpackDecoder(OmniEngineCoreOutputs)
|
|
with pytest.raises(Exception):
|
|
decoder.decode(encoded)
|