36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
import torch
|
|
|
|
from vllm_omni.engine.serialization import deserialize_additional_information, serialize_additional_information
|
|
|
|
|
|
def test_serialize_additional_information():
|
|
info = {
|
|
"tensor_f32": torch.rand(3, dtype=torch.float32),
|
|
"tensor_f64": torch.rand(3, 3, dtype=torch.float64),
|
|
"list": [1, 2, 3],
|
|
"scalar": 1,
|
|
}
|
|
|
|
payload = serialize_additional_information(info)
|
|
|
|
assert payload.entries["tensor_f32"].tensor_dtype == "float32"
|
|
assert payload.entries["tensor_f32"].tensor_shape == [3]
|
|
assert payload.entries["tensor_f32"].tensor_data == info["tensor_f32"].numpy().tobytes()
|
|
|
|
assert payload.entries["tensor_f64"].tensor_dtype == "float64"
|
|
assert payload.entries["tensor_f64"].tensor_shape == [3, 3]
|
|
assert payload.entries["tensor_f64"].tensor_data == info["tensor_f64"].numpy().tobytes()
|
|
|
|
assert payload.entries["list"].list_data == info["list"]
|
|
assert payload.entries["scalar"].scalar_data == info["scalar"]
|
|
|
|
deserialized = deserialize_additional_information(payload)
|
|
|
|
# compare tensors separately first
|
|
assert torch.equal(deserialized["tensor_f32"], info["tensor_f32"])
|
|
assert torch.equal(deserialized["tensor_f64"], info["tensor_f64"])
|
|
del deserialized["tensor_f32"], info["tensor_f32"]
|
|
del deserialized["tensor_f64"], info["tensor_f64"]
|
|
|
|
assert deserialized == info
|