609 lines
21 KiB
Python
609 lines
21 KiB
Python
"""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
|
|
DEFAULT_MAX_FRAGMENTS_PER_TENSOR = 64
|
|
DEFAULT_MAX_TENSORS_PER_BUNDLE = 64
|
|
DEFAULT_MAX_TENSOR_RANK = 8
|
|
DEFAULT_MAX_TENSOR_DIMENSION = (1 << 31) - 1
|
|
|
|
# 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,
|
|
*,
|
|
max_rank: int = DEFAULT_MAX_TENSOR_RANK,
|
|
max_dimension: int = DEFAULT_MAX_TENSOR_DIMENSION,
|
|
max_bytes: int | None = None,
|
|
) -> int:
|
|
"""Byte count a tensor of `shape` and `dtype` must occupy."""
|
|
if len(shape) > max_rank:
|
|
raise ProtocolError(f"tensor rank {len(shape)} exceeds limit {max_rank}")
|
|
if any(dim < 0 or dim > max_dimension for dim in shape):
|
|
raise ProtocolError(
|
|
f"dimension outside 0..{max_dimension} in shape {list(shape)}"
|
|
)
|
|
size = itemsize(dtype)
|
|
count = 1
|
|
for dim in shape:
|
|
count *= dim
|
|
if max_bytes is not None and count * size > max_bytes:
|
|
raise ProtocolError(f"tensor shape {list(shape)} exceeds byte limit {max_bytes}")
|
|
return count * size
|
|
|
|
|
|
# --- 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_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
|
|
max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
|
|
max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
|
|
) -> 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_chunk_bytes <= 0 or max_fragment_bytes <= 0 or max_fragments <= 0:
|
|
raise ProtocolError("tensor byte/count bounds must be positive")
|
|
declared = expected_bytes(shape, dtype, max_bytes=max_chunk_bytes)
|
|
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""]
|
|
if len(slices) > max_fragments:
|
|
raise ProtocolError(
|
|
f"tensor {name!r} needs {len(slices)} fragments, exceeding limit {max_fragments}"
|
|
)
|
|
|
|
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,
|
|
*,
|
|
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
|
|
max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
|
|
max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
|
|
) -> 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 max_chunk_bytes <= 0 or max_fragment_bytes <= 0 or max_fragments <= 0:
|
|
raise ProtocolError("negotiated byte/count bounds must be positive")
|
|
if tensor.total_bytes > max_chunk_bytes:
|
|
raise ProtocolError(
|
|
f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes, exceeding "
|
|
f"the {max_chunk_bytes}-byte negotiated chunk bound"
|
|
)
|
|
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")
|
|
declared = expected_bytes(
|
|
tensor.shape, tensor.dtype, max_bytes=max_chunk_bytes
|
|
)
|
|
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}"
|
|
)
|
|
|
|
if not tensor.fragments:
|
|
raise PayloadCorrupt(f"tensor {tensor.name!r} carries no fragments")
|
|
if len(tensor.fragments) > max_fragments:
|
|
raise PayloadCorrupt(
|
|
f"tensor {tensor.name!r} carries {len(tensor.fragments)} fragments, "
|
|
f"exceeding limit {max_fragments}"
|
|
)
|
|
if any(len(fragment.payload) > max_fragment_bytes for fragment in tensor.fragments):
|
|
raise PayloadCorrupt(
|
|
f"tensor {tensor.name!r} carries a fragment larger than "
|
|
f"{max_fragment_bytes} bytes"
|
|
)
|
|
wire_bytes = sum(len(fragment.payload) for fragment in tensor.fragments)
|
|
if wire_bytes > max_chunk_bytes:
|
|
raise PayloadCorrupt(
|
|
f"tensor {tensor.name!r} wire body exceeds the "
|
|
f"{max_chunk_bytes}-byte negotiated chunk bound"
|
|
)
|
|
|
|
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:
|
|
try:
|
|
data = decompress_activation(
|
|
bytes(body), "zstd", max_output_bytes=tensor.total_bytes
|
|
).body
|
|
except ValueError as exc:
|
|
raise PayloadCorrupt(
|
|
f"tensor {tensor.name!r} has invalid bounded zstd payload"
|
|
) from exc
|
|
elif tensor.compression == pb.COMPRESSION_NONE:
|
|
data = bytes(body)
|
|
else:
|
|
raise ProtocolError(
|
|
f"tensor {tensor.name!r} uses unspecified or 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)}"
|
|
)
|
|
|
|
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 != pb.CHECKSUM_ALGORITHM_NONE:
|
|
raise ProtocolError(
|
|
f"tensor {tensor.name!r} uses unspecified or unsupported checksum algorithm"
|
|
)
|
|
|
|
return data
|
|
|
|
|
|
def encode_bundle(
|
|
tensors: Iterable[pb.NamedTensor],
|
|
*,
|
|
architecture: int = pb.ARCHITECTURE_TYPE_UNSPECIFIED,
|
|
boundary_point: str = "",
|
|
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
|
|
max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
|
|
) -> pb.TensorBundle:
|
|
if max_chunk_bytes <= 0 or max_tensors <= 0:
|
|
raise ProtocolError("bundle byte/count bounds must be positive")
|
|
tensor_list = list(tensors)
|
|
if len(tensor_list) > max_tensors:
|
|
raise ProtocolError(
|
|
f"bundle carries {len(tensor_list)} tensors, exceeding limit {max_tensors}"
|
|
)
|
|
bundle = pb.TensorBundle(
|
|
bundle_version=BUNDLE_VERSION,
|
|
tensors=tensor_list,
|
|
architecture=architecture,
|
|
boundary_point=boundary_point,
|
|
)
|
|
if bundle.ByteSize() > max_chunk_bytes:
|
|
raise ProtocolError(
|
|
f"serialized tensor bundle exceeds the {max_chunk_bytes}-byte "
|
|
"negotiated chunk bound"
|
|
)
|
|
return bundle
|
|
|
|
|
|
def decode_bundle(
|
|
bundle: pb.TensorBundle,
|
|
*,
|
|
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
|
|
max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
|
|
max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
|
|
max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
|
|
) -> 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 not supported by this build "
|
|
f"({BUNDLE_VERSION})"
|
|
)
|
|
if (
|
|
max_chunk_bytes <= 0
|
|
or max_fragment_bytes <= 0
|
|
or max_fragments <= 0
|
|
or max_tensors <= 0
|
|
):
|
|
raise ProtocolError("negotiated byte/count bounds must be positive")
|
|
if len(bundle.tensors) > max_tensors:
|
|
raise ProtocolError(
|
|
f"bundle carries {len(bundle.tensors)} tensors, exceeding limit {max_tensors}"
|
|
)
|
|
if bundle.ByteSize() > max_chunk_bytes:
|
|
raise ProtocolError(
|
|
f"serialized tensor bundle exceeds the {max_chunk_bytes}-byte "
|
|
"negotiated chunk bound"
|
|
)
|
|
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,
|
|
max_chunk_bytes=max_chunk_bytes,
|
|
max_fragment_bytes=max_fragment_bytes,
|
|
max_fragments=max_fragments,
|
|
)
|
|
return payloads
|
|
|
|
|
|
def encode_decode_step(
|
|
bundle: pb.TensorBundle,
|
|
*,
|
|
idempotency_step: int,
|
|
position: int,
|
|
expected_past_len: int,
|
|
work_id: str,
|
|
deadline_unix_nanos: int = 0,
|
|
prefer_compact_one_tensor: bool = True,
|
|
) -> pb.DecodeStep:
|
|
"""Encode a decode boundary, retaining the deliberate compact fallback."""
|
|
step = pb.DecodeStep(
|
|
idempotency_step=idempotency_step,
|
|
position=position,
|
|
expected_past_len=expected_past_len,
|
|
work_id=work_id,
|
|
deadline_unix_nanos=deadline_unix_nanos,
|
|
)
|
|
if prefer_compact_one_tensor and len(bundle.tensors) == 1:
|
|
step.tensor.CopyFrom(bundle.tensors[0])
|
|
else:
|
|
step.bundle.CopyFrom(bundle)
|
|
return step
|
|
|
|
|
|
def validate_tail_result(result: pb.TailResult) -> None:
|
|
"""Fail closed unless a tail completion is bound to its exact recipe."""
|
|
identity = result.identity
|
|
required = (
|
|
identity.request_id,
|
|
identity.runtime_recipe_digest,
|
|
identity.chat_template_id,
|
|
identity.chat_template_version,
|
|
identity.reasoning_mode,
|
|
)
|
|
if not all(required) or identity.architecture == pb.ARCHITECTURE_TYPE_UNSPECIFIED:
|
|
raise ProtocolError("tail result lacks exact request/recipe/template identity")
|
|
if result.WhichOneof("output") not in {"logits", "sampled_token_id"}:
|
|
raise ProtocolError("tail result lacks logits or sampled token output")
|
|
|
|
|
|
def decode_step_bundle(
|
|
step: pb.DecodeStep,
|
|
*,
|
|
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
|
|
max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
|
|
max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
|
|
max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
|
|
) -> dict[str, bytes]:
|
|
"""Decode a fast-path boundary with the DGR-006 compatibility rule.
|
|
|
|
`bundle` is authoritative because it can carry architecture sidebands. The
|
|
old `tensor` field remains the compact representation for a certified
|
|
one-tensor boundary and is accepted by new readers during rollout.
|
|
"""
|
|
if step.HasField("bundle"):
|
|
return decode_bundle(
|
|
step.bundle,
|
|
max_chunk_bytes=max_chunk_bytes,
|
|
max_fragment_bytes=max_fragment_bytes,
|
|
max_fragments=max_fragments,
|
|
max_tensors=max_tensors,
|
|
)
|
|
if step.HasField("tensor"):
|
|
return decode_bundle(
|
|
encode_bundle([step.tensor], max_chunk_bytes=max_chunk_bytes),
|
|
max_chunk_bytes=max_chunk_bytes,
|
|
max_fragment_bytes=max_fragment_bytes,
|
|
max_fragments=max_fragments,
|
|
max_tensors=max_tensors,
|
|
)
|
|
raise ProtocolError("decode step carries neither TensorBundle nor legacy tensor")
|
|
|
|
|
|
def validate_session_message_size(
|
|
message: pb.SessionRequest | pb.SessionResponse,
|
|
*,
|
|
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
|
|
) -> int:
|
|
"""Reject an oversized complete stream frame, including protobuf overhead.
|
|
|
|
Bundle validation alone is insufficient because the envelope and oneof
|
|
framing are part of the same gRPC message. Senders call this immediately
|
|
before writing; receivers configure gRPC's receive limit to the same value
|
|
and call it again before semantic decoding.
|
|
"""
|
|
if max_chunk_bytes <= 0:
|
|
raise ProtocolError("max_chunk_bytes must be positive")
|
|
if not isinstance(message, (pb.SessionRequest, pb.SessionResponse)):
|
|
raise ProtocolError("size validation requires a session request or response")
|
|
size = message.ByteSize()
|
|
if size > max_chunk_bytes:
|
|
raise ProtocolError(
|
|
f"serialized session message is {size} bytes, exceeding the "
|
|
f"{max_chunk_bytes}-byte negotiated chunk bound"
|
|
)
|
|
return size
|
|
|
|
|
|
# --- 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
|
|
|
|
max_inflight_chunks = _min(
|
|
proposed.max_inflight_chunks,
|
|
limits.max_inflight_chunks,
|
|
DEFAULT_MAX_INFLIGHT_CHUNKS,
|
|
)
|
|
credits_granted = min(
|
|
_min(
|
|
proposed.credits_granted,
|
|
limits.credits_granted,
|
|
DEFAULT_MAX_INFLIGHT_CHUNKS,
|
|
),
|
|
max_inflight_chunks,
|
|
)
|
|
return pb.FlowControl(
|
|
credits_granted=credits_granted,
|
|
max_inflight_chunks=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,
|
|
),
|
|
)
|