"""DGR-002: generated-schema round-trip and compatibility tests. Covers the versioned gRPC Shard protocol (``packages/node/native/proto``): * Python round-trip across the full envelope, tensor bundle, and every service. * Proto3 forward/backward compatibility (unknown-field preservation, defaults). * Bounded-fragment tensor bundle framing + checksums. * Cross-language Python<->C++ round-trip when the C++ toolchain is available; otherwise the C++ test skips with an explicit reason (deterministic, GPU-free, model-download-free, API-credit-free by construction). """ from __future__ import annotations import shutil import subprocess import pytest # grpc_tools (grpcio-tools) is required to generate the stubs. It is present in # the project .venv; skip cleanly elsewhere rather than error. native_protocol = pytest.importorskip( "meshnet_node.native_protocol", reason="meshnet_node.native_protocol import failed", ) try: native_protocol.generate() _GEN_ERROR = None except native_protocol.ProtocGenerationError as exc: # pragma: no cover _GEN_ERROR = str(exc) pytestmark = pytest.mark.skipif( _GEN_ERROR is not None, reason=f"protobuf stubs unavailable: {_GEN_ERROR}", ) @pytest.fixture(scope="module") def pb2(): return native_protocol.load() # --------------------------------------------------------------------------- # Envelope / header round-trip and field coverage # --------------------------------------------------------------------------- def _full_header(pb2): return pb2.MessageHeader( schema_version=pb2.SCHEMA_VERSION_1, work_id="work-42", route_session_id="rs-7", route_epoch=9, fingerprint=pb2.ArtifactFingerprint( model_id="meta-llama/Llama-3.1-8B", revision="main", artifact_hash="sha256:deadbeef", quantization="Q4_K_M", runtime_recipe_fingerprint="recipe-123", ), shard_range=pb2.ShardRange( start_layer=8, end_layer=16, effective_start_layer=9, owns_embedding=False, owns_final_head=False, ), phase=pb2.PHASE_PREFILL, position=pb2.Position(start_position=0, token_count=12, sequence_length=12), idempotency_step=3, cache_expectation=pb2.CACHE_REUSE, compression=pb2.COMPRESSION_ZSTD, checksum=pb2.Checksum(algorithm=pb2.CHECKSUM_CRC32C, value=b"\x00\x01\x02\x03"), ) def test_message_header_carries_every_required_field(pb2): """The header carries every identifier the transport contract demands. Tags: protocol """ header = _full_header(pb2) raw = header.SerializeToString() back = pb2.MessageHeader() back.ParseFromString(raw) assert back.schema_version == pb2.SCHEMA_VERSION_1 assert back.work_id == "work-42" assert back.route_session_id == "rs-7" assert back.route_epoch == 9 assert back.fingerprint.artifact_hash == "sha256:deadbeef" assert back.fingerprint.runtime_recipe_fingerprint == "recipe-123" assert back.shard_range.effective_start_layer == 9 assert back.phase == pb2.PHASE_PREFILL assert back.position.token_count == 12 assert back.idempotency_step == 3 assert back.cache_expectation == pb2.CACHE_REUSE assert back.compression == pb2.COMPRESSION_ZSTD assert back.checksum.algorithm == pb2.CHECKSUM_CRC32C assert back.checksum.value == b"\x00\x01\x02\x03" def test_named_tensor_bundle_describes_shape_dtype_byteorder_and_fragments(pb2): """A tensor bundle round-trips name, shape, dtype, byte order and fragments. Tags: protocol """ bundle = pb2.TensorBundle( bundle_version=1, tensors=[ pb2.NamedTensor( name="hidden_states", shape=[2, 3, 4096], dtype=pb2.DTYPE_BF16, byte_order=pb2.BYTE_ORDER_LITTLE_ENDIAN, total_byte_length=16, compression=pb2.COMPRESSION_NONE, fragments=[ pb2.TensorFragment( fragment_index=0, fragment_count=2, byte_offset=0, data=b"\x00" * 8, ), pb2.TensorFragment( fragment_index=1, fragment_count=2, byte_offset=8, data=b"\x01" * 8, ), ], ) ], ) back = pb2.TensorBundle() back.ParseFromString(bundle.SerializeToString()) tensor = back.tensors[0] assert tensor.name == "hidden_states" assert list(tensor.shape) == [2, 3, 4096] assert tensor.dtype == pb2.DTYPE_BF16 assert tensor.byte_order == pb2.BYTE_ORDER_LITTLE_ENDIAN assert [f.byte_offset for f in tensor.fragments] == [0, 8] def test_session_stream_carries_open_prefill_decode_release_cancel(pb2): """The bidi stream oneof expresses every seam operation. Tags: protocol """ header = _full_header(pb2) frames = { "open": pb2.SessionActivation( open=pb2.SessionOpen( header=header, deadline_unix_nanos=1_000_000, max_prefill_tokens_per_chunk=256, max_fragment_bytes=1 << 20, initial_credit=pb2.FlowControl(credits=8, max_in_flight_bytes=1 << 24), ) ), "prefill": pb2.SessionActivation( prefill=pb2.PrefillChunk( header=header, chunk_index=0, chunk_count=2, final_chunk=False ) ), "decode": pb2.SessionActivation(decode=pb2.DecodeStep(header=header)), "release": pb2.SessionActivation( release=pb2.ReleaseRequest(header=header, reason="done") ), "cancel": pb2.SessionActivation( cancel=pb2.CancelRequest(header=header, reason="client abort") ), "flow_control": pb2.SessionActivation( flow_control=pb2.FlowControl(credits=4) ), } for name, frame in frames.items(): back = pb2.SessionActivation() back.ParseFromString(frame.SerializeToString()) assert back.WhichOneof("payload") == name def test_session_response_carries_structured_status_and_results(pb2): """Server frames carry accepted/result/status/acks with structured Status. Tags: protocol """ status = pb2.Status( code=8, message="resource exhausted", retry_class=pb2.RETRY_CLASS_RETRYABLE, details={"queue_depth": "128"}, ) resp = pb2.SessionResponse( result=pb2.ActivationResult( header=_full_header(pb2), outputs=pb2.TensorBundle(bundle_version=1), cache_result=pb2.CACHE_WRITTEN, status=status, ) ) back = pb2.SessionResponse() back.ParseFromString(resp.SerializeToString()) assert back.WhichOneof("payload") == "result" assert back.result.cache_result == pb2.CACHE_WRITTEN assert back.result.status.retry_class == pb2.RETRY_CLASS_RETRYABLE assert back.result.status.details["queue_depth"] == "128" def test_capability_and_health_round_trip(pb2): """Capability and health messages round-trip their admission fields. Tags: protocol """ cap = pb2.CapabilityResponse( schema_version=pb2.SCHEMA_VERSION_1, supported_schema_versions=[pb2.SCHEMA_VERSION_1], supported_architectures=["llama"], supported_quantizations=["Q4_K_M", "F16"], servable_range=pb2.ShardRange(start_layer=0, end_layer=16), budget=pb2.ResourceBudget( weight_bytes=1 << 32, kv_bytes=1 << 30, max_concurrent_sessions=4 ), supported_compression=[pb2.COMPRESSION_NONE, pb2.COMPRESSION_ZSTD], supported_checksums=[pb2.CHECKSUM_CRC32C, pb2.CHECKSUM_SHA256], ) cap_back = pb2.CapabilityResponse() cap_back.ParseFromString(cap.SerializeToString()) assert cap_back.budget.max_concurrent_sessions == 4 assert list(cap_back.supported_quantizations) == ["Q4_K_M", "F16"] health = pb2.HealthResponse( status=pb2.SERVING, active_sessions=2, queued_requests=1, kv_pressure=0.5 ) health_back = pb2.HealthResponse() health_back.ParseFromString(health.SerializeToString()) assert health_back.status == pb2.SERVING assert health_back.kv_pressure == pytest.approx(0.5) # --------------------------------------------------------------------------- # Compatibility # --------------------------------------------------------------------------- def test_unknown_fields_are_preserved_for_forward_compatibility(pb2): """An older reader tolerates and preserves fields it does not know. A newer sender may add a field; parsing into the current schema must not fail and must round-trip the unknown bytes. Tags: protocol, compatibility """ header = _full_header(pb2) raw = bytearray(header.SerializeToString()) # Append an unknown field: number 5000, wire type 2 (length-delimited). tag = (5000 << 3) | 2 raw += _encode_varint(tag) payload = b"future-field" raw += _encode_varint(len(payload)) raw += payload parsed = pb2.MessageHeader() # Parsing must not raise on the unknown field. parsed.ParseFromString(bytes(raw)) # Known fields survive intact. assert parsed.work_id == "work-42" assert parsed.route_epoch == 9 # The unknown bytes are preserved and re-emitted on re-serialization. This is # the behavioural compatibility guarantee; the introspection accessor # (UnknownFields()) is not implemented by the upb backend, so we assert the # observable outcome rather than the accessor. reserialized = parsed.SerializeToString() assert payload in reserialized assert _encode_varint(tag) in reserialized def test_defaults_are_stable_for_backward_compatibility(pb2): """A message from an older sender (missing new fields) reads as enum zero. Tags: protocol, compatibility """ empty = pb2.MessageHeader() back = pb2.MessageHeader() back.ParseFromString(empty.SerializeToString()) assert back.schema_version == pb2.SCHEMA_VERSION_UNSPECIFIED assert back.phase == pb2.PHASE_UNSPECIFIED assert back.cache_expectation == pb2.CACHE_EXPECTATION_UNSPECIFIED assert back.work_id == "" assert back.route_epoch == 0 # --------------------------------------------------------------------------- # Bounded-fragment helpers # --------------------------------------------------------------------------- def test_fragment_and_reassemble_round_trip_with_checksums(pb2): """Bounded fragmentation reassembles exactly and validates checksums. Tags: protocol """ payload = bytes((i * 7) % 256 for i in range(10_000)) tensor = native_protocol.fragment_tensor( name="hidden", shape=[1, 4096], dtype=pb2.DTYPE_F16, payload=payload, max_fragment_bytes=4096, checksum_algorithm=pb2.CHECKSUM_CRC32C, ) assert len(tensor.fragments) == 3 assert all(len(f.data) <= 4096 for f in tensor.fragments) # Survives a serialization round-trip before reassembly. back = pb2.NamedTensor() back.ParseFromString(tensor.SerializeToString()) assert native_protocol.reassemble_tensor(back) == payload def test_reassemble_detects_fragment_corruption(pb2): """A flipped fragment byte fails checksum verification. Tags: protocol """ payload = b"abcdefabcdef" * 100 tensor = native_protocol.fragment_tensor( name="t", shape=[len(payload)], dtype=pb2.DTYPE_U8, payload=payload, max_fragment_bytes=256, checksum_algorithm=pb2.CHECKSUM_SHA256, ) tensor.fragments[1].data = tensor.fragments[1].data[:-1] + b"\x00" with pytest.raises(ValueError, match="checksum mismatch"): native_protocol.reassemble_tensor(tensor) def test_checksum_algorithms_verify(pb2): """CRC32C, CRC32 and SHA256 all verify their own payloads. Tags: protocol """ data = b"the quick brown fox" for algo in (pb2.CHECKSUM_CRC32C, pb2.CHECKSUM_CRC32, pb2.CHECKSUM_SHA256): checksum = native_protocol.compute_checksum(algo, data) assert native_protocol.verify_checksum(checksum, data) assert not native_protocol.verify_checksum(checksum, data + b"!") def test_service_descriptor_exposes_all_operations(pb2): """The generated service defines capability/health/session/release/cancel. Requires the grpc runtime; skips cleanly without it. Tags: protocol """ grpc = pytest.importorskip("grpc", reason="grpc runtime not installed") assert grpc is not None grpc_mod = native_protocol.load_grpc() assert hasattr(grpc_mod, "ShardRuntimeStub") assert hasattr(grpc_mod, "ShardRuntimeServicer") # Confirm the streaming seam and unary ops exist on the servicer. servicer = grpc_mod.ShardRuntimeServicer for op in ("GetCapability", "Health", "ActivateSession", "Release", "Cancel"): assert hasattr(servicer, op), op # --------------------------------------------------------------------------- # Cross-language Python <-> C++ compatibility # --------------------------------------------------------------------------- def _cpp_toolchain_reason() -> str | None: """Return a skip reason if the C++ build toolchain is unavailable.""" for tool in ("cmake", "protoc"): if shutil.which(tool) is None: return f"{tool} not found on PATH" return None def _build_cpp_compatible_sample(pb2): """Python message matching what roundtrip_test.cpp CheckSample expects.""" header = pb2.MessageHeader( schema_version=pb2.SCHEMA_VERSION_1, work_id="w1", route_session_id="s1", route_epoch=3, phase=pb2.PHASE_PREFILL, idempotency_step=7, cache_expectation=pb2.CACHE_FRESH, compression=pb2.COMPRESSION_NONE, fingerprint=pb2.ArtifactFingerprint( model_id="meta-llama/Llama-3.1-8B", quantization="Q4_K_M", runtime_recipe_fingerprint="recipe-abc", ), shard_range=pb2.ShardRange( start_layer=0, end_layer=16, effective_start_layer=0, owns_embedding=True ), position=pb2.Position(start_position=0, token_count=5, sequence_length=5), ) return pb2.SessionActivation( prefill=pb2.PrefillChunk( header=header, chunk_index=0, chunk_count=1, final_chunk=True, activations=pb2.TensorBundle( bundle_version=1, tensors=[ pb2.NamedTensor( name="hidden", shape=[1, 4096], dtype=pb2.DTYPE_F16, byte_order=pb2.BYTE_ORDER_LITTLE_ENDIAN, total_byte_length=8, compression=pb2.COMPRESSION_NONE, fragments=[ pb2.TensorFragment( fragment_index=0, fragment_count=1, byte_offset=0, data=bytes(range(1, 9)), ) ], ) ], ), ) ) def test_cross_language_roundtrip_python_and_cpp(pb2, tmp_path): """Python and C++ parse each other's serialized frames (both directions). Builds the C++ round-trip binary via CMake, feeds it a Python-serialized fixture (C++ must parse it), and parses the C++-serialized output back in Python. Skips with an explicit reason when the C++ toolchain is absent. Tags: protocol, compatibility, cpp """ reason = _cpp_toolchain_reason() if reason is not None: pytest.skip(f"C++ toolchain unavailable: {reason}") native_root = native_protocol.PROTO_DIR.parent build_dir = tmp_path / "cpp-build" configure = subprocess.run( ["cmake", "-S", str(native_root), "-B", str(build_dir)], capture_output=True, text=True, ) if configure.returncode != 0: pytest.skip( "cmake configure failed (protobuf C++ dev likely missing):\n" + configure.stderr[-2000:] ) build = subprocess.run( ["cmake", "--build", str(build_dir), "--target", "shard_protocol_roundtrip_test"], capture_output=True, text=True, ) assert build.returncode == 0, f"C++ build failed:\n{build.stderr[-2000:]}" binary = build_dir / "shard_protocol_roundtrip_test" assert binary.exists(), "C++ test binary not produced" py_fixture = tmp_path / "py_sample.bin" cpp_out = tmp_path / "cpp_sample.bin" py_fixture.write_bytes(_build_cpp_compatible_sample(pb2).SerializeToString()) run = subprocess.run( [str(binary), "--selftest", "--read", str(py_fixture), "--write", str(cpp_out)], capture_output=True, text=True, ) assert run.returncode == 0, f"C++ round-trip failed:\n{run.stdout}\n{run.stderr}" # C++ parsed our bytes; now Python parses C++'s bytes and checks known fields. parsed = pb2.SessionActivation() parsed.ParseFromString(cpp_out.read_bytes()) assert parsed.WhichOneof("payload") == "prefill" assert parsed.prefill.header.work_id == "w1" assert parsed.prefill.header.route_epoch == 3 assert parsed.prefill.activations.tensors[0].name == "hidden" assert parsed.prefill.activations.tensors[0].dtype == pb2.DTYPE_F16 # --------------------------------------------------------------------------- # Local helpers # --------------------------------------------------------------------------- def _encode_varint(value: int) -> bytes: out = bytearray() while True: byte = value & 0x7F value >>= 7 if value: out.append(byte | 0x80) else: out.append(byte) return bytes(out)