popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: harden DGR-002 protocol bounds

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-13 17:30:54 +03:00
parent 30dcf953fe
commit d904c40f66
10 changed files with 470 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
.scratch/distributed-gguf-runtime/evidence/DGR-002/README.md
View File

@@ -38,7 +38,7 @@ New:
| Path | What |
|---|---|
| `packages/node/native/proto/shard_runtime.proto` | The schema (sha256 `077ee349…`, see `protocol.json`) |
| `packages/node/native/proto/shard_runtime.proto` | The schema (sha256 `9e211660…`, see `protocol.json`) |
| `packages/node/native/CMakeLists.txt` | C++ generation + build wiring + ctest |
| `packages/node/native/tests/test_shard_protocol_conformance.cpp` | C++ conformance test |
| `packages/node/native/testdata/*.binpb` | Committed cross-language vectors |
@@ -50,14 +50,19 @@ New:
| `scripts/generate_native_protocol.py` | Python generation, with `--check` |
| `scripts/generate_protocol_goldens.py` | Vector generation, with `--check` |
| `scripts/bootstrap_native_toolchain.sh` | Builds protobuf C++ from source |
| `tests/test_native_shard_protocol.py` | 29 Python tests |
| `tests/test_native_shard_protocol.py` | 45 Python tests |
Modified:
- `packages/node/pyproject.toml` — added `grpcio>=1.60`, `protobuf>=5`; new `proto`
extra pinning `grpcio-tools==1.82.1`.
- `packages/node/pyproject.toml` — added runtime floors `grpcio>=1.82.1` and
`protobuf>=7.35.0`, matching the committed generated-code requirements; new
`proto` extra pinning `grpcio-tools==1.82.1`.
- `packages/node/meshnet_node/activation_compression.py` — optional bounded zstd
output for untrusted protocol frames; existing callers remain compatible.
- `packages/node/meshnet_node/native_protocol/__init__.py` — exports negotiated
bound constants and whole-session-message validation.
No other working-tree file was touched. `git status` before this story was clean.
The canonical PRD marks only DGR-002 passed. `git status` before this story was clean.
## 3. Commands and real results
@@ -76,14 +81,45 @@ ctest --test-dir build/native --output-on-failure -> 1/1 Test #1: shard_prot
cmp build/native/cpp_roundtrip.binpb \
packages/node/native/testdata/session_request_golden.binpb -> identical (exit 0)
pytest -q tests/test_native_shard_protocol.py -> 29 passed
pytest -q (full suite) -> 712 passed, 12 skipped
pytest -q tests/test_native_shard_protocol.py -> 45 passed
pytest -q tests/test_native_shard_protocol.py \
tests/test_activation_compression.py -> 51 passed
pytest -q (final full suite) -> 728 passed, 12 skipped
pytest -q tests/test_tracker_routing.py::test_tracker_dashboard_can_cancel_inflight_proxy
(after an earlier flaky full-suite failure) -> 1 passed, 1 passed, 1 passed
clean minimum-runtime import + codec smoke test -> passed
grpcio==1.82.1, protobuf==7.35.0
compileall -q packages tests -> OK (exit 0)
git diff --check -> clean (exit 0)
```
The C++ lane was rebuilt from scratch (`rm -rf build/native`) using only the
documented commands, and reproduced the same result.
The C++ lane was rebuilt from scratch by Ralph (`rm -rf build/native`) using only
the documented commands, and reproduced the same result. During controller
review the user explicitly chose not to repeat the destructive build-directory
cleanup, so the independent controller relied on the recorded CMake/CTest run
while reproducing every Python/generation/full-suite gate.
### Controller review corrections
Independent controller review found and fixed two classes of issue before
integration:
1. Generated stubs required gRPC 1.82.1 and Protobuf 7.35.0, while the initial
package metadata allowed much older runtimes that could fail at import time.
2. Flow-control bounds were described but not enforced by the reference decoder.
Tensor declarations, shape rank/dimensions, fragment/tensor counts, fragments,
wire bodies, whole bundles, complete session messages (including envelope
overhead), and zstd window/output expansion are now fail-closed against the
negotiated/default bounds. Unspecified bundle versions, compression and
checksums are rejected rather than interpreted as valid data.
3. Negotiated initial credits could exceed `max_inflight_chunks`; credits are now
capped by the settled in-flight limit.
Controller results: protocol tests `45 passed`; protocol plus shared compression
tests `51 passed`; final full suite `728 passed, 12 skipped`. A clean environment
at the declared minimum gRPC/Protobuf runtime versions imported both generated
stub modules and round-tripped the codec. Generation checks, `compileall`, static
secret scan, and `git diff --check` all passed.
### Full-suite note — a pre-existing flaky test
@@ -104,10 +140,11 @@ pytest -q tests/test_tracker_routing.py::test_tracker_dashboard_can_cancel_infli
It is a timing race in proxy cancellation (a 3-second in-flight generation raced
against the cancel assertion), not a deterministic failure, and it touches no code
this story changes. The final full-suite run with DGR-002 applied was green
(`712 passed, 12 skipped`) — 683 pre-existing tests plus this story's 29. Flagged
for whoever owns the tracker cancel path; **not** fixed here, since silently
touching another story's code is out of scope.
this story changes. One controller full-suite run reported exactly that one failure
(`1 failed, 719 passed, 12 skipped`); three immediate isolated retries all passed
in 1.11 seconds, and the final exact-code full suite was green (`728 passed,
12 skipped`). It is flagged for whoever owns the tracker cancel path and is **not**
fixed here, since silently touching another story's code is out of scope.
## 4. Acceptance criteria
@@ -118,12 +155,12 @@ touching another story's code is out of scope.
| Bounded chunking for prefill; small decode fast path | `ChunkInfo` + `plan_prefill_chunks` (128-token bound, ADR-0008); `DecodeStep`; `test_prefill_is_split_into_bounded_token_aligned_chunks`, `test_decode_fast_path_is_much_smaller_than_a_full_envelope_chunk` |
| Envelope carries schema version, work id, session id, epoch, fingerprint, range/effective start, phase, position, idempotency step, cache expectation, compression, checksum | `Envelope` + `NamedTensor`; `test_envelope_carries_every_field_the_protocol_promises` asserts against the **descriptor**, so deleting a field from the `.proto` fails the test |
| Versioned named-tensor bundle: name, shape, dtype, byte order, fragments | `TensorBundle`/`NamedTensor`/`TensorFragment`; `test_named_tensor_bundle_is_versioned_and_fully_described`, `test_bundle_round_trips_multiple_named_tensors` |
| Round-trip + compatibility tests in Python and C++ | 29 Python tests; C++ `ctest` 1/1; cross-language byte equality |
| Targeted pytest passes | 29 passed |
| Round-trip + compatibility tests in Python and C++ | 45 Python tests; C++ `ctest` 1/1; cross-language byte equality |
| Targeted pytest passes | 45 passed |
| `compileall packages tests` | exit 0 |
| `git diff --check` | exit 0 |
| Default tests deterministic, download-free, credit-free, GPU-free | Pure in-memory protobuf; no model, no network, no GPU |
| Full deterministic pytest passes, or pre-existing failure recorded | 712 passed, 12 skipped; flaky pre-existing test documented above with clean-tree reproduction |
| Full deterministic pytest passes, or pre-existing failure recorded | Final exact-code run: 728 passed, 12 skipped; earlier sole flaky failure documented with clean-tree reproduction and 3/3 passing retries |
## 5. How the cross-language claim is actually earned

16
.scratch/distributed-gguf-runtime/evidence/DGR-002/commands.txt
View File

@@ -27,3 +27,19 @@ cmp build/native/cpp_roundtrip.binpb packages/node/native/testdata/session_reque
# --- repository gates
.venv/bin/python -m compileall -q packages tests
git diff --check
# --- independent controller review after Ralph
PYTHONPATH=packages/node .venv/bin/python -m pytest -q tests/test_native_shard_protocol.py
# -> 45 passed
PYTHONPATH=packages/node .venv/bin/python -m pytest -q \
tests/test_native_shard_protocol.py tests/test_activation_compression.py
# -> 51 passed
PYTHONPATH=packages/node .venv/bin/python -m pytest -q
# -> final exact-code run: 728 passed, 12 skipped
for i in 1 2 3; do PYTHONPATH=packages/node .venv/bin/python -m pytest -q \
tests/test_tracker_routing.py::test_tracker_dashboard_can_cancel_inflight_proxy; done
# -> 1 passed, 1 passed, 1 passed
# clean minimum-runtime venv: protobuf==7.35.0 grpcio==1.82.1
# generated pb2 + pb2_grpc imports and one-byte codec round trip -> passed
# The user chose to rely on Ralph's recorded successful C++ CMake/CTest run
# rather than repeat deletion of an isolated generated build directory.

15
.scratch/distributed-gguf-runtime/evidence/DGR-002/protocol.json
View File

@@ -2,7 +2,7 @@
"schema_version": "SCHEMA_VERSION_1",
"bundle_version": 1,
"proto_path": "packages/node/native/proto/shard_runtime.proto",
"proto_sha256": "077ee34961ab787c67457b1a373f73ec3d577f67086d6bc6a32435f37395b7d6",
"proto_sha256": "9e211660b3fcefc88bcdf3851c3571088c00349aacb5adc5ef45083c83d0cce2",
"protoc": "grpc_tools 1.82.1 (python) / protobuf 33.1 (C++)",
"service": {
"GetCapability": {
@@ -75,10 +75,21 @@
"max_prefill_chunk_tokens": 128,
"max_chunk_bytes": 4194304,
"max_fragment_bytes": 1048576,
"max_inflight_chunks": 8
"max_inflight_chunks": 8,
"max_fragments_per_tensor": 64,
"max_tensors_per_bundle": 64,
"max_tensor_rank": 8,
"max_tensor_dimension": 2147483647,
"whole_session_message_enforced": true
},
"golden_vectors": {
"session_request_golden.binpb": "c2c3df8a717ddeae7bd99624d2c7f34c09a518988de990237fe313b75cff0817",
"capability_report_golden.binpb": "71ac5f150775f398515b43a63596a5cbe8d2ad607e7e4de56bd44fbe7987080c"
},
"verification": {
"python_protocol_tests": "45 passed",
"python_protocol_and_compression_tests": "51 passed",
"full_suite": "728 passed, 12 skipped",
"minimum_runtime": "grpcio 1.82.1 / protobuf 7.35.0 passed import and codec smoke"
}
}

2
.scratch/distributed-gguf-runtime/prd.json
View File

@@ -54,7 +54,7 @@
"Update only this story issue to Status: done after every acceptance criterion and quality gate passes"
],
"priority": 1,
"passes": false,
"passes": true,
"notes": "Source issue: .scratch/distributed-gguf-runtime/issues/02-adopt-the-versioned-grpc-shard-protocol.md",
"dependsOn": []
},

22
packages/node/meshnet_node/activation_compression.py
View File

@@ -99,7 +99,12 @@ def compress_activation(body: bytes, policy: CompressionPolicy) -> CompressionRe
return CompressionResult(candidate, "zstd", len(body), len(candidate), time.monotonic() - started, "compressed")
def decompress_activation(body: bytes, encoding: str | None) -> CompressionResult:
def decompress_activation(
body: bytes,
encoding: str | None,
*,
max_output_bytes: int | None = None,
) -> CompressionResult:
"""Decode a modern zstd body or preserve a legacy raw body with metrics."""
started = time.monotonic()
if not encoding:
@@ -110,8 +115,23 @@ def decompress_activation(body: bytes, encoding: str | None) -> CompressionResul
import zstandard as zstd
except ImportError as exc:
raise ValueError("zstd support is unavailable") from exc
if max_output_bytes is not None and max_output_bytes < 0:
raise ValueError("max_output_bytes must be non-negative")
try:
if max_output_bytes is None:
raw = zstd.ZstdDecompressor().decompress(body)
else:
# Cap both decoder window allocation and bytes read. zstandard's
# max_window_size unit is KiB.
max_window_kib = max(1024, (max_output_bytes + 1023) // 1024)
decompressor = zstd.ZstdDecompressor(max_window_size=max_window_kib)
# `decompress(max_output_size=...)` may trust a frame's advertised
# content size. A bounded stream read enforces the limit regardless
# of frame metadata and detects trailing expansion with one byte.
with decompressor.stream_reader(body) as reader:
raw = reader.read(max_output_bytes + 1)
if len(raw) > max_output_bytes:
raise ValueError("zstd activation body exceeds its output limit")
except zstd.ZstdError as exc:
raise ValueError("invalid zstd activation body") from exc
return CompressionResult(raw, "zstd", len(body), len(raw), time.monotonic() - started, "decompressed")

10
packages/node/meshnet_node/native_protocol/__init__.py
View File

@@ -17,8 +17,12 @@ from .codec import (
BUNDLE_VERSION,
DEFAULT_MAX_CHUNK_BYTES,
DEFAULT_MAX_FRAGMENT_BYTES,
DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
DEFAULT_MAX_INFLIGHT_CHUNKS,
DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
DEFAULT_MAX_TENSORS_PER_BUNDLE,
DEFAULT_MAX_TENSOR_DIMENSION,
DEFAULT_MAX_TENSOR_RANK,
HIDDEN_STATES,
SCHEMA_VERSION,
PayloadCorrupt,
@@ -35,14 +39,19 @@ from .codec import (
itemsize,
negotiate_flow_control,
plan_prefill_chunks,
validate_session_message_size,
)
__all__ = [
"BUNDLE_VERSION",
"DEFAULT_MAX_CHUNK_BYTES",
"DEFAULT_MAX_FRAGMENT_BYTES",
"DEFAULT_MAX_FRAGMENTS_PER_TENSOR",
"DEFAULT_MAX_INFLIGHT_CHUNKS",
"DEFAULT_MAX_PREFILL_CHUNK_TOKENS",
"DEFAULT_MAX_TENSORS_PER_BUNDLE",
"DEFAULT_MAX_TENSOR_DIMENSION",
"DEFAULT_MAX_TENSOR_RANK",
"HIDDEN_STATES",
"SCHEMA_VERSION",
"PayloadCorrupt",
@@ -60,4 +69,5 @@ __all__ = [
"negotiate_flow_control",
"pb",
"plan_prefill_chunks",
"validate_session_message_size",
]

212
packages/node/meshnet_node/native_protocol/codec.py
View File

@@ -48,6 +48,10 @@ DEFAULT_MAX_CHUNK_BYTES = 4 * 1024 * 1024
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"
@@ -79,14 +83,28 @@ def itemsize(dtype: int) -> int:
raise ProtocolError(f"unsupported dtype {dtype}") from None
def expected_bytes(shape: Sequence[int], dtype: int) -> int:
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 any(dim < 0 for dim in shape):
raise ProtocolError(f"negative dimension in shape {list(shape)}")
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
return count * itemsize(dtype)
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 ----------------------------------------------------------------
@@ -136,16 +154,18 @@ def encode_tensor(
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_fragment_bytes <= 0:
raise ProtocolError("max_fragment_bytes must be positive")
declared = expected_bytes(shape, dtype)
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) "
@@ -179,6 +199,10 @@ def encode_tensor(
# 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):
@@ -194,20 +218,57 @@ def encode_tensor(
return tensor
def decode_tensor(tensor: pb.NamedTensor) -> bytes:
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
@@ -231,44 +292,90 @@ def decode_tensor(tensor: pb.NamedTensor) -> bytes:
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):
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 unsupported compression")
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)}"
)
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")
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]) -> 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:
def encode_bundle(
tensors: Iterable[pb.NamedTensor],
*,
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 version {bundle.bundle_version} is newer than this build "
f"understands ({BUNDLE_VERSION})"
f"bundle carries {len(tensor_list)} tensors, exceeding limit {max_tensors}"
)
bundle = pb.TensorBundle(bundle_version=BUNDLE_VERSION, tensors=tensor_list)
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:
@@ -276,10 +383,40 @@ def decode_bundle(bundle: pb.TensorBundle) -> dict[str, bytes]:
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)
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 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 ----------------------------------------------
@@ -363,15 +500,22 @@ def negotiate_flow_control(
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(
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
),

7
packages/node/native/proto/shard_runtime.proto
View File

@@ -94,9 +94,10 @@ message Checksum {
// One slice of a tensor's wire payload.
//
// Fragments bound the size of a single gRPC message. `byte_offset` lets a
// receiver verify the fragments tile the body exactly — no hole, no overlap —
// instead of trusting arrival order.
// Fragments bound individual payload pieces and make coverage independently
// verifiable. The negotiated `FlowControl.max_chunk_bytes` bounds the enclosing
// serialized stream message; `byte_offset` proves the pieces tile the tensor
// exactly — no hole, no overlap — instead of trusting arrival order.
//
// Offsets and payloads describe the *wire* body, which is the compressed frame
// when the parent tensor declares a compression. A zstd frame is not decodable

7
packages/node/pyproject.toml
View File

@@ -13,8 +13,11 @@ dependencies = [
"huggingface-hub>=0.20",
"accelerate>=0.28",
"bitsandbytes>=0.43",
"grpcio>=1.60",
"protobuf>=5",
# Keep these floors aligned with the committed grpcio-tools 1.82.1 output.
# Older runtimes satisfy broad constraints but fail while importing the
# generated modules before a node can report a useful protocol error.
"grpcio>=1.82.1",
"protobuf>=7.35.0",
"rich>=13",
"safetensors>=0.4",
"torch>=2.1",

167
tests/test_native_shard_protocol.py
View File

@@ -26,7 +26,9 @@ from google.protobuf import descriptor_pb2
from meshnet_node.activation_compression import CompressionPolicy
from meshnet_node.native_protocol import (
DEFAULT_MAX_CHUNK_BYTES,
DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
DEFAULT_MAX_TENSORS_PER_BUNDLE,
HIDDEN_STATES,
PayloadCorrupt,
ProtocolError,
@@ -39,6 +41,7 @@ from meshnet_node.native_protocol import (
negotiate_flow_control,
pb,
plan_prefill_chunks,
validate_session_message_size,
)
from meshnet_node.native_protocol import conformance
@@ -233,10 +236,161 @@ def test_bundle_from_a_newer_layout_is_refused():
bundle = encode_bundle([])
bundle.bundle_version = 99
with pytest.raises(ProtocolError, match="newer"):
with pytest.raises(ProtocolError, match="not supported"):
decode_bundle(bundle)
def test_bundle_with_unspecified_version_is_refused():
with pytest.raises(ProtocolError, match="version 0"):
decode_bundle(pb.TensorBundle())
def test_tensor_with_unspecified_compression_is_refused():
tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
tensor.compression = pb.COMPRESSION_UNSPECIFIED
with pytest.raises(ProtocolError, match="unspecified"):
decode_tensor(tensor)
def test_tensor_with_unspecified_checksum_is_refused():
tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
tensor.checksum.algorithm = pb.CHECKSUM_ALGORITHM_UNSPECIFIED
with pytest.raises(ProtocolError, match="unspecified"):
decode_tensor(tensor)
def test_declared_tensor_size_cannot_exceed_negotiated_chunk_bound():
tensor = pb.NamedTensor(
name=HIDDEN_STATES,
shape=[1],
dtype=pb.DTYPE_UINT8,
byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN,
total_bytes=DEFAULT_MAX_CHUNK_BYTES + 1,
compression=pb.COMPRESSION_NONE,
)
with pytest.raises(ProtocolError, match="negotiated chunk bound"):
decode_tensor(tensor)
def test_stricter_session_negotiation_overrides_global_defaults():
tensor = encode_tensor(
HIDDEN_STATES,
b"\xaa" * 32,
[32],
pb.DTYPE_UINT8,
max_fragment_bytes=32,
)
with pytest.raises(ProtocolError, match="16-byte negotiated"):
decode_tensor(tensor, max_chunk_bytes=16, max_fragment_bytes=16)
with pytest.raises(ProtocolError, match="16-byte"):
decode_bundle(
encode_bundle([tensor]),
max_chunk_bytes=16,
max_fragment_bytes=16,
)
def test_fragment_cannot_exceed_protocol_fragment_bound():
tensor = encode_tensor(
HIDDEN_STATES,
b"\xaa" * (1024 * 1024 + 1),
[1024 * 1024 + 1],
pb.DTYPE_UINT8,
max_fragment_bytes=1024 * 1024 + 1,
)
with pytest.raises(PayloadCorrupt, match="fragment larger"):
decode_tensor(tensor)
def test_fragment_count_is_bounded_before_sorting_or_allocation():
tensor = encode_tensor(
HIDDEN_STATES,
b"\xaa" * (DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1),
[DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1],
pb.DTYPE_UINT8,
max_fragment_bytes=1,
max_fragments=DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1,
)
with pytest.raises(PayloadCorrupt, match="fragments.*exceeding"):
decode_tensor(tensor)
def test_tensor_count_is_bounded_before_bundle_decoding():
tensors = [
pb.NamedTensor(name=f"t-{index}")
for index in range(DEFAULT_MAX_TENSORS_PER_BUNDLE + 1)
]
with pytest.raises(ProtocolError, match="tensors.*exceeding"):
encode_bundle(tensors)
@pytest.mark.parametrize(
("shape", "match"),
[
([1] * 9, "rank"),
([-1], "dimension outside"),
([1 << 31], "dimension outside"),
],
)
def test_shape_rank_and_dimensions_are_bounded_before_payload_work(shape, match):
with pytest.raises(ProtocolError, match=match):
encode_tensor(HIDDEN_STATES, b"\x00", shape, pb.DTYPE_UINT8)
def test_complete_session_message_includes_envelope_in_byte_bound():
small = pb.SessionRequest(chunk=pb.ActivationChunk(bundle=encode_bundle([])))
assert validate_session_message_size(small, max_chunk_bytes=64) == small.ByteSize()
oversized = pb.SessionRequest(
chunk=pb.ActivationChunk(
envelope=pb.Envelope(work_id="w" * 128),
bundle=encode_bundle([]),
)
)
with pytest.raises(ProtocolError, match="serialized session message"):
validate_session_message_size(oversized, max_chunk_bytes=64)
def test_bundle_serialized_size_is_bounded():
payload = b"\xaa" * (DEFAULT_MAX_CHUNK_BYTES // 2)
first = encode_tensor("first", payload, [len(payload)], pb.DTYPE_UINT8)
second = encode_tensor("second", payload, [len(payload)], pb.DTYPE_UINT8)
with pytest.raises(ProtocolError, match="serialized tensor bundle"):
encode_bundle([first, second])
def test_compressed_tensor_cannot_expand_past_declared_size():
pytest.importorskip("zstandard")
payload = b"\x00" * 65536
always = CompressionPolicy(min_input_bytes=0, min_savings_bytes=0, min_savings_ratio=0.0)
tensor = encode_tensor(
HIDDEN_STATES, payload, [len(payload)], pb.DTYPE_UINT8, policy=always
)
assert tensor.compression == pb.COMPRESSION_ZSTD
# Simulate a hostile frame that advertises a small output while expanding to
# the original 64 KiB. Shape and total_bytes agree, so only the bounded
# decompressor protects the receiver.
tensor.shape[:] = [1024]
tensor.total_bytes = 1024
with pytest.raises(PayloadCorrupt, match="bounded zstd"):
decode_tensor(tensor)
def test_generated_runtime_floors_match_committed_stubs():
metadata = (REPO_ROOT / "packages/node/pyproject.toml").read_text()
assert '"grpcio>=1.82.1"' in metadata
assert '"protobuf>=7.35.0"' in metadata
# --- Bounded prefill chunking and the decode fast path ----------------------
@@ -310,10 +464,21 @@ def test_flow_control_negotiation_takes_the_strictest_bound():
# A sender cannot talk a worker into unbounded queues by proposing a large
# window: neither peer can raise the other's ceiling.
assert settled.max_inflight_chunks == default_flow_control().max_inflight_chunks
assert settled.credits_granted <= settled.max_inflight_chunks
assert settled.max_chunk_bytes == DEFAULT_MAX_CHUNK_BYTES
assert settled.max_prefill_chunk_tokens == DEFAULT_MAX_PREFILL_CHUNK_TOKENS
def test_initial_credits_never_exceed_negotiated_inflight_limit():
proposed = pb.FlowControl(credits_granted=64, max_inflight_chunks=64)
limits = pb.FlowControl(credits_granted=64, max_inflight_chunks=8)
settled = negotiate_flow_control(proposed, limits)
assert settled.credits_granted == 8
assert settled.max_inflight_chunks == 8
# --- Compatibility ----------------------------------------------------------