"""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 hashlib 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] | str]: return { "requests": [r.hex() for r in self.requests], "responses": [r.hex() for r in self.responses], "requests_sha256": hashlib.sha256(b"".join(self.requests)).hexdigest(), "responses_sha256": hashlib.sha256(b"".join(self.responses)).hexdigest(), } 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()