329 lines
14 KiB
Python
329 lines
14 KiB
Python
import importlib.util
|
|
import math
|
|
import struct
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
MODULE_PATH = Path(__file__).resolve().parents[1] / "serverless_bridge" / "codec.py"
|
|
spec = importlib.util.spec_from_file_location("dist_media_codec_test", MODULE_PATH)
|
|
codec = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(codec)
|
|
|
|
DTYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
|
|
SHAPES = [
|
|
("IMAGE", (2, 3, 5, 1), "BHWC"),
|
|
("IMAGE", (2, 3, 5, 3), "BHWC"),
|
|
("IMAGE", (2, 3, 5, 4), "BHWC"),
|
|
("MASK", (2, 3, 5), "BHW"),
|
|
("AUDIO", (2, 1, 17), "BCT"),
|
|
("AUDIO", (2, 6, 17), "BCT"),
|
|
]
|
|
|
|
|
|
def audio(tensor, sample_rate=48000):
|
|
return {"waveform": tensor, "sample_rate": sample_rate}
|
|
|
|
|
|
def descriptor(_kind="MASK", **changes):
|
|
result = {
|
|
"encoding": "tensor-le-v1", "dtype": "float32",
|
|
"shape": [1, 1, 1], "layout": "BHW",
|
|
}
|
|
if _kind == "IMAGE":
|
|
result.update(shape=[1, 1, 1, 3], layout="BHWC")
|
|
elif _kind == "AUDIO":
|
|
result.update(layout="BCT", sample_rate=48000)
|
|
result.update(changes)
|
|
return result
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def forbid_gpu_queries(monkeypatch):
|
|
def forbidden(*args, **kwargs):
|
|
pytest.fail("codec must not query or initialize a GPU")
|
|
for name in ("is_available", "init", "_lazy_init", "device_count"):
|
|
monkeypatch.setattr(torch.cuda, name, forbidden)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
@pytest.mark.parametrize("kind,shape,layout", SHAPES)
|
|
def test_exact_roundtrip_preserves_dtype_shape_channels_and_bits(dtype, kind, shape, layout):
|
|
count = math.prod(shape)
|
|
tensor = (torch.arange(count, device="cpu", dtype=torch.float64) / 997 - 0.125).to(dtype).reshape(shape)
|
|
tensor.flatten()[0] = -0.0
|
|
tensor.requires_grad_(True)
|
|
value = audio(tensor) if kind == "AUDIO" else tensor
|
|
|
|
metadata, payload = codec.encode_payload(kind, value)
|
|
expected = {
|
|
"encoding": "tensor-le-v1", "dtype": str(dtype).removeprefix("torch."),
|
|
"shape": list(shape), "layout": layout,
|
|
}
|
|
if kind == "AUDIO":
|
|
expected["sample_rate"] = 48000
|
|
assert metadata == expected
|
|
assert type(payload) is bytes
|
|
assert len(payload) == count * tensor.element_size()
|
|
result = codec.decode_payload(kind, metadata, payload)
|
|
if kind == "AUDIO":
|
|
assert set(result) == {"waveform", "sample_rate"}
|
|
assert result["sample_rate"] == 48000
|
|
result = result["waveform"]
|
|
assert result.device.type == "cpu"
|
|
assert result.is_contiguous()
|
|
assert not result.requires_grad
|
|
assert result.dtype == dtype
|
|
assert result.shape == shape
|
|
assert torch.equal(result.view(torch.uint8), tensor.detach().view(torch.uint8))
|
|
|
|
|
|
@pytest.mark.parametrize("dtype,fmt,expected_bits", [
|
|
(torch.float16, "e", None), (torch.bfloat16, "H", [0x3F80, 0x8000, 0xC020]),
|
|
(torch.float32, "f", None), (torch.float64, "d", None),
|
|
])
|
|
def test_wire_bytes_are_little_endian(dtype, fmt, expected_bits):
|
|
values = [1.0, -0.0, -2.5]
|
|
tensor = torch.tensor(values, dtype=dtype, device="cpu").reshape(1, 1, 3)
|
|
metadata, payload = codec.encode_payload("MASK", tensor)
|
|
expected = struct.pack("<3" + fmt, *(expected_bits if expected_bits else values))
|
|
assert payload == expected
|
|
assert torch.equal(codec.decode_payload("MASK", metadata, expected).view(torch.uint8), tensor.view(torch.uint8))
|
|
|
|
|
|
def test_bfloat16_extremes_and_subnormal_bits_are_lossless():
|
|
bits = torch.tensor([0, 0x8000, 1, 0x8001, 0x3F81, 0x7F7F, 0xFF7F], dtype=torch.uint16, device="cpu")
|
|
value = bits.view(torch.bfloat16).reshape(1, 1, -1)
|
|
metadata, payload = codec.encode_payload("MASK", value)
|
|
assert payload == struct.pack("<7H", *bits.tolist())
|
|
assert torch.equal(codec.decode_payload("MASK", metadata, payload).view(torch.uint16).flatten(), bits)
|
|
|
|
|
|
def test_fractional_mask_is_not_quantized_clipped_or_normalized():
|
|
value = torch.tensor([0.1234567890123, 1.0000000000001, -2.5], dtype=torch.float64, device="cpu").reshape(1, 1, 3)
|
|
metadata, payload = codec.encode_payload("MASK", value)
|
|
assert torch.equal(codec.decode_payload("MASK", metadata, payload), value)
|
|
|
|
|
|
@pytest.mark.parametrize("text", ["", " ", "\t\n ", "café 雪 🎵", "e\u0301", "\x00\ufeff\r\n"])
|
|
def test_strings_preserve_exact_utf8(text):
|
|
metadata, payload = codec.encode_payload("STRING", text)
|
|
assert metadata == {"encoding": "utf8-v1"}
|
|
assert payload == text.encode("utf-8")
|
|
assert codec.decode_payload("STRING", metadata, payload) == text
|
|
|
|
|
|
def test_noncontiguous_and_lazy_negative_tensors_roundtrip():
|
|
value = torch.arange(24, dtype=torch.float32, device="cpu").reshape(2, 3, 4).transpose(1, 2)
|
|
for tensor in (value, torch._neg_view(value)):
|
|
metadata, payload = codec.encode_payload("MASK", tensor)
|
|
result = codec.decode_payload("MASK", metadata, payload)
|
|
assert result.is_contiguous()
|
|
assert torch.equal(result, tensor)
|
|
|
|
|
|
def test_decode_is_independent_and_ignores_default_device(monkeypatch):
|
|
source = torch.tensor([0.25], device="cpu").reshape(1, 1, 1)
|
|
metadata, payload = codec.encode_payload("MASK", source)
|
|
buffers = []
|
|
original = torch.frombuffer
|
|
def capture(buffer, *args, **kwargs):
|
|
buffers.append(buffer)
|
|
return original(buffer, *args, **kwargs)
|
|
monkeypatch.setattr(torch, "frombuffer", capture)
|
|
with torch.device("meta"):
|
|
first = codec.decode_payload("MASK", metadata, payload)
|
|
second = codec.decode_payload("MASK", metadata, payload)
|
|
assert first.device.type == second.device.type == "cpu"
|
|
first.fill_(9)
|
|
source.fill_(7)
|
|
for buffer in buffers:
|
|
if isinstance(buffer, bytearray):
|
|
buffer[:] = b"\x00" * len(buffer)
|
|
assert second.item() == 0.25
|
|
assert codec.decode_payload("MASK", metadata, payload).item() == 0.25
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["image", " MASK", "LATENT", "", None, 1, [], {}])
|
|
def test_unknown_kinds_raise_value_error(kind):
|
|
with pytest.raises(ValueError, match="kind"):
|
|
codec.encode_payload(kind, "")
|
|
with pytest.raises(ValueError, match="kind"):
|
|
codec.decode_payload(kind, {}, b"")
|
|
|
|
|
|
@pytest.mark.parametrize("kind,value", [
|
|
("STRING", None), ("STRING", b"text"), ("STRING", 123),
|
|
("IMAGE", [[[[0.5]]]]), ("MASK", None),
|
|
("AUDIO", torch.ones(1, 1, 1, device="cpu")), ("AUDIO", {}),
|
|
("AUDIO", {"waveform": "invalid", "sample_rate": 48000}),
|
|
("AUDIO", {"waveform": torch.ones(1, 1, 1, device="cpu")}),
|
|
])
|
|
def test_unsupported_values_raise_value_error(kind, value):
|
|
with pytest.raises(ValueError):
|
|
codec.encode_payload(kind, value)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.uint8, torch.int64, torch.bool, torch.complex64])
|
|
def test_unsupported_tensor_dtype(dtype):
|
|
with pytest.raises(ValueError, match="dtype"):
|
|
codec.encode_payload("MASK", torch.ones(1, 1, 1, dtype=dtype, device="cpu"))
|
|
|
|
|
|
@pytest.mark.parametrize("kind,shape", [
|
|
("IMAGE", (1, 2, 3)), ("IMAGE", (1, 1, 1, 2)), ("IMAGE", (1, 1, 1, 5)),
|
|
("MASK", (1, 2)), ("MASK", (1, 2, 3, 1)), ("AUDIO", (2, 4)),
|
|
("IMAGE", (1, 0, 2, 3)), ("MASK", (0, 1, 1)), ("AUDIO", (1, 0, 4)),
|
|
("AUDIO", (1, 2, 0)),
|
|
])
|
|
def test_invalid_tensor_shape(kind, shape):
|
|
tensor = torch.empty(shape, dtype=torch.float32, device="cpu")
|
|
with pytest.raises(ValueError, match="shape|rank|channel|dimension"):
|
|
codec.encode_payload(kind, audio(tensor) if kind == "AUDIO" else tensor)
|
|
|
|
|
|
def test_rejects_non_cpu_and_sparse_without_transferring(monkeypatch):
|
|
meta = torch.empty((1, 1, 1), device="meta")
|
|
sparse = torch.ones((1, 1, 1), device="cpu").to_sparse()
|
|
def forbidden(*args, **kwargs):
|
|
pytest.fail("codec must not move tensors between devices")
|
|
monkeypatch.setattr(torch.Tensor, "cpu", forbidden)
|
|
monkeypatch.setattr(torch.Tensor, "to", forbidden)
|
|
for tensor in (meta, sparse):
|
|
with pytest.raises(ValueError, match="CPU|layout|strided"):
|
|
codec.encode_payload("MASK", tensor)
|
|
|
|
|
|
@pytest.mark.parametrize("sample_rate", [True, False, 0, -1, 44100.0, "44100", None])
|
|
def test_invalid_sample_rate(sample_rate):
|
|
with pytest.raises(ValueError, match="sample_rate"):
|
|
codec.encode_payload("AUDIO", audio(torch.ones(1, 1, 1, device="cpu"), sample_rate))
|
|
with pytest.raises(ValueError, match="sample_rate"):
|
|
codec.decode_payload("AUDIO", descriptor("AUDIO", sample_rate=sample_rate), b"\x00" * 4)
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [None, [], "descriptor", {}, {"encoding": "utf8-v1"}])
|
|
def test_malformed_tensor_descriptor(bad):
|
|
with pytest.raises(ValueError, match="descriptor"):
|
|
codec.decode_payload("MASK", bad, b"\x00" * 4)
|
|
|
|
|
|
@pytest.mark.parametrize("changes", [
|
|
{"encoding": "pickle"}, {"encoding": None}, {"encoding": []},
|
|
{"dtype": "uint8"}, {"dtype": torch.float32}, {"dtype": []},
|
|
{"layout": "BHWC"}, {"layout": None}, {"layout": []},
|
|
{"shape": (1, 1, 1)}, {"shape": None}, {"shape": "1,1,1"},
|
|
{"shape": [1, 1]}, {"shape": [1, 1, 1, 1]}, {"shape": [1, 0, 1]},
|
|
{"shape": [1, -1, 1]}, {"shape": [1, True, 1]}, {"shape": [1, 1.0, 1]},
|
|
{"shape": [1, "1", 1]}, {"shape": [1, [], 1]}, {"kind": "MASK"},
|
|
])
|
|
def test_invalid_tensor_descriptor_fields(changes):
|
|
with pytest.raises(ValueError):
|
|
codec.decode_payload("MASK", descriptor(**changes), b"\x00" * 4)
|
|
|
|
|
|
@pytest.mark.parametrize("metadata", [{}, {"encoding": "tensor-le-v1"}, {"encoding": "utf8-v1", "kind": "STRING"}, {"encoding": []}])
|
|
def test_invalid_string_descriptor(metadata):
|
|
with pytest.raises(ValueError):
|
|
codec.decode_payload("STRING", metadata, b"")
|
|
|
|
|
|
@pytest.mark.parametrize("payload", [b"", b"\x00" * 3, b"\x00" * 5, b"\x00" * 8])
|
|
def test_shape_length_mismatch(payload):
|
|
with pytest.raises(ValueError, match="length"):
|
|
codec.decode_payload("MASK", descriptor(), payload)
|
|
|
|
|
|
@pytest.mark.parametrize("payload", [None, "text", [0, 0, 0, 0], 4])
|
|
def test_nonbyte_payload_rejected(payload):
|
|
with pytest.raises(ValueError, match="bytes"):
|
|
codec.decode_payload("MASK", descriptor(), payload)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")])
|
|
def test_nonfinite_values_fail_on_encode_and_decode(dtype, value):
|
|
tensor = torch.tensor([value], dtype=dtype, device="cpu").reshape(1, 1, 1)
|
|
with pytest.raises(ValueError, match="finite"):
|
|
codec.encode_payload("MASK", tensor)
|
|
fmt = {torch.float16: "e", torch.float32: "f", torch.float64: "d"}
|
|
if dtype == torch.bfloat16:
|
|
payload = struct.pack("<H", 0x7FC0 if math.isnan(value) else (0x7F80 if value > 0 else 0xFF80))
|
|
else:
|
|
payload = struct.pack("<" + fmt[dtype], value)
|
|
metadata = descriptor(dtype=str(dtype).removeprefix("torch."))
|
|
with pytest.raises(ValueError, match="finite"):
|
|
codec.decode_payload("MASK", metadata, payload)
|
|
|
|
|
|
def test_invalid_utf8_errors_do_not_echo_content():
|
|
with pytest.raises(ValueError) as error:
|
|
codec.encode_payload("STRING", "private-value\ud800")
|
|
assert "private-value" not in str(error.value)
|
|
with pytest.raises(ValueError) as error:
|
|
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"private-value\xff")
|
|
assert "private-value" not in str(error.value)
|
|
|
|
|
|
@pytest.mark.parametrize("max_bytes", [True, False, -1, 4.0, "4", None])
|
|
def test_invalid_byte_budget(max_bytes):
|
|
with pytest.raises(ValueError, match="max_bytes"):
|
|
codec.encode_payload("STRING", "", max_bytes=max_bytes)
|
|
with pytest.raises(ValueError, match="max_bytes"):
|
|
codec.decode_payload("STRING", {"encoding": "utf8-v1"}, b"", max_bytes=max_bytes)
|
|
|
|
|
|
def test_exact_byte_budget_and_zero_budget_for_empty_string():
|
|
tensor = torch.ones(1, 1, 1, device="cpu")
|
|
metadata, payload = codec.encode_payload("MASK", tensor, max_bytes=4)
|
|
assert codec.decode_payload("MASK", metadata, payload, max_bytes=4).item() == 1
|
|
with pytest.raises(ValueError, match="limit|budget"):
|
|
codec.encode_payload("MASK", tensor, max_bytes=3)
|
|
with pytest.raises(ValueError, match="limit|budget"):
|
|
codec.decode_payload("MASK", metadata, payload, max_bytes=3)
|
|
metadata, payload = codec.encode_payload("STRING", "", max_bytes=0)
|
|
assert codec.decode_payload("STRING", metadata, payload, max_bytes=0) == ""
|
|
for text in ("a", "é", "🎵"):
|
|
byte_count = len(text.encode("utf-8"))
|
|
metadata, payload = codec.encode_payload("STRING", text, max_bytes=byte_count)
|
|
assert codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count) == text
|
|
with pytest.raises(ValueError, match="limit|budget"):
|
|
codec.encode_payload("STRING", text, max_bytes=byte_count - 1)
|
|
with pytest.raises(ValueError, match="limit|budget"):
|
|
codec.decode_payload("STRING", metadata, payload, max_bytes=byte_count - 1)
|
|
|
|
|
|
def test_oversized_shape_rejected_before_tensor_allocation(monkeypatch):
|
|
assert codec.MAX_PAYLOAD_BYTES == 256 * 1024 * 1024
|
|
def forbidden(*args, **kwargs):
|
|
pytest.fail("oversized shape must fail before tensor allocation")
|
|
for name in ("frombuffer", "empty", "tensor", "zeros", "isfinite"):
|
|
monkeypatch.setattr(torch, name, forbidden)
|
|
for shape in ([1, 1, codec.MAX_PAYLOAD_BYTES], [2**100, 2**100, 1]):
|
|
with pytest.raises(ValueError, match="limit|budget"):
|
|
codec.decode_payload("MASK", descriptor(shape=shape), b"\x00" * 4)
|
|
|
|
|
|
def test_oversized_encode_rejected_before_materialization(monkeypatch):
|
|
huge_view = torch.ones(1, device="cpu").expand(1, 1, codec.MAX_PAYLOAD_BYTES)
|
|
def forbidden(*args, **kwargs):
|
|
pytest.fail("oversized tensor must fail before materialization or scanning")
|
|
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
|
|
monkeypatch.setattr(torch, "isfinite", forbidden)
|
|
with pytest.raises(ValueError, match="limit|budget"):
|
|
codec.encode_payload("MASK", huge_view)
|
|
|
|
|
|
@pytest.mark.parametrize("itemsize", [2, 4, 8])
|
|
def test_native_little_endian_conversion_handles_big_endian(monkeypatch, itemsize):
|
|
raw = bytes(range(itemsize * 3))
|
|
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="little"))
|
|
assert codec._native_little_endian(raw, itemsize) == raw
|
|
monkeypatch.setattr(codec, "sys", types.SimpleNamespace(byteorder="big"))
|
|
expected = b"".join(raw[i:i + itemsize][::-1] for i in range(0, len(raw), itemsize))
|
|
assert codec._native_little_endian(raw, itemsize) == expected
|
|
assert codec._native_little_endian(expected, itemsize) == raw
|