509 lines
18 KiB
Python
509 lines
18 KiB
Python
"""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)
|