190 lines
7.9 KiB
Python
190 lines
7.9 KiB
Python
"""Lossless, bounded CPU payloads shared with the standalone endpoint bridge.
|
|
|
|
Descriptors contain only kind-specific metadata. Tensor bytes are contiguous,
|
|
little-endian floating-point values, never a pickle or an image/audio container.
|
|
The optional byte budget can lower, but never raise, MAX_PAYLOAD_BYTES.
|
|
"""
|
|
|
|
import ctypes
|
|
import sys
|
|
|
|
import torch
|
|
|
|
|
|
MAX_PAYLOAD_BYTES = 256 * 1024 * 1024
|
|
|
|
_DTYPES = {
|
|
"float16": (torch.float16, 2),
|
|
"bfloat16": (torch.bfloat16, 2),
|
|
"float32": (torch.float32, 4),
|
|
"float64": (torch.float64, 8),
|
|
}
|
|
_DTYPE_NAMES = {dtype: name for name, (dtype, _) in _DTYPES.items()}
|
|
_LAYOUTS = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}
|
|
|
|
|
|
def _validate_kind(kind):
|
|
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO"):
|
|
raise ValueError("Unsupported payload kind")
|
|
|
|
|
|
def _byte_limit(max_bytes):
|
|
if type(max_bytes) is not int or max_bytes < 0:
|
|
raise ValueError("max_bytes must be a nonnegative integer")
|
|
return min(max_bytes, MAX_PAYLOAD_BYTES)
|
|
|
|
|
|
def _sample_rate(value):
|
|
if type(value) is not int or value <= 0:
|
|
raise ValueError("AUDIO sample_rate must be a positive integer")
|
|
return value
|
|
|
|
|
|
def _tensor_size(kind, shape, itemsize, limit):
|
|
if type(shape) is not list or len(shape) != len(_LAYOUTS[kind]):
|
|
raise ValueError("Tensor shape has an invalid rank")
|
|
if any(type(dimension) is not int or dimension <= 0 for dimension in shape):
|
|
raise ValueError("Tensor shape dimensions must be positive integers")
|
|
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
|
|
raise ValueError("IMAGE shape must have 1, 3 or 4 channels")
|
|
size = itemsize
|
|
for dimension in shape:
|
|
# Divide before multiplying, including for arbitrary-size Python ints.
|
|
if dimension > limit // size:
|
|
raise ValueError("Tensor payload exceeds byte limit")
|
|
size *= dimension
|
|
return size
|
|
|
|
|
|
def _native_little_endian(payload, itemsize):
|
|
"""Convert native to wire byte order, or vice versa (an involution)."""
|
|
if sys.byteorder == "little":
|
|
return payload
|
|
swapped = bytearray(len(payload))
|
|
for offset in range(itemsize):
|
|
swapped[offset::itemsize] = payload[itemsize - 1 - offset::itemsize]
|
|
return bytes(swapped)
|
|
|
|
|
|
def _encode_string(value, limit):
|
|
if not isinstance(value, str):
|
|
raise ValueError("STRING value must be text")
|
|
# Check the exact UTF-8 size before allocating the encoded payload.
|
|
if len(value) > limit:
|
|
raise ValueError("STRING payload exceeds byte limit")
|
|
size = 0
|
|
for character in value:
|
|
codepoint = ord(character)
|
|
if 0xD800 <= codepoint <= 0xDFFF:
|
|
raise ValueError("STRING value must be valid UTF-8 text")
|
|
size += 1 + (codepoint > 0x7F) + (codepoint > 0x7FF) + (codepoint > 0xFFFF)
|
|
if size > limit:
|
|
raise ValueError("STRING payload exceeds byte limit")
|
|
return {"encoding": "utf8-v1"}, value.encode("utf-8")
|
|
|
|
|
|
def encode_payload(kind: str, value, *, max_bytes: int = MAX_PAYLOAD_BYTES) -> tuple[dict, bytes]:
|
|
"""Encode a typed value without dtype conversion, normalization or GPU use.
|
|
|
|
IMAGE/MASK accept dense CPU tensors; AUDIO accepts a waveform CPU tensor
|
|
plus a positive integer sample_rate. STRING accepts exact Unicode text.
|
|
Invalid values and exceeded budgets raise nonsensitive ValueError messages.
|
|
"""
|
|
_validate_kind(kind)
|
|
limit = _byte_limit(max_bytes)
|
|
if kind == "STRING":
|
|
return _encode_string(value, limit)
|
|
|
|
sample_rate = None
|
|
if kind == "AUDIO":
|
|
if not isinstance(value, dict) or set(value) != {"waveform", "sample_rate"}:
|
|
raise ValueError("AUDIO value must contain waveform and sample_rate")
|
|
sample_rate = _sample_rate(value["sample_rate"])
|
|
value = value["waveform"]
|
|
if not isinstance(value, torch.Tensor):
|
|
raise ValueError("Tensor value must be a torch.Tensor")
|
|
if value.device.type != "cpu":
|
|
raise ValueError("Tensor value must be on CPU")
|
|
if value.layout != torch.strided or value.is_nested:
|
|
raise ValueError("Tensor value must have a dense strided layout")
|
|
if value.dtype not in _DTYPE_NAMES:
|
|
raise ValueError("Unsupported tensor dtype")
|
|
|
|
dtype_name = _DTYPE_NAMES[value.dtype]
|
|
_, itemsize = _DTYPES[dtype_name]
|
|
shape = list(value.shape)
|
|
size = _tensor_size(kind, shape, itemsize, limit)
|
|
tensor = value.detach().resolve_neg().contiguous()
|
|
if not torch.isfinite(tensor).all().item():
|
|
raise ValueError("Tensor values must be finite")
|
|
# A uint16 view preserves BF16 bits without conversion through float32.
|
|
if tensor.dtype == torch.bfloat16:
|
|
tensor = tensor.view(torch.uint16)
|
|
# The live, contiguous CPU tensor owns this pointer for the entire copy.
|
|
payload = _native_little_endian(ctypes.string_at(tensor.data_ptr(), size), itemsize)
|
|
descriptor = {
|
|
"encoding": "tensor-le-v1", "dtype": dtype_name,
|
|
"shape": shape, "layout": _LAYOUTS[kind],
|
|
}
|
|
if kind == "AUDIO":
|
|
descriptor["sample_rate"] = sample_rate
|
|
return descriptor, payload
|
|
|
|
|
|
def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: int = MAX_PAYLOAD_BYTES):
|
|
"""Decode validated wire bytes into independent contiguous CPU tensors/text.
|
|
|
|
All metadata, shape products and byte lengths are checked before tensor
|
|
allocation. No GPU device discovery, transfers or object deserialization
|
|
occur. AUDIO returns {'waveform': tensor, 'sample_rate': int}.
|
|
"""
|
|
_validate_kind(kind)
|
|
limit = _byte_limit(max_bytes)
|
|
if not isinstance(descriptor, dict):
|
|
raise ValueError("Payload descriptor must be a dictionary")
|
|
if not isinstance(payload, bytes):
|
|
raise ValueError("Payload must be bytes")
|
|
if len(payload) > limit:
|
|
raise ValueError("Payload exceeds byte limit")
|
|
if kind == "IMAGE" and descriptor.get("encoding") == "image-transfer-v1":
|
|
# Lazy and IMAGE-only: exact tensors and other kinds stay unchanged.
|
|
from .image_codec import decode_image
|
|
return decode_image(descriptor, payload, max_bytes=limit)
|
|
if kind == "STRING":
|
|
if set(descriptor) != {"encoding"} or descriptor["encoding"] != "utf8-v1":
|
|
raise ValueError("Invalid STRING descriptor encoding or fields")
|
|
try:
|
|
return payload.decode("utf-8")
|
|
except UnicodeDecodeError:
|
|
raise ValueError("STRING payload must be valid UTF-8 bytes") from None
|
|
|
|
fields = {"encoding", "dtype", "shape", "layout"}
|
|
if kind == "AUDIO":
|
|
fields.add("sample_rate")
|
|
if set(descriptor) != fields:
|
|
raise ValueError("Invalid tensor descriptor fields")
|
|
if descriptor["encoding"] != "tensor-le-v1":
|
|
raise ValueError("Unsupported tensor descriptor encoding")
|
|
if descriptor["layout"] != _LAYOUTS[kind]:
|
|
raise ValueError("Invalid tensor descriptor layout")
|
|
dtype_name = descriptor["dtype"]
|
|
if type(dtype_name) is not str or dtype_name not in _DTYPES:
|
|
raise ValueError("Unsupported tensor descriptor dtype")
|
|
dtype, itemsize = _DTYPES[dtype_name]
|
|
size = _tensor_size(kind, descriptor["shape"], itemsize, limit)
|
|
if len(payload) != size:
|
|
raise ValueError("Tensor payload length does not match shape and dtype")
|
|
sample_rate = _sample_rate(descriptor["sample_rate"]) if kind == "AUDIO" else None
|
|
|
|
# frombuffer is always CPU. Copy immutable wire bytes into a writable private
|
|
# buffer, then clone so the returned tensor owns its storage independently.
|
|
buffer = bytearray(_native_little_endian(payload, itemsize))
|
|
storage_dtype = torch.uint16 if dtype == torch.bfloat16 else dtype
|
|
tensor = torch.frombuffer(buffer, dtype=storage_dtype).view(dtype)
|
|
if not torch.isfinite(tensor).all().item():
|
|
raise ValueError("Tensor payload values must be finite")
|
|
tensor = tensor.clone().reshape(descriptor["shape"])
|
|
if kind == "AUDIO":
|
|
return {"waveform": tensor, "sample_rate": sample_rate}
|
|
return tensor
|