5177db25b0
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.
530 lines
22 KiB
Python
530 lines
22 KiB
Python
"""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()
|