413 lines
16 KiB
Python
413 lines
16 KiB
Python
"""Conformance tests for the native Shard protocol (ADR-0020, DGR-002).
|
|
|
|
Three layers are tested, and they are not the same thing:
|
|
|
|
1. The *schema* — asserted against the descriptor, not against the Python
|
|
helpers. If a field the protocol promises to carry were dropped from the
|
|
`.proto`, a test that only exercised the codec would still pass.
|
|
2. The *codec* — that a payload which is corrupt, short, holed, or byte-swapped
|
|
is rejected rather than fed to a forward pass.
|
|
3. *Compatibility* — that an old build preserves fields a newer peer added, and
|
|
that the committed cross-language vectors still encode as promised.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pathlib
|
|
import subprocess
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("google.protobuf", reason="protobuf runtime is required")
|
|
|
|
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_PREFILL_CHUNK_TOKENS,
|
|
HIDDEN_STATES,
|
|
PayloadCorrupt,
|
|
ProtocolError,
|
|
checksum_of,
|
|
decode_bundle,
|
|
decode_tensor,
|
|
default_flow_control,
|
|
encode_bundle,
|
|
encode_tensor,
|
|
negotiate_flow_control,
|
|
pb,
|
|
plan_prefill_chunks,
|
|
)
|
|
from meshnet_node.native_protocol import conformance
|
|
|
|
REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
|
|
CPP_ROUNDTRIP = REPO_ROOT / "build/native" / conformance.CPP_ROUNDTRIP
|
|
|
|
|
|
# --- The schema itself ------------------------------------------------------
|
|
|
|
|
|
def test_service_exposes_capability_health_session_release_and_cancel():
|
|
service = pb.DESCRIPTOR.services_by_name["ShardRuntime"]
|
|
assert set(service.methods_by_name) == {
|
|
"GetCapability",
|
|
"Health",
|
|
"Session",
|
|
"Release",
|
|
"Cancel",
|
|
}
|
|
|
|
|
|
def test_session_is_one_long_lived_bidirectional_stream():
|
|
session = pb.DESCRIPTOR.services_by_name["ShardRuntime"].methods_by_name["Session"]
|
|
assert session.client_streaming, "the seam must stream requests"
|
|
assert session.server_streaming, "the seam must stream responses"
|
|
|
|
# Cancellation must not have to travel down a stream that flow control has
|
|
# wedged, so it also exists as its own unary call.
|
|
cancel = pb.DESCRIPTOR.services_by_name["ShardRuntime"].methods_by_name["Cancel"]
|
|
assert not cancel.client_streaming and not cancel.server_streaming
|
|
|
|
|
|
def test_envelope_carries_every_field_the_protocol_promises():
|
|
# Asserted against the descriptor: this is the acceptance criterion, and it
|
|
# must fail if the .proto drops a field, not merely if the codec stops
|
|
# setting one.
|
|
fields = set(pb.Envelope.DESCRIPTOR.fields_by_name)
|
|
assert {
|
|
"schema_version",
|
|
"work_id",
|
|
"route_session_id",
|
|
"route_epoch",
|
|
"fingerprint",
|
|
"shard_range",
|
|
"phase",
|
|
"position",
|
|
"idempotency_step",
|
|
"cache_expectation",
|
|
"deadline_unix_nanos",
|
|
"chunk",
|
|
} <= fields
|
|
|
|
assert {"model_artifact_digest", "runtime_recipe_digest"} <= set(
|
|
pb.Fingerprint.DESCRIPTOR.fields_by_name
|
|
)
|
|
# Overlap-safe start (ADR-0012) is a distinct field from the registered one.
|
|
assert {"start_layer", "end_layer", "effective_start_layer"} <= set(
|
|
pb.ShardRange.DESCRIPTOR.fields_by_name
|
|
)
|
|
assert {"mode", "expected_past_len"} <= set(
|
|
pb.CacheExpectation.DESCRIPTOR.fields_by_name
|
|
)
|
|
|
|
|
|
def test_named_tensor_bundle_is_versioned_and_fully_described():
|
|
assert "bundle_version" in pb.TensorBundle.DESCRIPTOR.fields_by_name
|
|
assert {"name", "shape", "dtype", "byte_order", "compression", "checksum", "fragments"} <= set(
|
|
pb.NamedTensor.DESCRIPTOR.fields_by_name
|
|
)
|
|
|
|
|
|
def test_phases_and_error_codes_cover_the_lifecycle():
|
|
assert {"PHASE_PREFILL", "PHASE_DECODE", "PHASE_RELEASE", "PHASE_CANCEL"} <= set(
|
|
pb.Phase.keys()
|
|
)
|
|
# A cache miss is a first-class, recoverable outcome (ADR-0022), not a crash.
|
|
assert {
|
|
"ERROR_CODE_CACHE_MISS",
|
|
"ERROR_CODE_FINGERPRINT_MISMATCH",
|
|
"ERROR_CODE_EPOCH_STALE",
|
|
"ERROR_CODE_SCHEMA_UNSUPPORTED",
|
|
"ERROR_CODE_DEADLINE_EXCEEDED",
|
|
"ERROR_CODE_FLOW_CONTROL_VIOLATION",
|
|
} <= set(pb.ErrorCode.keys())
|
|
|
|
|
|
# --- Tensor bundle round-trip ----------------------------------------------
|
|
|
|
|
|
def test_tensor_round_trips_through_fragments():
|
|
payload = bytes(range(256)) * 4
|
|
tensor = encode_tensor(
|
|
HIDDEN_STATES, payload, [1, 64, 8], pb.DTYPE_BFLOAT16, max_fragment_bytes=100
|
|
)
|
|
|
|
assert len(tensor.fragments) > 1, "the bound must actually split the payload"
|
|
assert tensor.total_bytes == len(payload)
|
|
assert decode_tensor(tensor) == payload
|
|
|
|
|
|
def test_bundle_round_trips_multiple_named_tensors():
|
|
# An architecture boundary may need more than one tensor; that is why the
|
|
# payload is a named bundle rather than a bare buffer.
|
|
hidden = encode_tensor(HIDDEN_STATES, b"\x01\x02" * 8, [1, 8, 1], pb.DTYPE_BFLOAT16)
|
|
positions = encode_tensor(
|
|
"position_ids", (7).to_bytes(4, "little") * 8, [8], pb.DTYPE_INT32
|
|
)
|
|
bundle = encode_bundle([hidden, positions])
|
|
|
|
restored = decode_bundle(pb.TensorBundle.FromString(bundle.SerializeToString()))
|
|
assert restored == {
|
|
HIDDEN_STATES: b"\x01\x02" * 8,
|
|
"position_ids": (7).to_bytes(4, "little") * 8,
|
|
}
|
|
|
|
|
|
def test_compressed_tensor_round_trips_and_keeps_its_uncompressed_checksum():
|
|
pytest.importorskip("zstandard")
|
|
# Highly compressible, and over the policy's minimum input size.
|
|
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, [1, 4096, 8], pb.DTYPE_BFLOAT16, policy=always
|
|
)
|
|
|
|
assert tensor.compression == pb.COMPRESSION_ZSTD
|
|
assert sum(len(f.payload) for f in tensor.fragments) < len(payload)
|
|
# The checksum covers the uncompressed bytes, so it stays valid whether or
|
|
# not a hop chose to compress.
|
|
assert tensor.checksum == checksum_of(payload)
|
|
assert decode_tensor(tensor) == payload
|
|
|
|
|
|
# --- The codec refuses what it cannot account for --------------------------
|
|
|
|
|
|
def test_corrupt_payload_is_rejected_by_checksum():
|
|
tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
|
|
# Flip one byte, as a lossy relay or a bad NIC would.
|
|
tensor.fragments[0].payload = b"\xab" + tensor.fragments[0].payload[1:]
|
|
|
|
with pytest.raises(PayloadCorrupt, match="CRC32C"):
|
|
decode_tensor(tensor)
|
|
|
|
|
|
def test_missing_fragment_is_rejected_rather_than_silently_truncated():
|
|
tensor = encode_tensor(
|
|
HIDDEN_STATES, b"\xaa" * 64, [1, 32, 1], pb.DTYPE_BFLOAT16, max_fragment_bytes=16
|
|
)
|
|
del tensor.fragments[1]
|
|
|
|
with pytest.raises(PayloadCorrupt):
|
|
decode_tensor(tensor)
|
|
|
|
|
|
def test_fragment_hole_is_rejected():
|
|
tensor = encode_tensor(
|
|
HIDDEN_STATES, b"\xaa" * 64, [1, 32, 1], pb.DTYPE_BFLOAT16, max_fragment_bytes=16
|
|
)
|
|
# A gap in coverage: offsets no longer tile the body.
|
|
tensor.fragments[2].byte_offset += 4
|
|
|
|
with pytest.raises(PayloadCorrupt, match="expected"):
|
|
decode_tensor(tensor)
|
|
|
|
|
|
def test_shape_that_disagrees_with_payload_is_rejected_at_encode():
|
|
with pytest.raises(ProtocolError, match="carries"):
|
|
encode_tensor(HIDDEN_STATES, b"\x01\x02", [1, 8, 8], pb.DTYPE_BFLOAT16)
|
|
|
|
|
|
def test_shape_that_disagrees_with_declared_bytes_is_rejected_at_decode():
|
|
tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
|
|
# A peer claiming a larger tensor than its bytes describe.
|
|
tensor.shape[1] = 32
|
|
|
|
with pytest.raises(PayloadCorrupt, match="implies"):
|
|
decode_tensor(tensor)
|
|
|
|
|
|
def test_big_endian_tensor_is_rejected_loudly():
|
|
tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
|
|
tensor.byte_order = pb.BYTE_ORDER_BIG_ENDIAN
|
|
|
|
# Byte-swapped activations would be plausible-looking garbage, so this is an
|
|
# error rather than a best-effort read.
|
|
with pytest.raises(ProtocolError, match="big-endian"):
|
|
decode_tensor(tensor)
|
|
|
|
|
|
def test_bundle_from_a_newer_layout_is_refused():
|
|
bundle = encode_bundle([])
|
|
bundle.bundle_version = 99
|
|
|
|
with pytest.raises(ProtocolError, match="newer"):
|
|
decode_bundle(bundle)
|
|
|
|
|
|
# --- Bounded prefill chunking and the decode fast path ----------------------
|
|
|
|
|
|
def test_prefill_is_split_into_bounded_token_aligned_chunks():
|
|
chunks = plan_prefill_chunks(2048)
|
|
|
|
assert all(c.token_count <= DEFAULT_MAX_PREFILL_CHUNK_TOKENS for c in chunks)
|
|
assert sum(c.token_count for c in chunks) == 2048
|
|
# Contiguous, token-aligned, and no split falls mid-token.
|
|
assert [c.first_position for c in chunks] == [
|
|
i * DEFAULT_MAX_PREFILL_CHUNK_TOKENS for i in range(len(chunks))
|
|
]
|
|
assert [c.final_chunk for c in chunks] == [False] * (len(chunks) - 1) + [True]
|
|
|
|
|
|
def test_final_prefill_chunk_carries_the_remainder():
|
|
chunks = plan_prefill_chunks(300, max_tokens=128)
|
|
|
|
assert [c.token_count for c in chunks] == [128, 128, 44]
|
|
assert chunks[-1].chunk_info().final_chunk
|
|
assert chunks[0].position().first_position == 0
|
|
assert chunks[-1].position().first_position == 256
|
|
|
|
|
|
def test_empty_prefill_is_refused():
|
|
with pytest.raises(ProtocolError):
|
|
plan_prefill_chunks(0)
|
|
|
|
|
|
def test_decode_fast_path_is_much_smaller_than_a_full_envelope_chunk():
|
|
hidden = b"\x01\x02" * 8 # one token, hidden=8, bfloat16
|
|
tensor = encode_tensor(HIDDEN_STATES, hidden, [1, 1, 8], pb.DTYPE_BFLOAT16)
|
|
|
|
fast = pb.SessionRequest(
|
|
decode=pb.DecodeStep(
|
|
idempotency_step=9,
|
|
position=1024,
|
|
expected_past_len=1024,
|
|
tensor=tensor,
|
|
work_id="work-7f3a",
|
|
)
|
|
)
|
|
# The same single token carried the long way, repeating identity that the
|
|
# handshake already fixed for the life of the stream.
|
|
full = pb.SessionRequest(
|
|
chunk=pb.ActivationChunk(
|
|
envelope=conformance.canonical_session_request().chunk.envelope,
|
|
bundle=encode_bundle([tensor]),
|
|
)
|
|
)
|
|
|
|
assert len(fast.SerializeToString()) * 2 < len(full.SerializeToString())
|
|
assert decode_tensor(fast.decode.tensor) == hidden
|
|
|
|
|
|
def test_flow_control_defaults_bound_the_queue_and_the_message():
|
|
limits = default_flow_control()
|
|
|
|
assert limits.max_prefill_chunk_tokens == DEFAULT_MAX_PREFILL_CHUNK_TOKENS
|
|
assert limits.max_chunk_bytes == DEFAULT_MAX_CHUNK_BYTES
|
|
assert limits.max_inflight_chunks > 0
|
|
|
|
|
|
def test_flow_control_negotiation_takes_the_strictest_bound():
|
|
proposed = pb.FlowControl(
|
|
credits_granted=64, max_inflight_chunks=64, max_chunk_bytes=64 << 20,
|
|
max_prefill_chunk_tokens=1024,
|
|
)
|
|
settled = negotiate_flow_control(proposed, default_flow_control())
|
|
|
|
# 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.max_chunk_bytes == DEFAULT_MAX_CHUNK_BYTES
|
|
assert settled.max_prefill_chunk_tokens == DEFAULT_MAX_PREFILL_CHUNK_TOKENS
|
|
|
|
|
|
# --- Compatibility ----------------------------------------------------------
|
|
|
|
|
|
def test_committed_vectors_still_encode_as_promised():
|
|
# The C++ test asserts against these exact bytes. If a schema change alters
|
|
# the canonical encoding, it must be acknowledged by regenerating them.
|
|
golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()
|
|
assert conformance.serialize(conformance.canonical_session_request()) == golden
|
|
|
|
report = (conformance.TESTDATA_DIR / conformance.GOLDEN_CAPABILITY_REPORT).read_bytes()
|
|
assert conformance.serialize(conformance.canonical_capability_report()) == report
|
|
|
|
|
|
def test_golden_session_request_round_trips_with_every_field_intact():
|
|
golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()
|
|
request = pb.SessionRequest.FromString(golden)
|
|
envelope = request.chunk.envelope
|
|
|
|
assert envelope.work_id == conformance.WORK_ID
|
|
assert envelope.route_session_id == conformance.ROUTE_SESSION_ID
|
|
assert envelope.route_epoch == conformance.ROUTE_EPOCH
|
|
assert envelope.idempotency_step == conformance.IDEMPOTENCY_STEP
|
|
assert envelope.shard_range.effective_start_layer == conformance.EFFECTIVE_START_LAYER
|
|
assert envelope.phase == pb.PHASE_PREFILL
|
|
assert envelope.cache_expectation.mode == pb.CACHE_MODE_PREFILL
|
|
assert envelope.deadline_unix_nanos == conformance.DEADLINE_UNIX_NANOS
|
|
assert decode_bundle(request.chunk.bundle) == {
|
|
HIDDEN_STATES: conformance.canonical_payload()
|
|
}
|
|
assert request.SerializeToString(deterministic=True) == golden
|
|
|
|
|
|
def test_unknown_fields_from_a_newer_peer_survive_a_forwarding_hop():
|
|
# A Shard forwards activations onward. If it silently dropped fields a newer
|
|
# peer added, it would corrupt a route it is merely a waypoint on.
|
|
golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()
|
|
future_field = b"\xb8\xe0\x04\xb9\x60" # field 9999, varint 12345
|
|
request = pb.SessionRequest.FromString(golden + future_field)
|
|
|
|
assert request.chunk.envelope.work_id == conformance.WORK_ID
|
|
assert request.SerializeToString() == golden + future_field
|
|
|
|
|
|
def test_a_message_missing_newer_field_groups_still_parses():
|
|
sparse = pb.SessionRequest(chunk=pb.ActivationChunk(envelope=pb.Envelope(work_id="w")))
|
|
parsed = pb.SessionRequest.FromString(sparse.SerializeToString())
|
|
|
|
assert parsed.chunk.envelope.work_id == "w"
|
|
assert parsed.chunk.envelope.route_epoch == 0
|
|
assert parsed.chunk.envelope.phase == pb.PHASE_UNSPECIFIED
|
|
|
|
|
|
def test_retired_fragment_field_stays_reserved():
|
|
# `uncompressed_size` (field 5) was removed because NamedTensor.total_bytes
|
|
# is the single source of truth. The number stays reserved so it can never
|
|
# be recycled for a different meaning — a recycled number is the one schema
|
|
# change that old and new peers cannot detect, because the bytes still parse.
|
|
descriptor = descriptor_pb2.DescriptorProto()
|
|
pb.TensorFragment.DESCRIPTOR.CopyToProto(descriptor)
|
|
|
|
assert 5 not in {field.number for field in descriptor.field}
|
|
assert any(r.start <= 5 < r.end for r in descriptor.reserved_range)
|
|
assert "uncompressed_size" in descriptor.reserved_name
|
|
|
|
|
|
def test_a_peer_still_sending_the_retired_field_does_not_corrupt_the_tensor():
|
|
# An older peer that still sets field 5 must be parsed, not rejected: the
|
|
# value lands in unknown fields and the payload is unaffected.
|
|
tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
|
|
wire = tensor.SerializeToString()
|
|
fragment = pb.TensorFragment.FromString(tensor.fragments[0].SerializeToString() + b"\x28\x20")
|
|
|
|
assert fragment.payload == tensor.fragments[0].payload
|
|
assert decode_tensor(pb.NamedTensor.FromString(wire)) == b"\xaa" * 32
|
|
|
|
|
|
def test_generated_python_stubs_match_the_proto():
|
|
pytest.importorskip("grpc_tools", reason="protoc toolchain is required to verify")
|
|
result = subprocess.run(
|
|
[sys.executable, str(REPO_ROOT / "scripts/generate_native_protocol.py"), "--check"],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
assert result.returncode == 0, result.stdout + result.stderr
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not CPP_ROUNDTRIP.is_file(),
|
|
reason="build the C++ conformance test to check cross-language agreement",
|
|
)
|
|
def test_cpp_and_python_agree_byte_for_byte():
|
|
# Written by the C++ conformance test: it parsed the golden bytes into its
|
|
# own object model and serialized them back. Byte equality means both
|
|
# languages encode every field of this schema identically.
|
|
golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()
|
|
assert CPP_ROUNDTRIP.read_bytes() == golden
|