feat: DGR-002 - Adopt the versioned gRPC Shard protocol
This commit is contained in:
Dobromir Popov
383
packages/node/meshnet_node/native_protocol/codec.py
Normal file
383
packages/node/meshnet_node/native_protocol/codec.py
Normal file
@@ -0,0 +1,383 @@
|
||||
"""Encode and decode the native Shard protocol's named-tensor bundles.
|
||||
|
||||
The generated stubs give us message *structure*; they cannot enforce the
|
||||
invariants that keep a distributed forward correct. A bundle whose declared
|
||||
shape disagrees with its byte count, whose fragments leave a hole, or whose
|
||||
checksum does not match is not a slightly-wrong activation — it is silently
|
||||
wrong tokens for the rest of the generation. So decoding is validating: every
|
||||
path into a tensor's bytes goes through :func:`decode_tensor`, which refuses a
|
||||
payload it cannot fully account for.
|
||||
|
||||
Compression is a transport optimisation and is decided by the same policy layer
|
||||
the existing HTTP seam already uses (``activation_compression``), so a node's
|
||||
tuned thresholds apply to both transports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import struct
|
||||
from typing import Iterable, Sequence
|
||||
|
||||
from ..activation_compression import (
|
||||
CompressionPolicy,
|
||||
compress_activation,
|
||||
decompress_activation,
|
||||
)
|
||||
from .generated import shard_runtime_pb2 as pb
|
||||
|
||||
# The schema generation this build speaks. A peer offering something else is
|
||||
# rejected at the handshake rather than being half-understood.
|
||||
SCHEMA_VERSION = pb.SCHEMA_VERSION_1
|
||||
|
||||
# Generation of the tensor-bundle layout, versioned independently of the
|
||||
# protocol so a boundary payload can evolve without a protocol bump.
|
||||
BUNDLE_VERSION = 1
|
||||
|
||||
# Token-aligned prefill chunk bound. 128 tokens is the size ADR-0008 already
|
||||
# uses on the HTTP seam; keeping it identical means seam bytes stay comparable
|
||||
# across transports.
|
||||
DEFAULT_MAX_PREFILL_CHUNK_TOKENS = 128
|
||||
|
||||
# gRPC's default maximum receive size. Fragmenting below it keeps us inside the
|
||||
# default limits of any conformant peer instead of requiring every client to
|
||||
# raise its window.
|
||||
DEFAULT_MAX_CHUNK_BYTES = 4 * 1024 * 1024
|
||||
|
||||
# Leave room for envelope and framing overhead inside one chunk message.
|
||||
DEFAULT_MAX_FRAGMENT_BYTES = 1024 * 1024
|
||||
|
||||
DEFAULT_MAX_INFLIGHT_CHUNKS = 8
|
||||
|
||||
# Canonical boundary tensor name for a dense transformer hidden state.
|
||||
HIDDEN_STATES = "hidden_states"
|
||||
|
||||
_DTYPE_ITEMSIZE: dict[int, int] = {
|
||||
pb.DTYPE_BFLOAT16: 2,
|
||||
pb.DTYPE_FLOAT16: 2,
|
||||
pb.DTYPE_FLOAT32: 4,
|
||||
pb.DTYPE_INT32: 4,
|
||||
pb.DTYPE_INT64: 8,
|
||||
pb.DTYPE_UINT8: 1,
|
||||
pb.DTYPE_INT8: 1,
|
||||
pb.DTYPE_BOOL: 1,
|
||||
}
|
||||
|
||||
|
||||
class ProtocolError(Exception):
|
||||
"""A peer sent something this build cannot safely interpret."""
|
||||
|
||||
|
||||
class PayloadCorrupt(ProtocolError):
|
||||
"""A tensor payload failed validation: size, coverage, or checksum."""
|
||||
|
||||
|
||||
def itemsize(dtype: int) -> int:
|
||||
try:
|
||||
return _DTYPE_ITEMSIZE[dtype]
|
||||
except KeyError:
|
||||
raise ProtocolError(f"unsupported dtype {dtype}") from None
|
||||
|
||||
|
||||
def expected_bytes(shape: Sequence[int], dtype: int) -> int:
|
||||
"""Byte count a tensor of `shape` and `dtype` must occupy."""
|
||||
if any(dim < 0 for dim in shape):
|
||||
raise ProtocolError(f"negative dimension in shape {list(shape)}")
|
||||
count = 1
|
||||
for dim in shape:
|
||||
count *= dim
|
||||
return count * itemsize(dtype)
|
||||
|
||||
|
||||
# --- CRC32C ----------------------------------------------------------------
|
||||
#
|
||||
# CRC32C (Castagnoli), not zlib's CRC32: it is the checksum gRPC, and the
|
||||
# storage systems these payloads pass through, already use, and hardware
|
||||
# implements it. `google_crc32c` is used when present; the table fallback keeps
|
||||
# the default test suite dependency-free.
|
||||
|
||||
_CRC32C_POLY = 0x82F63B78
|
||||
_CRC32C_TABLE: list[int] = []
|
||||
for _i in range(256):
|
||||
_c = _i
|
||||
for _ in range(8):
|
||||
_c = (_c >> 1) ^ (_CRC32C_POLY if _c & 1 else 0)
|
||||
_CRC32C_TABLE.append(_c)
|
||||
|
||||
try: # pragma: no cover - depends on an optional native package
|
||||
from google_crc32c import value as _fast_crc32c
|
||||
except ImportError: # pragma: no cover
|
||||
_fast_crc32c = None
|
||||
|
||||
|
||||
def crc32c(data: bytes) -> int:
|
||||
if _fast_crc32c is not None: # pragma: no cover - optional fast path
|
||||
return _fast_crc32c(data)
|
||||
crc = 0xFFFFFFFF
|
||||
for byte in data:
|
||||
crc = (crc >> 8) ^ _CRC32C_TABLE[(crc ^ byte) & 0xFF]
|
||||
return crc ^ 0xFFFFFFFF
|
||||
|
||||
|
||||
def checksum_of(data: bytes) -> pb.Checksum:
|
||||
return pb.Checksum(
|
||||
algorithm=pb.CHECKSUM_ALGORITHM_CRC32C,
|
||||
value=struct.pack(">I", crc32c(data)),
|
||||
)
|
||||
|
||||
|
||||
# --- Tensors ---------------------------------------------------------------
|
||||
|
||||
|
||||
def encode_tensor(
|
||||
name: str,
|
||||
data: bytes,
|
||||
shape: Sequence[int],
|
||||
dtype: int = pb.DTYPE_BFLOAT16,
|
||||
*,
|
||||
policy: CompressionPolicy | None = None,
|
||||
max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
|
||||
) -> pb.NamedTensor:
|
||||
"""Build a NamedTensor, compressing and fragmenting as needed.
|
||||
|
||||
`data` is the uncompressed little-endian payload. The checksum is taken over
|
||||
it *before* compression so it stays valid whichever framing a hop chooses.
|
||||
"""
|
||||
if max_fragment_bytes <= 0:
|
||||
raise ProtocolError("max_fragment_bytes must be positive")
|
||||
declared = expected_bytes(shape, dtype)
|
||||
if len(data) != declared:
|
||||
raise ProtocolError(
|
||||
f"tensor {name!r} declares shape {list(shape)} ({declared} bytes) "
|
||||
f"but carries {len(data)} bytes"
|
||||
)
|
||||
|
||||
body = data
|
||||
compression = pb.COMPRESSION_NONE
|
||||
if policy is not None:
|
||||
result = compress_activation(data, policy)
|
||||
if result.compressed:
|
||||
body = result.body
|
||||
compression = pb.COMPRESSION_ZSTD
|
||||
|
||||
tensor = pb.NamedTensor(
|
||||
name=name,
|
||||
shape=list(shape),
|
||||
dtype=dtype,
|
||||
byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN,
|
||||
total_bytes=len(data),
|
||||
compression=compression,
|
||||
checksum=checksum_of(data),
|
||||
)
|
||||
|
||||
# Fragment the wire body (compressed if we compressed). Offsets walk the
|
||||
# wire body so a receiver can verify coverage without assuming arrival
|
||||
# order; a zstd frame is not decodable per fragment, so reassembly comes
|
||||
# first and decompression happens once, in decode_tensor.
|
||||
slices = [body[i : i + max_fragment_bytes] for i in range(0, len(body), max_fragment_bytes)]
|
||||
if not slices:
|
||||
# A zero-element tensor is legal (e.g. an empty mask) and still needs a
|
||||
# fragment, so coverage checks have something to verify.
|
||||
slices = [b""]
|
||||
|
||||
offset = 0
|
||||
for index, piece in enumerate(slices):
|
||||
tensor.fragments.append(
|
||||
pb.TensorFragment(
|
||||
fragment_index=index,
|
||||
fragment_count=len(slices),
|
||||
byte_offset=offset,
|
||||
payload=piece,
|
||||
)
|
||||
)
|
||||
offset += len(piece)
|
||||
return tensor
|
||||
|
||||
|
||||
def decode_tensor(tensor: pb.NamedTensor) -> bytes:
|
||||
"""Reassemble, decompress and validate a NamedTensor's payload.
|
||||
|
||||
Raises PayloadCorrupt rather than returning a payload it cannot fully
|
||||
account for: a hole in the fragments or a bad checksum means the activation
|
||||
is not what the sender computed, and continuing would corrupt the route.
|
||||
"""
|
||||
if tensor.byte_order == pb.BYTE_ORDER_BIG_ENDIAN:
|
||||
raise ProtocolError(f"tensor {tensor.name!r} is big-endian; wire order is little-endian")
|
||||
if tensor.byte_order != pb.BYTE_ORDER_LITTLE_ENDIAN:
|
||||
raise ProtocolError(f"tensor {tensor.name!r} declares no byte order")
|
||||
|
||||
if not tensor.fragments:
|
||||
raise PayloadCorrupt(f"tensor {tensor.name!r} carries no fragments")
|
||||
|
||||
fragments = sorted(tensor.fragments, key=lambda f: f.byte_offset)
|
||||
count = fragments[0].fragment_count
|
||||
if any(f.fragment_count != count for f in fragments):
|
||||
raise PayloadCorrupt(f"tensor {tensor.name!r} has inconsistent fragment_count")
|
||||
if len(fragments) != count:
|
||||
raise PayloadCorrupt(
|
||||
f"tensor {tensor.name!r} expects {count} fragments but carries {len(fragments)}"
|
||||
)
|
||||
if {f.fragment_index for f in fragments} != set(range(count)):
|
||||
raise PayloadCorrupt(f"tensor {tensor.name!r} has duplicate or missing fragment indices")
|
||||
|
||||
# Contiguity: offsets must tile the body exactly, with no hole and no overlap.
|
||||
body = bytearray()
|
||||
for fragment in fragments:
|
||||
if fragment.byte_offset != len(body):
|
||||
raise PayloadCorrupt(
|
||||
f"tensor {tensor.name!r} fragment {fragment.fragment_index} starts at "
|
||||
f"{fragment.byte_offset}, expected {len(body)}"
|
||||
)
|
||||
body.extend(fragment.payload)
|
||||
|
||||
if tensor.compression == pb.COMPRESSION_ZSTD:
|
||||
data = decompress_activation(bytes(body), "zstd").body
|
||||
elif tensor.compression in (pb.COMPRESSION_NONE, pb.COMPRESSION_UNSPECIFIED):
|
||||
data = bytes(body)
|
||||
else:
|
||||
raise ProtocolError(f"tensor {tensor.name!r} uses unsupported compression")
|
||||
|
||||
if len(data) != tensor.total_bytes:
|
||||
raise PayloadCorrupt(
|
||||
f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes but "
|
||||
f"reassembled {len(data)}"
|
||||
)
|
||||
declared = expected_bytes(tensor.shape, tensor.dtype)
|
||||
if declared != tensor.total_bytes:
|
||||
raise PayloadCorrupt(
|
||||
f"tensor {tensor.name!r} shape {list(tensor.shape)} implies {declared} bytes "
|
||||
f"but declares {tensor.total_bytes}"
|
||||
)
|
||||
|
||||
algorithm = tensor.checksum.algorithm
|
||||
if algorithm == pb.CHECKSUM_ALGORITHM_CRC32C:
|
||||
if tensor.checksum.value != struct.pack(">I", crc32c(data)):
|
||||
raise PayloadCorrupt(f"tensor {tensor.name!r} failed its CRC32C check")
|
||||
elif algorithm not in (pb.CHECKSUM_ALGORITHM_NONE, pb.CHECKSUM_ALGORITHM_UNSPECIFIED):
|
||||
raise ProtocolError(f"tensor {tensor.name!r} uses unsupported checksum algorithm")
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def encode_bundle(tensors: Iterable[pb.NamedTensor]) -> pb.TensorBundle:
|
||||
return pb.TensorBundle(bundle_version=BUNDLE_VERSION, tensors=list(tensors))
|
||||
|
||||
|
||||
def decode_bundle(bundle: pb.TensorBundle) -> dict[str, bytes]:
|
||||
"""Validate every tensor in a bundle and return name -> payload."""
|
||||
if bundle.bundle_version > BUNDLE_VERSION:
|
||||
raise ProtocolError(
|
||||
f"bundle version {bundle.bundle_version} is newer than this build "
|
||||
f"understands ({BUNDLE_VERSION})"
|
||||
)
|
||||
payloads: dict[str, bytes] = {}
|
||||
for tensor in bundle.tensors:
|
||||
if not tensor.name:
|
||||
raise ProtocolError("bundle carries an unnamed tensor")
|
||||
if tensor.name in payloads:
|
||||
raise ProtocolError(f"bundle carries duplicate tensor {tensor.name!r}")
|
||||
payloads[tensor.name] = decode_tensor(tensor)
|
||||
return payloads
|
||||
|
||||
|
||||
# --- Bounded prefill chunking ----------------------------------------------
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PrefillChunk:
|
||||
"""One token-aligned slice of a prefill."""
|
||||
|
||||
chunk_index: int
|
||||
chunk_count: int
|
||||
first_position: int
|
||||
token_count: int
|
||||
|
||||
@property
|
||||
def final_chunk(self) -> bool:
|
||||
return self.chunk_index == self.chunk_count - 1
|
||||
|
||||
def chunk_info(self) -> pb.ChunkInfo:
|
||||
return pb.ChunkInfo(
|
||||
chunk_index=self.chunk_index,
|
||||
chunk_count=self.chunk_count,
|
||||
final_chunk=self.final_chunk,
|
||||
)
|
||||
|
||||
def position(self) -> pb.PositionSpan:
|
||||
return pb.PositionSpan(
|
||||
first_position=self.first_position, token_count=self.token_count
|
||||
)
|
||||
|
||||
|
||||
def plan_prefill_chunks(
|
||||
total_tokens: int,
|
||||
*,
|
||||
first_position: int = 0,
|
||||
max_tokens: int = DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
|
||||
) -> list[PrefillChunk]:
|
||||
"""Split a prefill into bounded, token-aligned chunks.
|
||||
|
||||
Splits fall on token boundaries only (ADR-0008): a fragment of a token's
|
||||
hidden state is not a thing a receiver can execute.
|
||||
"""
|
||||
if total_tokens <= 0:
|
||||
raise ProtocolError("a prefill must carry at least one token")
|
||||
if max_tokens <= 0:
|
||||
raise ProtocolError("max_tokens must be positive")
|
||||
|
||||
count = (total_tokens + max_tokens - 1) // max_tokens
|
||||
chunks = []
|
||||
for index in range(count):
|
||||
offset = index * max_tokens
|
||||
chunks.append(
|
||||
PrefillChunk(
|
||||
chunk_index=index,
|
||||
chunk_count=count,
|
||||
first_position=first_position + offset,
|
||||
token_count=min(max_tokens, total_tokens - offset),
|
||||
)
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
def default_flow_control() -> pb.FlowControl:
|
||||
return pb.FlowControl(
|
||||
credits_granted=DEFAULT_MAX_INFLIGHT_CHUNKS,
|
||||
max_inflight_chunks=DEFAULT_MAX_INFLIGHT_CHUNKS,
|
||||
max_chunk_bytes=DEFAULT_MAX_CHUNK_BYTES,
|
||||
max_prefill_chunk_tokens=DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
|
||||
)
|
||||
|
||||
|
||||
def negotiate_flow_control(
|
||||
proposed: pb.FlowControl, limits: pb.FlowControl
|
||||
) -> pb.FlowControl:
|
||||
"""Settle a stream's limits: the strictest bound of either peer wins.
|
||||
|
||||
Taking the minimum means neither peer can raise the other's ceiling, so a
|
||||
misconfigured — or hostile — sender cannot talk a worker into unbounded
|
||||
queues by proposing a large window.
|
||||
"""
|
||||
|
||||
def _min(a: int, b: int, fallback: int) -> int:
|
||||
candidates = [v for v in (a, b) if v > 0]
|
||||
return min(candidates) if candidates else fallback
|
||||
|
||||
return pb.FlowControl(
|
||||
credits_granted=_min(
|
||||
proposed.credits_granted, limits.credits_granted, DEFAULT_MAX_INFLIGHT_CHUNKS
|
||||
),
|
||||
max_inflight_chunks=_min(
|
||||
proposed.max_inflight_chunks,
|
||||
limits.max_inflight_chunks,
|
||||
DEFAULT_MAX_INFLIGHT_CHUNKS,
|
||||
),
|
||||
max_chunk_bytes=_min(
|
||||
proposed.max_chunk_bytes, limits.max_chunk_bytes, DEFAULT_MAX_CHUNK_BYTES
|
||||
),
|
||||
max_prefill_chunk_tokens=_min(
|
||||
proposed.max_prefill_chunk_tokens,
|
||||
limits.max_prefill_chunk_tokens,
|
||||
DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
|
||||
),
|
||||
)
|
||||
Reference in New Issue
Block a user