Files
neuron-tai/packages/node/meshnet_node/native_protocol/conformance.py
2026-07-13 16:00:49 +03:00

142 lines
5.2 KiB
Python

"""Canonical conformance vectors for the native Shard protocol.
Two independently-written codecs that each round-trip their own output prove
nothing about each other. These vectors are the shared reference: Python builds
the canonical message, the bytes are committed under
`packages/node/native/testdata/`, and the C++ test parses those exact bytes and
asserts the same field values. A change that alters the wire meaning of a field
breaks the vector in both languages instead of drifting silently in one.
The vector deliberately exercises every field group the protocol promises to
carry — identity, epoch, fingerprint, range, phase, position, idempotency,
cache expectation, deadline, chunking, compression, checksum and a multi-
fragment named tensor — so it doubles as an executable inventory of the
contract.
"""
from __future__ import annotations
import pathlib
from . import codec
from .generated import shard_runtime_pb2 as pb
# Committed vectors live beside the schema, under `packages/node/native/`.
# parents[2] is `packages/node`: native_protocol -> meshnet_node -> node.
TESTDATA_DIR = pathlib.Path(__file__).resolve().parents[2] / "native/testdata"
GOLDEN_SESSION_REQUEST = "session_request_golden.binpb"
GOLDEN_CAPABILITY_REPORT = "capability_report_golden.binpb"
# Written by the C++ conformance test into its build tree; the Python test picks
# it up when present to prove the two languages agree byte-for-byte.
CPP_ROUNDTRIP = "cpp_roundtrip.binpb"
# Fixed, non-default values. Every one is chosen to be distinguishable from a
# proto3 default so an unset field can never masquerade as a correct one.
WORK_ID = "work-7f3a"
ROUTE_SESSION_ID = "rs-2b91"
ROUTE_EPOCH = 7
IDEMPOTENCY_STEP = 42
FIRST_POSITION = 256
TOKEN_COUNT = 128
EXPECTED_PAST_LEN = 256
DEADLINE_UNIX_NANOS = 1_800_000_000_000_000_000
MODEL_ARTIFACT_DIGEST = "sha256:1f0c9d2e"
RUNTIME_RECIPE_DIGEST = "sha256:ab77e410"
RECIPE_ID = "llama-gguf-q4km-rocm"
RECIPE_VERSION = "3"
CATALOGUE_VERSION = "2026.07.1"
START_LAYER = 12
END_LAYER = 24
EFFECTIVE_START_LAYER = 16
HIDDEN_SIZE = 8
# A payload big enough to force more than one fragment at the bound below, so
# the vector actually exercises reassembly rather than the one-fragment path.
FRAGMENT_BYTES = 64
TENSOR_SHAPE = [1, TOKEN_COUNT, HIDDEN_SIZE]
def canonical_payload() -> bytes:
"""Deterministic bfloat16-sized payload for the canonical tensor."""
total = codec.expected_bytes(TENSOR_SHAPE, pb.DTYPE_BFLOAT16)
return bytes((i * 7 + 11) % 256 for i in range(total))
def canonical_session_request() -> pb.SessionRequest:
"""The canonical prefill chunk carried on a session stream."""
tensor = codec.encode_tensor(
codec.HIDDEN_STATES,
canonical_payload(),
TENSOR_SHAPE,
pb.DTYPE_BFLOAT16,
max_fragment_bytes=FRAGMENT_BYTES,
)
envelope = pb.Envelope(
schema_version=pb.SCHEMA_VERSION_1,
work_id=WORK_ID,
route_session_id=ROUTE_SESSION_ID,
route_epoch=ROUTE_EPOCH,
fingerprint=pb.Fingerprint(
model_artifact_digest=MODEL_ARTIFACT_DIGEST,
runtime_recipe_digest=RUNTIME_RECIPE_DIGEST,
recipe_id=RECIPE_ID,
recipe_version=RECIPE_VERSION,
catalogue_version=CATALOGUE_VERSION,
),
shard_range=pb.ShardRange(
start_layer=START_LAYER,
end_layer=END_LAYER,
effective_start_layer=EFFECTIVE_START_LAYER,
),
phase=pb.PHASE_PREFILL,
position=pb.PositionSpan(
first_position=FIRST_POSITION, token_count=TOKEN_COUNT
),
idempotency_step=IDEMPOTENCY_STEP,
cache_expectation=pb.CacheExpectation(
mode=pb.CACHE_MODE_PREFILL, expected_past_len=EXPECTED_PAST_LEN
),
deadline_unix_nanos=DEADLINE_UNIX_NANOS,
chunk=pb.ChunkInfo(chunk_index=1, chunk_count=3, final_chunk=False),
)
return pb.SessionRequest(
chunk=pb.ActivationChunk(
envelope=envelope, bundle=codec.encode_bundle([tensor])
)
)
def canonical_capability_report() -> pb.CapabilityReport:
"""The canonical capability report a worker answers admission with."""
return pb.CapabilityReport(
schema_version=pb.SCHEMA_VERSION_1,
fingerprint=pb.Fingerprint(
model_artifact_digest=MODEL_ARTIFACT_DIGEST,
runtime_recipe_digest=RUNTIME_RECIPE_DIGEST,
recipe_id=RECIPE_ID,
recipe_version=RECIPE_VERSION,
catalogue_version=CATALOGUE_VERSION,
),
shard_range=pb.ShardRange(
start_layer=START_LAYER,
end_layer=END_LAYER,
effective_start_layer=EFFECTIVE_START_LAYER,
),
backend="rocm",
device="gfx1151",
validated=True,
max_concurrent_sessions=4,
max_context_tokens=8192,
flow_control=codec.default_flow_control(),
accepted_compression=[pb.COMPRESSION_NONE, pb.COMPRESSION_ZSTD],
supported_schema_versions=[pb.SCHEMA_VERSION_1],
validated_at_unix_nanos=DEADLINE_UNIX_NANOS,
)
def serialize(message) -> bytes:
"""Serialize deterministically, so committed golden bytes are stable."""
return message.SerializeToString(deterministic=True)