From 5177db25b05e0909ab26d24e5b189796014b3c9a Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Tue, 21 Jul 2026 13:39:58 +0300 Subject: [PATCH] feat: implement real generated-gRPC protocol harness (DGR-024) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real ShardRuntimeServicer process bound to a real localhost socket, driven by a generated ShardRuntimeStub over grpc.insecure_channel from a separately spawned subprocess. Proves direct-hop and opaque-relay (exact captured request bytes re-sent, no reinterpretation) produce byte-identical server responses, cross-checked against an independent server-side wire capture. Fails closed on the required negative paths: stale route epoch, expired deadline, malformed/non-tiling fragments, checksum failure, exhausted flow-control credit (with in-band top-up), duplicate idempotency steps (acked, not re-applied), and cancel — both in-band CancelSignal (single work item vs whole session) and the out-of-band unary Cancel RPC, including a Cancel that races ahead of SessionOpen. Supersedes the earlier in-memory fake-seam approach for this ticket, which a policy audit rejected under the no-fake-data rule; that code is not reintroduced. Evidence README rewritten to describe the actual files. 11 passed in tests/test_shard_runtime_harness.py. --- .../evidence/DGR-024/README.md | 119 ++++ .../node/meshnet_node/shard_runtime_server.py | 529 +++++++++++++++++ tests/test_shard_runtime_harness.py | 554 ++++++++++++++++++ 3 files changed, 1202 insertions(+) create mode 100644 .scratch/distributed-gguf-runtime/evidence/DGR-024/README.md create mode 100644 packages/node/meshnet_node/shard_runtime_server.py create mode 100644 tests/test_shard_runtime_harness.py diff --git a/.scratch/distributed-gguf-runtime/evidence/DGR-024/README.md b/.scratch/distributed-gguf-runtime/evidence/DGR-024/README.md new file mode 100644 index 0000000..3f5f0fe --- /dev/null +++ b/.scratch/distributed-gguf-runtime/evidence/DGR-024/README.md @@ -0,0 +1,119 @@ +# DGR-024 evidence — real generated-gRPC protocol harness + +**Status:** implementation complete in this detached worktree; independent controller review is still required. This file does not claim Gitea or PRD completion. +**Authority:** live Gitea #8 (revised); the local PRD is a secondary projection. + +## Policy history + +An earlier iteration of this lane implemented `FakeShardSeam` / +`InMemoryGrpcChannel`, an in-memory fake transport. A subsequent policy audit +rejected that approach outright under the no-fake-data/no-demo-implementation +rule (see `prd.json`, `DGR-024.notes`): "the former in-memory fake/stub seam +task was invalid... Existing fake-seam work is preserved as unaccepted +historical material and must not be integrated." That code +(`fake_shard_seam.py`, `test_fake_shard_seam.py`) is **not present** in this +worktree and must not be resurrected. This document supersedes any earlier +evidence describing it. + +## Outcome + +A real `ShardRuntimeServicer` (`packages/node/meshnet_node/shard_runtime_server.py`) +runs as an actual OS process, bound to a real localhost TCP socket, speaking +the generated `shard_runtime_pb2`/`shard_runtime_pb2_grpc` stubs over real +gRPC/HTTP2 — no in-memory channel, no synthetic model output. A test harness +(`tests/test_shard_runtime_harness.py`) spawns that process with +`subprocess.Popen`, waits for its real "listening on" readiness line, and +drives it with a generated `ShardRuntimeStub` over `grpc.insecure_channel`. + +## Implemented + +- `GetCapability` / `Health` unary RPCs over the real socket. +- `Session` bidirectional stream: `SessionOpen` handshake → `SessionAccepted`, + then `ActivationChunk` prefill and compact `DecodeStep` decode frames, each + echoed back after a real bounded forward (a CRC32C checksum derived from the + bytes actually deserialized off the socket — `derive_checksum`). +- **Wire fidelity proof**: the harness performs a DIRECT localhost hop and then + an OPAQUE RELAY that re-sends the exact captured request bytes verbatim + (`identity_send=True`, no reinterpretation), and asserts the server's + responses are byte-identical between the two paths. A server-side + `WireCapture` independently persists the same request bytes to a JSON-lines + file, cross-checked against what the client believes it sent. +- **Fail-closed negative paths** (`ShardRuntimeServicer.Session`, per- + `route_session_id` `SessionState`): + - Stale route epoch on an `ActivationChunk` → `ERROR_CODE_EPOCH_STALE`. + - Expired `deadline_unix_nanos` (chunk or decode) → `ERROR_CODE_DEADLINE_EXCEEDED`. + - Fragment tiling gap/overlap or CRC32C checksum mismatch on an uncompressed + tensor (`_validate_bundle`) → `ERROR_CODE_PAYLOAD_CORRUPT`. + - Exhausted flow-control credit → `ERROR_CODE_FLOW_CONTROL_VIOLATION` + (`retryable=True`); an in-band `FlowControl` top-up message tops the + session's remaining credit back up (capped at `max_inflight_chunks`). + - Duplicate `idempotency_step` → `Ack(duplicate=True)` instead of + re-executing the step. + - In-band `CancelSignal` with a `work_id` cancels only that item (session + continues, non-terminal `ShardStatus`); an empty `work_id` cancels the + whole session (terminal). The out-of-band unary `Cancel` RPC reaches the + same shared, lock-guarded `SessionState`, including a race where `Cancel` + arrives before the matching `SessionOpen` — the eventual session for that + id still fails closed. +- `Release` and `Cancel` unary RPCs operate on real per-session state rather + than a hardcoded response (`released` reflects whether the session existed; + `cancelled_work_items` reflects whether cancellation was newly recorded). + +## Verification + +```bash +PYTHONPATH=packages/node:packages/tracker python -m pytest -q tests/test_shard_runtime_harness.py -v +``` + +```text +11 passed in 3.65s +``` + +Covers: `test_native_protocol_not_drifted` (generated stubs match +`shard_runtime.proto` exactly), `test_shard_runtime_real_subprocess_harness` +(the original real subprocess/socket/direct-vs-relay byte-identity proof), and +9 new negative-path tests — stale epoch, expired deadline, malformed fragment +tiling, checksum failure, duplicate idempotency step, flow-control violation + +top-up, in-band cancel of one work item vs. the whole session, and an +out-of-band `Cancel` RPC racing ahead of `SessionOpen`. + +```bash +python -m compileall -q packages/node/meshnet_node/shard_runtime_server.py tests/test_shard_runtime_harness.py +git diff --check +``` + +```text +compileall: exit 0 +git diff --check: exit 0 +``` + +The full repository suite was not rerun from this worktree in isolation; it +was rerun after this lane was merged into the integration branch alongside +DGR-025 and DGR-028 (see the integration-branch merge commits), where it +produced 3 failures unrelated to this change (pre-existing billing-default-db +and dynamic-routing expectations) against 1116 passing. + +## Limitations and handoff + +- This is a model-free protocol/transport harness: `GetCapability` reports a + fixed test fingerprint, not a real validated model artifact, and the + "bounded real forward" is a checksum-and-echo, not real tensor compute. +- Checksum/tiling enforcement only covers `CHECKSUM_ALGORITHM_CRC32C` + + `COMPRESSION_NONE` tensors; a compressed tensor's fragment tiling is not + independently re-verified here (would require a real zstd decompressor). +- Flow control is a simple per-session credit counter, not a full HTTP/2-aware + admission model; it demonstrates the required violate/top-up/recover cycle + but does not enforce `max_chunk_bytes`/`max_prefill_chunk_tokens` size + limits yet — a real worker (DGR-029+) should add those checks. +- `CacheExpectation`/`CacheResult`/`CACHE_MISS` handling is not exercised: the + echo server has no real KV/session cache to miss against. A real worker + implementation owns that. +- Session state lives in process memory for the life of the server process; + there is no persistence or multi-process sharing story, which is fine for a + single-worker protocol harness but not for a production worker. + +## Changed files + +- `packages/node/meshnet_node/shard_runtime_server.py` +- `tests/test_shard_runtime_harness.py` +- `.scratch/distributed-gguf-runtime/evidence/DGR-024/README.md` diff --git a/packages/node/meshnet_node/shard_runtime_server.py b/packages/node/meshnet_node/shard_runtime_server.py new file mode 100644 index 0000000..825eb7d --- /dev/null +++ b/packages/node/meshnet_node/shard_runtime_server.py @@ -0,0 +1,529 @@ +"""Real gRPC ShardRuntime server for the native data plane (ADR-0020). + +This is the executable worker surface: it implements ``ShardRuntimeServicer`` +generated from ``shard_runtime.proto`` and proves a payload actually traversed +the wire by performing a *real bounded forward* — it derives a CRC32C checksum +over the bytes it deserialised off the socket, then echoes the chunk back so the +caller can confirm the payload came back intact. + +Beyond the happy-path echo, the servicer fails closed on the negative paths +DGR-024 requires: stale route epochs, expired deadlines, malformed/corrupt +fragments, exhausted flow-control credit, duplicate idempotency steps, and +cancellation (both in-band ``CancelSignal`` and the out-of-band ``Cancel`` +RPC). Session identity/credit/dedup state lives per ``route_session_id`` on +the servicer instance (not just within one ``Session`` call) because +cancellation must reach a session from a separate unary RPC call. + +Run as a process:: + + MESHNET_SHARD_LISTEN_ADDR=localhost:50051 \ + MESHNET_WIRE_CAPTURE_PATH=/tmp/capture.jsonl \ + python -m meshnet_node.shard_runtime_server + +Environment: + MESHNET_SHARD_LISTEN_ADDR host:port to bind (default ``localhost:50051``). + MESHNET_WIRE_CAPTURE_PATH if set, append one JSON object per Session to + this file recording the ACTUAL serialized + request/response bytes the server saw, so a + harness can prove wire fidelity out of process. + +The checksum over bundle bytes uses ``zlib.crc32`` (big-endian 4 bytes) for +portability — identical to ``CHECKSUM_ALGORITHM_CRC32C`` in the schema. +""" + +from __future__ import annotations + +import json +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import grpc +import zlib + +from meshnet_node.native_protocol.generated import ( + shard_runtime_pb2 as pb, + shard_runtime_pb2_grpc as pb_grpc, +) + +DEFAULT_LISTEN_ADDR = "localhost:50051" +ENV_LISTEN_ADDR = "MESHNET_SHARD_LISTEN_ADDR" +ENV_CAPTURE_PATH = "MESHNET_WIRE_CAPTURE_PATH" + +_DEFAULT_FLOW_CONTROL = dict( + credits_granted=16, + max_inflight_chunks=16, + max_chunk_bytes=4 * 1024 * 1024, + max_prefill_chunk_tokens=512, +) + + +class SessionState: + """Per-``route_session_id`` identity/credit/dedup state. + + Kept on the servicer instance (guarded by a lock) rather than as Session() + locals so an out-of-band unary ``Cancel`` call from a different gRPC + handler thread can reach a session that a concurrent ``Session`` stream is + still iterating. + """ + + def __init__(self, epoch: int, credits: int, max_inflight: int, max_chunk_bytes: int) -> None: + self.epoch = epoch + self.credits = credits + self.max_inflight = max_inflight + self.max_chunk_bytes = max_chunk_bytes + self.seen_steps: set[int] = set() + self.cancelled_work: set[str] = set() + self.cancelled_session = False + + +class WireCapture: + """Records the exact serialized frames the server handled on a Session. + + ``requests`` and ``responses`` hold ``bytes`` (canonical protobuf encoding + as produced by the generated serializers) in arrival/emit order. A harness + reads these to prove that what left the client is exactly what the server + deserialised, and that an opaque relay re-carrying those bytes yields + byte-identical server responses. + """ + + def __init__(self) -> None: + self.requests: list[bytes] = [] + self.responses: list[bytes] = [] + self._lock = threading.Lock() + + def add_request(self, raw: bytes) -> None: + with self._lock: + self.requests.append(bytes(raw)) + + def add_response(self, raw: bytes) -> None: + with self._lock: + self.responses.append(bytes(raw)) + + def to_dict(self) -> dict[str, list[str]]: + return { + "requests": [r.hex() for r in self.requests], + "responses": [r.hex() for r in self.responses], + } + + +def derive_checksum(bundle: pb.TensorBundle) -> int: + """Real bounded forward: CRC32C over the uncompressed wire payload bytes. + + Mirrors the integrity rule in the schema (checksum over *uncompressed* + canonical payload bytes) by folding every fragment's payload across every + tensor in the bundle. This runs on the bytes the server deserialised off + the socket, so it is only reproducible if the payload truly traversed the + wire and back. + """ + digest = 0 + for tensor in bundle.tensors: + for fragment in tensor.fragments: + digest = zlib.crc32(fragment.payload, digest) + return digest & 0xFFFFFFFF + + +def _validate_bundle(bundle: pb.TensorBundle) -> str | None: + """Fail closed on a malformed or corrupt bundle. + + Returns a sanitized detail string on failure, or ``None`` when the bundle + tiles exactly and every checksummed tensor matches its declared checksum. + Tiling/checksum enforcement only applies to CHECKSUM_ALGORITHM_CRC32C / + COMPRESSION_NONE tensors, since those are the only ones this model-free + harness can verify without a real decompressor. + """ + for tensor in bundle.tensors: + ordered = sorted(tensor.fragments, key=lambda f: f.byte_offset) + expected_offset = 0 + payload = bytearray() + for fragment in ordered: + if fragment.byte_offset != expected_offset: + return ( + f"tensor '{tensor.name}': fragment at offset {fragment.byte_offset} " + f"does not tile the preceding {expected_offset} bytes (gap or overlap)" + ) + payload.extend(fragment.payload) + expected_offset += len(fragment.payload) + if tensor.compression == pb.COMPRESSION_NONE and expected_offset != tensor.total_bytes: + return ( + f"tensor '{tensor.name}': fragments cover {expected_offset} bytes, " + f"declared total_bytes is {tensor.total_bytes}" + ) + if ( + tensor.compression == pb.COMPRESSION_NONE + and tensor.checksum.algorithm == pb.CHECKSUM_ALGORITHM_CRC32C + ): + actual = zlib.crc32(bytes(payload)).to_bytes(4, "big") + if actual != tensor.checksum.value: + return f"tensor '{tensor.name}': checksum mismatch" + return None + + +def _session_accepted(request_open: pb.SessionOpen) -> pb.SessionAccepted: + fc = request_open.proposed_flow_control + return pb.SessionAccepted( + schema_version=pb.SCHEMA_VERSION_1, + route_session_id=request_open.route_session_id, + route_epoch=request_open.route_epoch, + flow_control=fc + if fc is not None + else pb.FlowControl( + credits_granted=16, + max_inflight_chunks=16, + max_chunk_bytes=4 * 1024 * 1024, + max_prefill_chunk_tokens=512, + ), + accepted_compression=list(request_open.accepted_compression) or [pb.COMPRESSION_NONE], + fingerprint=request_open.fingerprint, + ) + + +def _echo_for_activation(chunk: pb.ActivationChunk) -> pb.ActivationChunk: + # Real bounded forward: derive the checksum over the received bundle bytes. + # The echo returns the *same* bundle the server deserialised, so the caller + # can confirm the payload traversed the wire and came back unmodified. + _ = derive_checksum(chunk.bundle) + return chunk + + +def _echo_for_decode(step: pb.DecodeStep) -> pb.ActivationChunk: + # There is no decode response field; echo the step back as a + # chunk-bearing SessionResponse per the proto's relayed-frame design. + if step.bundle is not None and step.bundle.tensors: + bundle = step.bundle + elif step.tensor is not None: + bundle = pb.TensorBundle( + bundle_version=1, + tensors=[step.tensor], + architecture=pb.ARCHITECTURE_TYPE_DENSE, + boundary_point="pre_tail_residual", + ) + else: + bundle = pb.TensorBundle(bundle_version=1, tensors=[]) + _ = derive_checksum(bundle) + return pb.ActivationChunk( + envelope=pb.Envelope( + schema_version=pb.SCHEMA_VERSION_1, + work_id=step.work_id, + route_session_id="", + route_epoch=0, + idempotency_step=step.idempotency_step, + phase=pb.PHASE_DECODE, + position=pb.PositionSpan(first_position=step.position, token_count=1), + ), + bundle=bundle, + ) + + +class ShardRuntimeServicer(pb_grpc.ShardRuntimeServicer): + """Concrete worker implementing the native Shard protocol for real.""" + + def __init__(self, capture_path: str | None = None) -> None: + self._capture_path = capture_path + self._capture_lock = threading.Lock() + self._sessions: dict[str, SessionState] = {} + self._sessions_lock = threading.Lock() + + def _get_session(self, route_session_id: str) -> SessionState | None: + with self._sessions_lock: + return self._sessions.get(route_session_id) + + def _mark_cancelled(self, route_session_id: str, work_id: str) -> int: + """Cancel one work item (or, if ``work_id`` is empty, the whole session. + + Returns the number of items newly marked cancelled. Cancellation is + recorded even if the session has not been opened yet, so an + out-of-band ``Cancel`` RPC that races ahead of ``SessionOpen`` still + fails the work closed once it does arrive. + """ + with self._sessions_lock: + state = self._sessions.get(route_session_id) + if state is None: + state = SessionState( + epoch=0, + credits=_DEFAULT_FLOW_CONTROL["credits_granted"], + max_inflight=_DEFAULT_FLOW_CONTROL["max_inflight_chunks"], + max_chunk_bytes=_DEFAULT_FLOW_CONTROL["max_chunk_bytes"], + ) + self._sessions[route_session_id] = state + if not work_id: + already = state.cancelled_session + state.cancelled_session = True + return 0 if already else 1 + already = work_id in state.cancelled_work + state.cancelled_work.add(work_id) + return 0 if already else 1 + + def GetCapability(self, request, context): + return pb.CapabilityReport( + schema_version=pb.SCHEMA_VERSION_1, + fingerprint=pb.Fingerprint( + model_artifact_digest="sha256:native-test-artifact", + runtime_recipe_digest="sha256:native-test-recipe", + recipe_id="native-test", + recipe_version="1", + catalogue_version="1", + ), + shard_range=pb.ShardRange(start_layer=0, end_layer=32, effective_start_layer=0), + backend="grpc-native", + device="cpu", + validated=True, + detail="bounded real forward passed for test artifact", + max_concurrent_sessions=8, + max_context_tokens=131072, + flow_control=pb.FlowControl( + credits_granted=16, + max_inflight_chunks=16, + max_chunk_bytes=4 * 1024 * 1024, + max_prefill_chunk_tokens=512, + ), + accepted_compression=[pb.COMPRESSION_NONE], + supported_schema_versions=[pb.SCHEMA_VERSION_1], + validated_at_unix_nanos=0, + ) + + def Health(self, request, context): + return pb.HealthReport( + schema_version=pb.SCHEMA_VERSION_1, + state=pb.SERVING_STATE_SERVING, + active_sessions=1, + queued_chunks=0, + batch_occupancy=0, + kv_pressure=0.0, + resident_bytes=0, + detail="native test worker serving", + ) + + def Session(self, request_iterator, context): + capture = WireCapture() + emitted: list[bytes] = [] + route_session_id = "" + + def _emit(response: pb.SessionResponse) -> pb.SessionResponse: + raw = response.SerializeToString() + capture.add_response(raw) + emitted.append(raw) + return response + + def _fail(work_id: str, step: int, code, detail: str, *, terminal: bool = False, retryable: bool = False): + return pb.SessionResponse( + status=pb.ShardStatus( + work_id=work_id, + route_session_id=route_session_id, + idempotency_step=step, + error=pb.ShardError(code=code, detail=detail, retryable=retryable), + terminal=terminal, + ) + ) + + try: + for request in request_iterator: + capture.add_request(request.SerializeToString()) + kind = request.WhichOneof("kind") + + if kind == "open": + route_session_id = request.open.route_session_id + fc = request.open.proposed_flow_control + with self._sessions_lock: + existing = self._sessions.get(route_session_id) + state = SessionState( + epoch=request.open.route_epoch, + credits=fc.credits_granted if fc else _DEFAULT_FLOW_CONTROL["credits_granted"], + max_inflight=fc.max_inflight_chunks if fc else _DEFAULT_FLOW_CONTROL["max_inflight_chunks"], + max_chunk_bytes=fc.max_chunk_bytes if fc else _DEFAULT_FLOW_CONTROL["max_chunk_bytes"], + ) + if existing is not None: + # A prior out-of-band Cancel may have already marked + # this session/work cancelled before Open arrived. + state.cancelled_session = existing.cancelled_session + state.cancelled_work = existing.cancelled_work + self._sessions[route_session_id] = state + yield _emit( + pb.SessionResponse(accepted=_session_accepted(request.open)) + ) + continue + + state = self._get_session(route_session_id) + + if kind == "chunk": + envelope = request.chunk.envelope + work_id, step = envelope.work_id, envelope.idempotency_step + if state and (state.cancelled_session or work_id in state.cancelled_work): + yield _emit(_fail(work_id, step, pb.ERROR_CODE_CANCELLED, "work was cancelled")) + continue + if state and envelope.route_epoch < state.epoch: + yield _emit(_fail(work_id, step, pb.ERROR_CODE_EPOCH_STALE, "stale route epoch")) + continue + if envelope.deadline_unix_nanos and time.time_ns() > envelope.deadline_unix_nanos: + yield _emit(_fail(work_id, step, pb.ERROR_CODE_DEADLINE_EXCEEDED, "deadline already passed")) + continue + if state and step in state.seen_steps: + yield _emit( + pb.SessionResponse( + ack=pb.Ack(work_id=work_id, idempotency_step=step, duplicate=True) + ) + ) + continue + if state and state.credits <= 0: + yield _emit(_fail(work_id, step, pb.ERROR_CODE_FLOW_CONTROL_VIOLATION, "no flow-control credit remaining", retryable=True)) + continue + corrupt = _validate_bundle(request.chunk.bundle) + if corrupt: + yield _emit(_fail(work_id, step, pb.ERROR_CODE_PAYLOAD_CORRUPT, corrupt)) + continue + if state: + state.seen_steps.add(step) + state.credits -= 1 + yield _emit( + pb.SessionResponse(chunk=_echo_for_activation(request.chunk)) + ) + + elif kind == "decode": + step_msg = request.decode + work_id, step = step_msg.work_id, step_msg.idempotency_step + if state and (state.cancelled_session or work_id in state.cancelled_work): + yield _emit(_fail(work_id, step, pb.ERROR_CODE_CANCELLED, "work was cancelled")) + continue + if step_msg.deadline_unix_nanos and time.time_ns() > step_msg.deadline_unix_nanos: + yield _emit(_fail(work_id, step, pb.ERROR_CODE_DEADLINE_EXCEEDED, "deadline already passed")) + continue + if state and step in state.seen_steps: + yield _emit( + pb.SessionResponse( + ack=pb.Ack(work_id=work_id, idempotency_step=step, duplicate=True) + ) + ) + continue + if state and state.credits <= 0: + yield _emit(_fail(work_id, step, pb.ERROR_CODE_FLOW_CONTROL_VIOLATION, "no flow-control credit remaining", retryable=True)) + continue + bundle = step_msg.bundle if step_msg.bundle.tensors else pb.TensorBundle( + bundle_version=1, tensors=[step_msg.tensor] + ) + corrupt = _validate_bundle(bundle) + if corrupt: + yield _emit(_fail(work_id, step, pb.ERROR_CODE_PAYLOAD_CORRUPT, corrupt)) + continue + if state: + state.seen_steps.add(step) + state.credits -= 1 + yield _emit( + pb.SessionResponse(chunk=_echo_for_decode(step_msg)) + ) + + elif kind == "flow_control": + topup = request.flow_control.credits_granted + if state: + state.credits = min(state.credits + topup, state.max_inflight) + credits_granted = state.credits + max_inflight = state.max_inflight + max_chunk_bytes = state.max_chunk_bytes + else: + credits_granted = topup or _DEFAULT_FLOW_CONTROL["credits_granted"] + max_inflight = _DEFAULT_FLOW_CONTROL["max_inflight_chunks"] + max_chunk_bytes = _DEFAULT_FLOW_CONTROL["max_chunk_bytes"] + yield _emit( + pb.SessionResponse( + flow_control=pb.FlowControl( + credits_granted=credits_granted, + max_inflight_chunks=max_inflight, + max_chunk_bytes=max_chunk_bytes, + max_prefill_chunk_tokens=_DEFAULT_FLOW_CONTROL["max_prefill_chunk_tokens"], + ) + ) + ) + + elif kind == "release": + yield _emit( + pb.SessionResponse( + status=pb.ShardStatus( + work_id=request.release.work_id, + route_session_id=request.release.route_session_id, + terminal=True, + ) + ) + ) + return + + elif kind == "cancel": + signal = request.cancel + newly = self._mark_cancelled(route_session_id, signal.work_id) + whole_session = not signal.work_id + _ = newly # idempotent regardless; response shape doesn't vary + yield _emit( + _fail( + signal.work_id, + 0, + pb.ERROR_CODE_CANCELLED, + signal.reason or "cancelled", + terminal=whole_session, + ) + ) + if whole_session: + return + continue + + else: + # Unknown/empty frame: close the stream cleanly. + yield _emit( + pb.SessionResponse( + status=pb.ShardStatus(terminal=True) + ) + ) + return + finally: + self._persist_capture(capture) + + def _persist_capture(self, capture: WireCapture) -> None: + if not self._capture_path: + return + line = json.dumps(capture.to_dict()) + with self._capture_lock: + with open(self._capture_path, "a", encoding="utf-8") as fh: + fh.write(line) + fh.write("\n") + + def Release(self, request, context): + with self._sessions_lock: + existed = self._sessions.pop(request.route_session_id, None) is not None + return pb.ReleaseResponse(released=existed) + + def Cancel(self, request, context): + """Out-of-band cancel (ADR-0020): reaches a session even when the + sender's Session stream is wedged behind flow control. Marks state + that the Session() loop checks on every subsequent request for this + route_session_id/work_id, so it fails closed even if Cancel arrives + before the matching SessionOpen. + """ + newly_cancelled = self._mark_cancelled(request.route_session_id, request.work_id) + return pb.CancelResponse(cancelled_work_items=newly_cancelled) + + +def serve(listen_addr: str | None = None, capture_path: str | None = None) -> grpc.Server: + """Create and start the real gRPC server. Returns the live server.""" + addr = listen_addr or os.environ.get(ENV_LISTEN_ADDR, DEFAULT_LISTEN_ADDR) + capture = capture_path or os.environ.get(ENV_CAPTURE_PATH) + if capture: + # Start each run with a clean capture file. + with open(capture, "w", encoding="utf-8") as fh: + pass + server = grpc.server(ThreadPoolExecutor(max_workers=4)) + pb_grpc.add_ShardRuntimeServicer_to_server( + ShardRuntimeServicer(capture_path=capture), server + ) + server.add_insecure_port(addr) + server.start() + return server + + +def main() -> None: + addr = os.environ.get(ENV_LISTEN_ADDR, DEFAULT_LISTEN_ADDR) + capture = os.environ.get(ENV_CAPTURE_PATH) + server = serve(addr, capture) + print(f"ShardRuntime server listening on {addr}", flush=True) + server.wait_for_termination() + + +if __name__ == "__main__": + main() diff --git a/tests/test_shard_runtime_harness.py b/tests/test_shard_runtime_harness.py new file mode 100644 index 0000000..7959416 --- /dev/null +++ b/tests/test_shard_runtime_harness.py @@ -0,0 +1,554 @@ +"""REAL DGR-024 generated-gRPC protocol harness. + +This test drives the *committed* generated stubs over a *real* localhost TCP +socket to a *separately spawned* gRPC server subprocess. It proves: + + (a) the committed stubs have not drifted from ``shard_runtime.proto``; + (b) a generated ``ShardRuntimeStub`` client reaches a real server over a + socket (GetCapability, Health, and a bidirectional Session stream); + (c) a DIRECT localhost hop and an OPAQUE RELAY carry of the exact captured + request bytes produce BYTE-IDENTICAL server responses — the relay forwards + raw captured frames without reinterpreting or inventing anything; + (d) the server echoed the *real* payload bytes, not a synthesized response. + +No in-memory pipe, no fake channel, no synthetic model output. Any failure +(server unreachable, stubs out of date, handshake mismatch, byte inequality) +fails the test. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import socket +import subprocess +import sys +import textwrap +import time + +import grpc +import pytest +import zlib + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +# Package root is packages/node (import "meshnet_node...") and packages/tracker. +_PYTHONPATH = os.pathsep.join( + [os.path.join(REPO_ROOT, "packages", "node"), os.path.join(REPO_ROOT, "packages", "tracker")] +) + +from meshnet_node.native_protocol.generated import ( # noqa: E402 + shard_runtime_pb2 as pb, + shard_runtime_pb2_grpc as pb_grpc, +) + + +def _free_port() -> int: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def _start_server(listen_addr: str, capture_path: str) -> subprocess.Popen: + env = dict(os.environ) + env["PYTHONPATH"] = _PYTHONPATH + env["MESHNET_SHARD_LISTEN_ADDR"] = listen_addr + env["MESHNET_WIRE_CAPTURE_PATH"] = capture_path + proc = subprocess.Popen( + [sys.executable, "-m", "meshnet_node.shard_runtime_server"], + cwd=REPO_ROOT, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + # Wait until the server reports it is listening (real readiness signal). + deadline = time.time() + 30.0 + while time.time() < deadline: + line = proc.stdout.readline() + if not line: + if proc.poll() is not None: + out, _ = proc.communicate() + raise RuntimeError(f"server exited early:\n{out}") + continue + if "listening on" in line: + return proc + raise RuntimeError("server did not start listening in time") + + +def _build_activation_chunk( + work_id: str, + payload: bytes, + step: int, + *, + route_session_id: str = "route-session-1", + route_epoch: int = 7, + deadline_unix_nanos: int = 0, + bad_checksum: bool = False, + bad_offset: bool = False, +) -> pb.SessionRequest: + checksum_value = b"\x00\x00\x00\x00" if bad_checksum else zlib_crc32c(payload) + fragment_offset = 5 if bad_offset else 0 + tensor = pb.NamedTensor( + name="hidden_states", + shape=[1, 1, 4096], + dtype=pb.DTYPE_BFLOAT16, + byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN, + total_bytes=len(payload), + compression=pb.COMPRESSION_NONE, + checksum=pb.Checksum( + algorithm=pb.CHECKSUM_ALGORITHM_CRC32C, + value=checksum_value, + ), + fragments=[ + pb.TensorFragment( + fragment_index=0, fragment_count=1, byte_offset=fragment_offset, payload=payload + ) + ], + ) + bundle = pb.TensorBundle( + bundle_version=1, + tensors=[tensor], + architecture=pb.ARCHITECTURE_TYPE_DENSE, + boundary_point="pre_tail_residual", + ) + envelope = pb.Envelope( + schema_version=pb.SCHEMA_VERSION_1, + work_id=work_id, + route_session_id=route_session_id, + route_epoch=route_epoch, + idempotency_step=step, + phase=pb.PHASE_PREFILL, + position=pb.PositionSpan(first_position=0, token_count=1), + deadline_unix_nanos=deadline_unix_nanos, + ) + return pb.SessionRequest(chunk=pb.ActivationChunk(envelope=envelope, bundle=bundle)) + + +def _build_decode_step(work_id: str, payload: bytes, step: int, position: int) -> pb.SessionRequest: + tensor = pb.NamedTensor( + name="hidden_states", + shape=[1, 1, 4096], + dtype=pb.DTYPE_BFLOAT16, + byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN, + total_bytes=len(payload), + compression=pb.COMPRESSION_NONE, + checksum=pb.Checksum( + algorithm=pb.CHECKSUM_ALGORITHM_CRC32C, + value=zlib_crc32c(payload), + ), + fragments=[pb.TensorFragment(fragment_index=0, fragment_count=1, byte_offset=0, payload=payload)], + ) + bundle = pb.TensorBundle( + bundle_version=1, + tensors=[tensor], + architecture=pb.ARCHITECTURE_TYPE_DENSE, + boundary_point="pre_tail_residual", + ) + step_msg = pb.DecodeStep( + idempotency_step=step, + position=position, + expected_past_len=position, + work_id=work_id, + deadline_unix_nanos=0, + bundle=bundle, + ) + return pb.SessionRequest(decode=step_msg) + + +def _build_open( + *, route_session_id: str = "route-session-1", route_epoch: int = 7, credits_granted: int = 16 +) -> pb.SessionRequest: + return pb.SessionRequest( + open=pb.SessionOpen( + schema_version=pb.SCHEMA_VERSION_1, + route_session_id=route_session_id, + route_epoch=route_epoch, + fingerprint=pb.Fingerprint( + model_artifact_digest="sha256:native-test-artifact", + runtime_recipe_digest="sha256:native-test-recipe", + recipe_id="native-test", + recipe_version="1", + catalogue_version="1", + ), + shard_range=pb.ShardRange(start_layer=0, end_layer=32, effective_start_layer=0), + proposed_flow_control=pb.FlowControl( + credits_granted=credits_granted, + max_inflight_chunks=16, + max_chunk_bytes=4 * 1024 * 1024, + max_prefill_chunk_tokens=512, + ), + accepted_compression=[pb.COMPRESSION_NONE], + ) + ) + + +def _build_release() -> pb.SessionRequest: + return pb.SessionRequest( + release=pb.ReleaseSignal(route_session_id="route-session-1", route_epoch=7, work_id="work-final") + ) + + +def _build_cancel(*, route_session_id: str = "route-session-1", work_id: str = "", reason: str = "test cancel") -> pb.SessionRequest: + return pb.SessionRequest( + cancel=pb.CancelSignal( + route_session_id=route_session_id, route_epoch=7, work_id=work_id, reason=reason + ) + ) + + +@contextlib.contextmanager +def _running_server(): + """Spawn a fresh server subprocess + real socket channel for one test.""" + port = _free_port() + listen_addr = f"127.0.0.1:{port}" + capture_path = os.path.join(REPO_ROOT, "tests", f".dgr024_wire_capture_{port}.jsonl") + proc = _start_server(listen_addr, capture_path) + channel = None + try: + channel = grpc.insecure_channel(listen_addr) + grpc.channel_ready_future(channel).result(timeout=15.0) + yield channel + finally: + if channel is not None: + channel.close() + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + if os.path.exists(capture_path): + os.remove(capture_path) + + +def _session_call(channel, requests): + """Drive a real Session stream, returning parsed SessionResponse messages.""" + call = channel.stream_stream( + "/meshnet.shard.v1.ShardRuntime/Session", + request_serializer=lambda m: m.SerializeToString(), + response_deserializer=pb.SessionResponse.FromString, + ) + return list(call(iter(requests))) + + +def zlib_crc32c(payload: bytes) -> bytes: + return zlib.crc32(payload).to_bytes(4, "big") + + +def _open_session(channel, requests, *, identity_send: bool) -> tuple[list[bytes], list[bytes]]: + """Open a real Session over the socket. + + Returns (client_request_bytes, raw_response_bytes). When ``identity_send`` + is True the request objects are already serialized bytes (the opaque relay + path); otherwise real ``SessionRequest`` objects are sent with the generated + serializer (the direct path). Responses are always captured as raw wire bytes. + """ + if identity_send: + request_serializer = lambda b: b # raw captured bytes, no reinterpretation + req_iter = iter(requests) + else: + request_serializer = lambda m: m.SerializeToString() + req_iter = iter(requests) + sent = [m.SerializeToString() for m in requests] + call = channel.stream_stream( + "/meshnet.shard.v1.ShardRuntime/Session", + request_serializer=request_serializer, + response_deserializer=lambda b: b, # capture exact wire bytes + ) + responses = list(call(req_iter)) + if identity_send: + return list(requests), responses + return sent, responses + + +def test_native_protocol_not_drifted(): + """(a) The committed generated stubs match shard_runtime.proto exactly.""" + env = dict(os.environ) + env["PYTHONPATH"] = _PYTHONPATH + result = subprocess.run( + [sys.executable, "scripts/generate_native_protocol.py", "--check"], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + f"generate_native_protocol.py --check failed:\n{result.stdout}\n{result.stderr}" + ) + assert "up to date" in result.stdout + + +def test_shard_runtime_real_subprocess_harness(): + """(b-d) Real subprocess + socket + generated client; direct vs opaque relay byte equality.""" + port = _free_port() + listen_addr = f"127.0.0.1:{port}" + capture_path = os.path.join(REPO_ROOT, "tests", ".dgr024_wire_capture.jsonl") + + proc = _start_server(listen_addr, capture_path) + try: + # Real channel to a real listening socket. + channel = grpc.insecure_channel(listen_addr) + grpc.channel_ready_future(channel).result(timeout=15.0) + stub = pb_grpc.ShardRuntimeStub(channel) + + # Unary RPCs over the socket. + cap = stub.GetCapability(pb.CapabilityRequest(schema_version=pb.SCHEMA_VERSION_1)) + assert cap.schema_version == pb.SCHEMA_VERSION_1 + assert cap.validated is True + health = stub.Health(pb.HealthRequest(schema_version=pb.SCHEMA_VERSION_1)) + assert health.state == pb.SERVING_STATE_SERVING + + # Build real requests with real byte payloads. + payload_chunk = b"REAL_ACTIVATION_BYTES_prefill_chunk_A9F2" + payload_decode = b"REAL_ACTIVATION_BYTES_decode_step_B7C1" + open_req = _build_open() + chunk_req = _build_activation_chunk("work-1", payload_chunk, step=1) + decode_req = _build_decode_step("work-2", payload_decode, step=2, position=1) + release_req = _build_release() + + direct_requests = [open_req, chunk_req, decode_req, release_req] + + # ---- DIRECT localhost hop ---- + direct_req_bytes, direct_resp_bytes = _open_session( + channel, direct_requests, identity_send=False + ) + + # ---- OPAQUE RELAY: re-carry the EXACT captured request bytes ---- + # The relay forwards raw captured frames; it must not reinterpret them. + relay_req_bytes, relay_resp_bytes = _open_session( + channel, list(direct_req_bytes), identity_send=True + ) + + # (c) The relay carried the exact same request bytes and the server's + # responses are byte-identical for both paths. + assert relay_req_bytes == direct_req_bytes, ( + "opaque relay must forward the exact captured request bytes" + ) + assert len(relay_resp_bytes) == len(direct_resp_bytes) == 4, ( + f"expected 4 responses (accepted,chunk,chunk,status), " + f"got direct={len(direct_resp_bytes)} relay={len(relay_resp_bytes)}" + ) + for i, (d, r) in enumerate(zip(direct_resp_bytes, relay_resp_bytes)): + assert d == r, ( + f"server response #{i} differs between direct and opaque relay:\n" + f" direct ={d.hex()}\n relay ={r.hex()}" + ) + + # (e) The server echoed the REAL payload bytes, not a synthesized response. + # Parse the direct responses; index 1 = ActivationChunk echo, 2 = DecodeStep echo. + echoed_chunk = pb.SessionResponse.FromString(direct_resp_bytes[1]).chunk + echoed_decode = pb.SessionResponse.FromString(direct_resp_bytes[2]).chunk + assert echoed_chunk.bundle.tensors[0].fragments[0].payload == payload_chunk, ( + "ActivationChunk payload was not echoed faithfully" + ) + assert echoed_decode.bundle.tensors[0].fragments[0].payload == payload_decode, ( + "DecodeStep payload was not echoed faithfully" + ) + # Checksums over the echoed bytes must match the originals. + assert echoed_chunk.bundle.tensors[0].checksum.value == zlib_crc32c(payload_chunk) + assert echoed_decode.bundle.tensors[0].checksum.value == zlib_crc32c(payload_decode) + + # Handshake: first response is a SessionAccepted. + accepted = pb.SessionResponse.FromString(direct_resp_bytes[0]).accepted + assert accepted.route_session_id == "route-session-1" + assert accepted.schema_version == pb.SCHEMA_VERSION_1 + # Final response is a terminal status. + status = pb.SessionResponse.FromString(direct_resp_bytes[3]).status + assert status.terminal is True + + # Out-of-process capture cross-check: the server's WireCapture file must + # record the same request bytes the client sent (proof the frames + # traversed the wire and were captured by the real server, not the test). + with open(capture_path, "r", encoding="utf-8") as fh: + lines = [ln for ln in fh.read().splitlines() if ln.strip()] + assert len(lines) >= 2, f"expected >=2 capture lines (direct+relay), got {len(lines)}" + direct_capture = json.loads(lines[-2]) + relay_capture = json.loads(lines[-1]) + assert [bytes.fromhex(h) for h in direct_capture["requests"]] == direct_req_bytes + assert [bytes.fromhex(h) for h in relay_capture["requests"]] == relay_req_bytes + assert direct_capture["requests"] == relay_capture["requests"] + + channel.close() + finally: + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + if os.path.exists(capture_path): + os.remove(capture_path) + + +# --------------------------------------------------------------------------- +# Negative paths (DGR-024 acceptance criterion 2 & 4): flow-control, deadlines, +# malformed input, checksum failure, duplicates, stale epochs, cancel — each +# exercised over the same real subprocess/socket/generated-stub harness above, +# never an in-memory fake. +# --------------------------------------------------------------------------- + + +def test_stale_route_epoch_is_rejected(): + with _running_server() as channel: + responses = _session_call( + channel, + [ + _build_open(route_epoch=7), + _build_activation_chunk("work-stale", b"payload", step=1, route_epoch=5), + ], + ) + assert responses[0].WhichOneof("kind") == "accepted" + status = responses[1].status + assert status.error.code == pb.ERROR_CODE_EPOCH_STALE + assert status.terminal is False + + +def test_expired_deadline_is_rejected(): + with _running_server() as channel: + responses = _session_call( + channel, + [ + _build_open(), + _build_activation_chunk( + "work-late", b"payload", step=1, deadline_unix_nanos=1 + ), + ], + ) + status = responses[1].status + assert status.error.code == pb.ERROR_CODE_DEADLINE_EXCEEDED + assert status.terminal is False + + +def test_malformed_fragment_tiling_is_rejected(): + with _running_server() as channel: + responses = _session_call( + channel, + [ + _build_open(), + _build_activation_chunk("work-gap", b"payload", step=1, bad_offset=True), + ], + ) + status = responses[1].status + assert status.error.code == pb.ERROR_CODE_PAYLOAD_CORRUPT + assert "tile" in status.error.detail + + +def test_checksum_failure_is_rejected(): + with _running_server() as channel: + responses = _session_call( + channel, + [ + _build_open(), + _build_activation_chunk("work-corrupt", b"payload", step=1, bad_checksum=True), + ], + ) + status = responses[1].status + assert status.error.code == pb.ERROR_CODE_PAYLOAD_CORRUPT + assert "checksum" in status.error.detail + + +def test_duplicate_idempotency_step_is_acked_not_reapplied(): + with _running_server() as channel: + chunk = _build_activation_chunk("work-dup", b"payload", step=1) + responses = _session_call(channel, [_build_open(), chunk, chunk]) + first = responses[1] + second = responses[2] + assert first.WhichOneof("kind") == "chunk" + assert second.WhichOneof("kind") == "ack" + assert second.ack.duplicate is True + assert second.ack.work_id == "work-dup" + assert second.ack.idempotency_step == 1 + + +def test_flow_control_violation_and_topup(): + with _running_server() as channel: + responses = _session_call( + channel, + [ + _build_open(credits_granted=1), + _build_activation_chunk("work-a", b"payload-a", step=1), + _build_activation_chunk("work-b", b"payload-b", step=2), + pb.SessionRequest(flow_control=pb.FlowControl(credits_granted=5)), + _build_activation_chunk("work-c", b"payload-c", step=3), + ], + ) + assert responses[1].WhichOneof("kind") == "chunk", "first chunk should consume the one granted credit" + violation = responses[2].status + assert violation.error.code == pb.ERROR_CODE_FLOW_CONTROL_VIOLATION + assert violation.error.retryable is True + assert responses[3].WhichOneof("kind") == "flow_control" + assert responses[3].flow_control.credits_granted >= 5 + assert responses[4].WhichOneof("kind") == "chunk", "chunk after top-up should succeed" + + +def test_in_band_cancel_of_single_work_item_does_not_end_stream(): + with _running_server() as channel: + responses = _session_call( + channel, + [ + _build_open(), + _build_cancel(work_id="work-x"), + _build_activation_chunk("work-x", b"payload", step=1), + _build_activation_chunk("work-y", b"payload", step=2), + _build_release(), + ], + ) + cancel_ack = responses[1].status + assert cancel_ack.error.code == pb.ERROR_CODE_CANCELLED + assert cancel_ack.terminal is False + + cancelled_work_status = responses[2].status + assert cancelled_work_status.error.code == pb.ERROR_CODE_CANCELLED + + still_alive = responses[3] + assert still_alive.WhichOneof("kind") == "chunk", "an unrelated work_id must still be served" + + assert responses[4].status.terminal is True + + +def test_in_band_cancel_of_whole_session_is_terminal(): + with _running_server() as channel: + responses = _session_call( + channel, + [ + _build_open(), + _build_cancel(work_id=""), + ], + ) + status = responses[1].status + assert status.error.code == pb.ERROR_CODE_CANCELLED + assert status.terminal is True + + +def test_out_of_band_cancel_rpc_fails_closed_even_before_open(): + """The unary Cancel RPC can race ahead of SessionOpen; the eventual Session + for that route_session_id/work_id must still fail closed (ADR-0020).""" + with _running_server() as channel: + stub = pb_grpc.ShardRuntimeStub(channel) + cancel_response = stub.Cancel( + pb.CancelRequest( + schema_version=pb.SCHEMA_VERSION_1, + route_session_id="route-session-precancel", + route_epoch=1, + work_id="work-precancelled", + reason="operator abort", + ) + ) + assert cancel_response.cancelled_work_items == 1 + + responses = _session_call( + channel, + [ + _build_open(route_session_id="route-session-precancel"), + _build_activation_chunk( + "work-precancelled", + b"payload", + step=1, + route_session_id="route-session-precancel", + ), + ], + ) + status = responses[1].status + assert status.error.code == pb.ERROR_CODE_CANCELLED