252 lines
11 KiB
Python
252 lines
11 KiB
Python
|
"""Tests for data_entry_keys: TypedDict payload structure, flatten/unflatten, serialize/deserialize."""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from vllm_omni.data_entry_keys import (
|
||
|
|
OmniPayload,
|
||
|
|
deserialize_payload,
|
||
|
|
flatten_payload,
|
||
|
|
serialize_payload,
|
||
|
|
unflatten_payload,
|
||
|
|
)
|
||
|
|
from vllm_omni.engine import AdditionalInformationPayload
|
||
|
|
|
||
|
|
|
||
|
|
class TestOmniPayload:
|
||
|
|
def test_nested_payload_structure(self):
|
||
|
|
"""Verify OmniPayload can be constructed with nested dicts."""
|
||
|
|
payload: OmniPayload = {
|
||
|
|
"hidden_states": {"output": torch.tensor([1.0])},
|
||
|
|
"embed": {"prefill": torch.tensor([2.0])},
|
||
|
|
"codes": {"audio": torch.tensor([3.0])},
|
||
|
|
"ids": {"all": [1, 2, 3]},
|
||
|
|
"meta": {"finished": torch.tensor(True, dtype=torch.bool)},
|
||
|
|
}
|
||
|
|
assert torch.equal(payload["hidden_states"]["output"], torch.tensor([1.0]))
|
||
|
|
assert torch.equal(payload["embed"]["prefill"], torch.tensor([2.0]))
|
||
|
|
assert torch.equal(payload["codes"]["audio"], torch.tensor([3.0]))
|
||
|
|
assert payload["ids"]["all"] == [1, 2, 3]
|
||
|
|
assert payload["meta"]["finished"].item() is True
|
||
|
|
|
||
|
|
def test_partial_payload(self):
|
||
|
|
"""OmniPayload fields are all optional (total=False)."""
|
||
|
|
payload: OmniPayload = {"meta": {"finished": torch.tensor(False, dtype=torch.bool)}}
|
||
|
|
assert payload["meta"]["finished"].item() is False
|
||
|
|
|
||
|
|
def test_empty_payload(self):
|
||
|
|
payload: OmniPayload = {}
|
||
|
|
assert len(payload) == 0
|
||
|
|
|
||
|
|
|
||
|
|
class TestFlattenPayload:
|
||
|
|
def test_basic_nested_to_dotted(self):
|
||
|
|
nested = {
|
||
|
|
"codes": {"audio": torch.tensor([1.0])},
|
||
|
|
"meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 5},
|
||
|
|
}
|
||
|
|
flat = flatten_payload(nested)
|
||
|
|
assert torch.equal(flat["codes.audio"], torch.tensor([1.0]))
|
||
|
|
assert flat["meta.finished"].item() is True
|
||
|
|
assert flat["meta.left_context_size"] == 5
|
||
|
|
assert "codes" not in flat
|
||
|
|
assert "meta" not in flat
|
||
|
|
|
||
|
|
def test_top_level_keys_preserved(self):
|
||
|
|
nested = {
|
||
|
|
"latent": torch.tensor([9.0]),
|
||
|
|
"generated_len": 42,
|
||
|
|
}
|
||
|
|
flat = flatten_payload(nested)
|
||
|
|
assert torch.equal(flat["latent"], torch.tensor([9.0]))
|
||
|
|
assert flat["generated_len"] == 42
|
||
|
|
|
||
|
|
def test_hidden_states_layers_expanded(self):
|
||
|
|
nested = {
|
||
|
|
"hidden_states": {
|
||
|
|
"output": torch.tensor([1.0]),
|
||
|
|
"layers": {
|
||
|
|
0: torch.tensor([2.0]),
|
||
|
|
24: torch.tensor([3.0]),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
flat = flatten_payload(nested)
|
||
|
|
assert torch.equal(flat["hidden_states.output"], torch.tensor([1.0]))
|
||
|
|
assert torch.equal(flat["hidden_states.layer_0"], torch.tensor([2.0]))
|
||
|
|
assert torch.equal(flat["hidden_states.layer_24"], torch.tensor([3.0]))
|
||
|
|
assert "hidden_states.layers" not in flat
|
||
|
|
|
||
|
|
def test_empty_payload(self):
|
||
|
|
assert flatten_payload({}) == {}
|
||
|
|
|
||
|
|
def test_mixed_nested_and_top_level(self):
|
||
|
|
nested: OmniPayload = {
|
||
|
|
"codes": {"audio": torch.tensor([1.0])},
|
||
|
|
"latent": torch.tensor([2.0]),
|
||
|
|
"meta": {"finished": torch.tensor(False, dtype=torch.bool)},
|
||
|
|
}
|
||
|
|
flat = flatten_payload(nested)
|
||
|
|
assert set(flat.keys()) == {"codes.audio", "latent", "meta.finished"}
|
||
|
|
|
||
|
|
|
||
|
|
class TestUnflattenPayload:
|
||
|
|
def test_basic_dotted_to_nested(self):
|
||
|
|
flat = {
|
||
|
|
"codes.audio": torch.tensor([1.0]),
|
||
|
|
"meta.finished": torch.tensor(True, dtype=torch.bool),
|
||
|
|
"meta.left_context_size": 5,
|
||
|
|
}
|
||
|
|
nested = unflatten_payload(flat)
|
||
|
|
assert torch.equal(nested["codes"]["audio"], torch.tensor([1.0]))
|
||
|
|
assert nested["meta"]["finished"].item() is True
|
||
|
|
assert nested["meta"]["left_context_size"] == 5
|
||
|
|
|
||
|
|
def test_top_level_keys_preserved(self):
|
||
|
|
flat = {"latent": torch.tensor([9.0]), "generated_len": 42}
|
||
|
|
nested = unflatten_payload(flat)
|
||
|
|
assert torch.equal(nested["latent"], torch.tensor([9.0]))
|
||
|
|
assert nested["generated_len"] == 42
|
||
|
|
|
||
|
|
def test_hidden_states_layers_collected(self):
|
||
|
|
flat = {
|
||
|
|
"hidden_states.output": torch.tensor([1.0]),
|
||
|
|
"hidden_states.layer_0": torch.tensor([2.0]),
|
||
|
|
"hidden_states.layer_24": torch.tensor([3.0]),
|
||
|
|
}
|
||
|
|
nested = unflatten_payload(flat)
|
||
|
|
assert torch.equal(nested["hidden_states"]["output"], torch.tensor([1.0]))
|
||
|
|
assert torch.equal(nested["hidden_states"]["layers"][0], torch.tensor([2.0]))
|
||
|
|
assert torch.equal(nested["hidden_states"]["layers"][24], torch.tensor([3.0]))
|
||
|
|
|
||
|
|
def test_empty_payload(self):
|
||
|
|
assert unflatten_payload({}) == {}
|
||
|
|
|
||
|
|
|
||
|
|
class TestFlattenUnflattenRoundTrip:
|
||
|
|
def test_round_trip_simple(self):
|
||
|
|
original: OmniPayload = {
|
||
|
|
"codes": {"audio": torch.tensor([1.0, 2.0])},
|
||
|
|
"meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 10},
|
||
|
|
"ids": {"prompt": [1, 2, 3]},
|
||
|
|
"latent": torch.tensor([5.0]),
|
||
|
|
}
|
||
|
|
restored = unflatten_payload(flatten_payload(original))
|
||
|
|
assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
|
||
|
|
assert restored["meta"]["finished"].item() is True
|
||
|
|
assert restored["meta"]["left_context_size"] == 10
|
||
|
|
assert restored["ids"]["prompt"] == [1, 2, 3]
|
||
|
|
assert torch.equal(restored["latent"], original["latent"])
|
||
|
|
|
||
|
|
def test_round_trip_with_layers(self):
|
||
|
|
original = {
|
||
|
|
"hidden_states": {
|
||
|
|
"output": torch.tensor([1.0]),
|
||
|
|
"layers": {0: torch.tensor([2.0]), 24: torch.tensor([3.0])},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
restored = unflatten_payload(flatten_payload(original))
|
||
|
|
assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
|
||
|
|
assert torch.equal(restored["hidden_states"]["layers"][0], torch.tensor([2.0]))
|
||
|
|
assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
|
||
|
|
|
||
|
|
def test_round_trip_all_categories(self):
|
||
|
|
original: OmniPayload = {
|
||
|
|
"hidden_states": {"output": torch.tensor([1.0]), "last": torch.tensor([2.0])},
|
||
|
|
"embed": {"prefill": torch.tensor([3.0]), "tts_bos": torch.tensor([4.0])},
|
||
|
|
"codes": {"audio": torch.tensor([5.0]), "ref": torch.tensor([6.0])},
|
||
|
|
"ids": {"all": [1, 2], "prompt": [3, 4]},
|
||
|
|
"meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 8},
|
||
|
|
}
|
||
|
|
restored = unflatten_payload(flatten_payload(original))
|
||
|
|
assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
|
||
|
|
assert torch.equal(restored["hidden_states"]["last"], torch.tensor([2.0]))
|
||
|
|
assert torch.equal(restored["embed"]["prefill"], torch.tensor([3.0]))
|
||
|
|
assert torch.equal(restored["embed"]["tts_bos"], torch.tensor([4.0]))
|
||
|
|
assert torch.equal(restored["codes"]["audio"], torch.tensor([5.0]))
|
||
|
|
assert torch.equal(restored["codes"]["ref"], torch.tensor([6.0]))
|
||
|
|
assert restored["ids"]["all"] == [1, 2]
|
||
|
|
assert restored["ids"]["prompt"] == [3, 4]
|
||
|
|
assert restored["meta"]["finished"].item() is False
|
||
|
|
assert restored["meta"]["ar_width"] == 8
|
||
|
|
|
||
|
|
|
||
|
|
class TestSerializeDeserializePayload:
|
||
|
|
def test_tensor_round_trip(self):
|
||
|
|
original: OmniPayload = {
|
||
|
|
"hidden_states": {"output": torch.tensor([[1.0, 2.0], [3.0, 4.0]])},
|
||
|
|
}
|
||
|
|
wire = serialize_payload(original)
|
||
|
|
assert isinstance(wire, AdditionalInformationPayload)
|
||
|
|
restored = deserialize_payload(wire)
|
||
|
|
assert torch.equal(restored["hidden_states"]["output"], original["hidden_states"]["output"])
|
||
|
|
|
||
|
|
def test_list_round_trip(self):
|
||
|
|
original: OmniPayload = {
|
||
|
|
"ids": {"prompt": [10, 20, 30]},
|
||
|
|
}
|
||
|
|
wire = serialize_payload(original)
|
||
|
|
restored = deserialize_payload(wire)
|
||
|
|
assert restored["ids"]["prompt"] == [10, 20, 30]
|
||
|
|
|
||
|
|
def test_finished_tensor_round_trip(self):
|
||
|
|
original: OmniPayload = {
|
||
|
|
"meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 5},
|
||
|
|
}
|
||
|
|
wire = serialize_payload(original)
|
||
|
|
restored = deserialize_payload(wire)
|
||
|
|
assert isinstance(restored["meta"]["finished"], torch.Tensor)
|
||
|
|
assert restored["meta"]["finished"].dtype == torch.bool
|
||
|
|
assert restored["meta"]["finished"].item() is True
|
||
|
|
assert restored["meta"]["left_context_size"] == 5
|
||
|
|
|
||
|
|
def test_mixed_types_round_trip(self):
|
||
|
|
original: OmniPayload = {
|
||
|
|
"hidden_states": {"output": torch.tensor([1.0, 2.0])},
|
||
|
|
"ids": {"all": [1, 2, 3]},
|
||
|
|
"meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 4},
|
||
|
|
"codes": {"audio": torch.tensor([3.0])},
|
||
|
|
}
|
||
|
|
wire = serialize_payload(original)
|
||
|
|
restored = deserialize_payload(wire)
|
||
|
|
assert torch.equal(restored["hidden_states"]["output"], original["hidden_states"]["output"])
|
||
|
|
assert restored["ids"]["all"] == [1, 2, 3]
|
||
|
|
assert restored["meta"]["finished"].item() is False
|
||
|
|
assert restored["meta"]["ar_width"] == 4
|
||
|
|
assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
|
||
|
|
|
||
|
|
def test_hidden_states_layers_round_trip(self):
|
||
|
|
original = {
|
||
|
|
"hidden_states": {
|
||
|
|
"output": torch.tensor([1.0]),
|
||
|
|
"layers": {0: torch.tensor([2.0]), 24: torch.tensor([3.0])},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
wire = serialize_payload(original)
|
||
|
|
restored = deserialize_payload(wire)
|
||
|
|
assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
|
||
|
|
assert torch.equal(restored["hidden_states"]["layers"][0], torch.tensor([2.0]))
|
||
|
|
assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
|
||
|
|
|
||
|
|
def test_tensor_dtype_preserved(self):
|
||
|
|
# bfloat16 excluded: numpy() doesn't support it; callers must cast before serializing.
|
||
|
|
for dtype in [torch.float16, torch.float32, torch.int64, torch.int32, torch.bool]:
|
||
|
|
original: OmniPayload = {"codes": {"audio": torch.tensor([1], dtype=dtype)}}
|
||
|
|
wire = serialize_payload(original)
|
||
|
|
restored = deserialize_payload(wire)
|
||
|
|
assert restored["codes"]["audio"].dtype == dtype, f"dtype mismatch for {dtype}"
|
||
|
|
|
||
|
|
def test_tensor_shape_preserved(self):
|
||
|
|
t = torch.randn(3, 4, 5)
|
||
|
|
original: OmniPayload = {"hidden_states": {"output": t}}
|
||
|
|
wire = serialize_payload(original)
|
||
|
|
restored = deserialize_payload(wire)
|
||
|
|
assert restored["hidden_states"]["output"].shape == (3, 4, 5)
|
||
|
|
assert torch.allclose(restored["hidden_states"]["output"], t)
|
||
|
|
|
||
|
|
def test_empty_payload_returns_none(self):
|
||
|
|
assert serialize_payload({}) is None
|
||
|
|
|
||
|
|
def test_none_values_skipped(self):
|
||
|
|
original: OmniPayload = {"meta": {"finished": None}}
|
||
|
|
wire = serialize_payload(original)
|
||
|
|
assert wire is None
|