diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-002/README.md b/.scratch/distributed-gguf-runtime/evidence/DGR-002/README.md index 11de0cf..cf1d00f 100644 --- a/.scratch/distributed-gguf-runtime/evidence/DGR-002/README.md +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-002/README.md @@ -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 diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-002/commands.txt b/.scratch/distributed-gguf-runtime/evidence/DGR-002/commands.txt index 2c197b7..42a6b0a 100644 --- a/.scratch/distributed-gguf-runtime/evidence/DGR-002/commands.txt +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-002/commands.txt @@ -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. diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-002/protocol.json b/.scratch/distributed-gguf-runtime/evidence/DGR-002/protocol.json index 2dda264..748b314 100644 --- a/.scratch/distributed-gguf-runtime/evidence/DGR-002/protocol.json +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-002/protocol.json @@ -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" } } diff --git a/.scratch/distributed-gguf-runtime/prd.json b/.scratch/distributed-gguf-runtime/prd.json index 11d9b11..323eaae 100644 --- a/.scratch/distributed-gguf-runtime/prd.json +++ b/.scratch/distributed-gguf-runtime/prd.json @@ -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": [] }, diff --git a/packages/node/meshnet_node/activation_compression.py b/packages/node/meshnet_node/activation_compression.py index 87e4226..ffd0c8a 100644 --- a/packages/node/meshnet_node/activation_compression.py +++ b/packages/node/meshnet_node/activation_compression.py @@ -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: - raw = zstd.ZstdDecompressor().decompress(body) + 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") diff --git a/packages/node/meshnet_node/native_protocol/__init__.py b/packages/node/meshnet_node/native_protocol/__init__.py index 235f9f5..16e00a0 100644 --- a/packages/node/meshnet_node/native_protocol/__init__.py +++ b/packages/node/meshnet_node/native_protocol/__init__.py @@ -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", ] diff --git a/packages/node/meshnet_node/native_protocol/codec.py b/packages/node/meshnet_node/native_protocol/codec.py index 00eb4db..80d6ad1 100644 --- a/packages/node/meshnet_node/native_protocol/codec.py +++ b/packages/node/meshnet_node/native_protocol/codec.py @@ -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( - proposed.max_inflight_chunks, - limits.max_inflight_chunks, + 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 ), diff --git a/packages/node/native/proto/shard_runtime.proto b/packages/node/native/proto/shard_runtime.proto index 71a00c9..cf97d96 100644 --- a/packages/node/native/proto/shard_runtime.proto +++ b/packages/node/native/proto/shard_runtime.proto @@ -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 diff --git a/packages/node/pyproject.toml b/packages/node/pyproject.toml index a11d2cb..41fd484 100644 --- a/packages/node/pyproject.toml +++ b/packages/node/pyproject.toml @@ -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", diff --git a/tests/test_native_shard_protocol.py b/tests/test_native_shard_protocol.py index ab36e3e..bd50f5c 100644 --- a/tests/test_native_shard_protocol.py +++ b/tests/test_native_shard_protocol.py @@ -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 ----------------------------------------------------------