diff --git a/.claude/memory/MEMORY.md b/.claude/memory/MEMORY.md index 2b77780..443a395 100644 --- a/.claude/memory/MEMORY.md +++ b/.claude/memory/MEMORY.md @@ -6,3 +6,4 @@ - **Alpha hardening** — `.scratch/alpha-hardening/` (22 issues, ADRs 0016–0019, [README](../.scratch/alpha-hardening/README.md), [handoff](../.scratch/alpha-hardening/handoff.md)) - [Alpha hardening navigation](alpha-hardening-navigation.md) — locked fraud/auth decisions, Bucket-1 order, handoff pointers - **Node capability admission** — `.scratch/node-capability-admission/` (P0 plan: generic doctor/real-forward validation, fail-closed readiness, tracker admission gate; [PRD](../.scratch/node-capability-admission/PRD.md), [README](../.scratch/node-capability-admission/README.md), ADR-0023) +- **Distributed relay performance** — relay `/rpc` requester sockets are persistent per Route Session and Activation Seam as of 2026-07-10; `request_id` remains unique per activation while `X-Meshnet-Session` remains stable for KV state. Next low-risk priorities: persistent direct/loopback HTTP, seam byte/latency telemetry, then trace-driven zstd tuning. diff --git a/.scratch/distributed-inference-performance/PRD.md b/.scratch/distributed-inference-performance/PRD.md new file mode 100644 index 0000000..4e025d9 --- /dev/null +++ b/.scratch/distributed-inference-performance/PRD.md @@ -0,0 +1,40 @@ +# PRD: Distributed inference performance + +## Problem + +Distributed decode already avoids full-prompt recomputation when the local KV +path is active, but each Activation Seam can still pay transport and data-plane +overhead for every generated token. Relay logs show a new `request_id` per +token; that is correct correlation, but the old relay implementation also +opened a new WebSocket per token. Direct hops and relay bridge forwarding use +fresh HTTP requests as well. Without timing and byte measurements, compression, +copy, and buffering choices cannot be ranked safely. + +## Outcome + +For a cached Route Session, connection setup is amortized across the session, +decode payloads remain one-step activations, progress reporting is bounded, and +the benchmark can attribute latency to model execution, serialization, relay, +HTTP, queueing, and backpressure. Optimizations must preserve output tokens, +KV semantics, failure behavior, and compatibility with legacy one-shot peers. + +## Non-goals + +- No speculative decoding or multi-token model execution in this feature. +- No QUIC/WebRTC/custom transport rewrite. +- No centralized Hot KV State. +- No silent reuse of a `request_id`; each activation remains independently + traceable. + +## Acceptance criteria + +- A reproducible local two-node and relay benchmark reports per-token and + per-seam timing plus bytes. +- Cached decode does not perform a new TCP/WebSocket connection per token. +- Direct and relay-to-local HTTP paths reuse connections safely or document why + a path cannot do so. +- Compression and copy decisions are based on recorded traces, not guesses. +- Slow prefill consumers apply bounded backpressure rather than unbounded body + buffering. +- A benchmark regression threshold catches a meaningful transport slowdown. + diff --git a/.scratch/distributed-inference-performance/README.md b/.scratch/distributed-inference-performance/README.md new file mode 100644 index 0000000..e8ec6ec --- /dev/null +++ b/.scratch/distributed-inference-performance/README.md @@ -0,0 +1,37 @@ +# Distributed inference performance + +Status: draft scratch package. + +This feature measures and reduces avoidable overhead around the existing +Route Session, prefill/decode, local Hot KV State, and binary activation path. +It does not replace the distributed GGUF runtime plan. The goal is to make +transport and data movement cheap enough that model execution, rather than +connection setup, logging, or serialization, dominates token latency. + +## Scope + +- Baseline per-token compute, seam, connection, serialization, and queue time. +- Keep one connection alive for a Route Session wherever protocol semantics allow. +- Add bounded, actionable Generation Telemetry for each Activation Seam. +- Tune compression and buffer conversion from measured activation traces. +- Add bounded prefill backpressure and an end-to-end benchmark gate. + +## Existing decisions preserved + +- `X-Meshnet-Session` is stable for one Route Session. +- `request_id` remains unique per activation request for correlation. +- Hot KV State remains local to each Shard node. +- v1 activation transfer remains binary HTTP-shaped traffic. +- Streaming output remains preferred and telemetry remains mandatory. + +## Task order + +1. 01 — baseline and profiling harness +2. 02 — persistent relay compatibility hardening +3. 03 — direct and bridge HTTP keep-alive +4. 04 — seam telemetry and bounded progress reporting +5. 05 — adaptive activation compression +6. 06 — activation framing and copy reduction +7. 07 — prefill chunk backpressure +8. 08 — end-to-end performance gate + diff --git a/.scratch/distributed-inference-performance/issues/01-baseline-profiling-harness.md b/.scratch/distributed-inference-performance/issues/01-baseline-profiling-harness.md new file mode 100644 index 0000000..0847134 --- /dev/null +++ b/.scratch/distributed-inference-performance/issues/01-baseline-profiling-harness.md @@ -0,0 +1,27 @@ +Status: ready-for-agent + +# 01 — Baseline and profiling harness + +## What to build + +Create a deterministic stub-backed benchmark for a Route Session that measures +prefill and cached decode across direct and relay paths. Attribute time to model +execution, activation encoding/decoding, compression, connection setup, relay +queueing, local HTTP forwarding, and end-to-end seam latency. Record payload +sizes and connection counts without requiring a real model or external host. + +## Acceptance criteria + +- [ ] The harness runs a fixed prompt and fixed generated-token count through a + two-node route in direct and relay modes. +- [ ] It reports p50/p95 per-token latency, per-hop latency, payload bytes, + compression ratio, connection attempts, and queue wait. +- [ ] It distinguishes prefill from decode and cached from stateless mode. +- [ ] It emits machine-readable JSON suitable for CI artifacts and a concise + human-readable summary. +- [ ] A test fixture can assert connection attempts and output token identity. + +## Blocked by + +None - can start immediately. + diff --git a/.scratch/distributed-inference-performance/issues/02-relay-session-compatibility.md b/.scratch/distributed-inference-performance/issues/02-relay-session-compatibility.md new file mode 100644 index 0000000..38cc7d6 --- /dev/null +++ b/.scratch/distributed-inference-performance/issues/02-relay-session-compatibility.md @@ -0,0 +1,29 @@ +Status: ready-for-agent + +# 02 — Persistent relay compatibility hardening + +## What to build + +Harden the persistent `/rpc/` connection used by one Route Session. +Preserve unique request correlation while allowing sequential binary and JSON +requests on one socket. Handle peer disconnects, requester cancellation, +legacy one-request relays, timeout cleanup, and generation-end close without +leaking pending RPC entries or accidentally replaying a model mutation. + +## Acceptance criteria + +- [ ] A cached decode session uses one requester connection per relay Activation + Seam and sends one unique request id per activation. +- [ ] Legacy relays that close after one response fail over clearly without + corrupting the Route Session or replaying an uncertain request. +- [ ] Relay and bridge cleanup removes pending request state on normal close, + cancellation, timeout, and peer disconnect. +- [ ] Concurrent Route Sessions do not share a non-thread-safe socket; responses + remain matched by request id. +- [ ] Tests cover two sequential binary requests, JSON compatibility, close, + timeout, disconnect, cancellation, and no leaked pending entries. + +## Blocked by + +- 01 — Baseline and profiling harness. + diff --git a/.scratch/distributed-inference-performance/issues/03-http-keepalive.md b/.scratch/distributed-inference-performance/issues/03-http-keepalive.md new file mode 100644 index 0000000..590ccd7 --- /dev/null +++ b/.scratch/distributed-inference-performance/issues/03-http-keepalive.md @@ -0,0 +1,28 @@ +Status: ready-for-agent + +# 03 — Direct and bridge HTTP keep-alive + +## What to build + +Amortize TCP connection setup for direct node hops and for the relay bridge's +local request into the shard server. Use bounded per-session or per-worker +connection ownership, explicit response lengths, and safe invalidation on +errors. Do not share a connection across concurrent requests unless the client +supports serialization. + +## Acceptance criteria + +- [ ] Direct cached decode reuses a connection to each downstream HTTP node. +- [ ] Relay bridge forwarding reuses loopback HTTP connections without blocking + unrelated worker requests. +- [ ] HTTP/1.1 framing is correct for success, error, empty, streamed, and + cancellation responses; no request hangs waiting for EOF. +- [ ] Broken or stale connections are discarded and the current request follows + the existing safe failure/fallback policy. +- [ ] Benchmark 01 shows connection attempts are independent of generated token + count for a healthy session. + +## Blocked by + +- 01 — Baseline and profiling harness. + diff --git a/.scratch/distributed-inference-performance/issues/04-seam-telemetry.md b/.scratch/distributed-inference-performance/issues/04-seam-telemetry.md new file mode 100644 index 0000000..c2fe3ed --- /dev/null +++ b/.scratch/distributed-inference-performance/issues/04-seam-telemetry.md @@ -0,0 +1,28 @@ +Status: ready-for-agent + +# 04 — Activation Seam telemetry and bounded progress reporting + +## What to build + +Expose structured timing and byte counters for each Activation Seam while +keeping per-token progress overhead bounded. Report route/session, phase, +hop/node, queue wait, model time, encode/decode time, compression time, wire +bytes, response bytes, and connection reuse. Aggregate high-cardinality events +instead of flushing a log line for every token. + +## Acceptance criteria + +- [ ] Generation Telemetry includes prefill/decode seam latency and rolling + tokens/sec without changing token output or cache behavior. +- [ ] Every request can be correlated by stable Route Session plus unique + activation request id. +- [ ] Counters are sampled or aggregated so telemetry work is bounded and does + not perform network I/O in the model hot loop. +- [ ] Logs summarize decode progress by session and retain actionable failure + context. +- [ ] Tests verify counters, aggregation cadence, and cleanup at session close. + +## Blocked by + +- 01 — Baseline and profiling harness. + diff --git a/.scratch/distributed-inference-performance/issues/05-adaptive-compression.md b/.scratch/distributed-inference-performance/issues/05-adaptive-compression.md new file mode 100644 index 0000000..46cc2de --- /dev/null +++ b/.scratch/distributed-inference-performance/issues/05-adaptive-compression.md @@ -0,0 +1,26 @@ +Status: ready-for-agent + +# 05 — Trace-driven activation compression + +## What to build + +Make zstd decisions from activation size, measured compression ratio, CPU cost, +and route conditions. Keep small decode payloads on the fast path, avoid +compressing data that does not shrink, and expose enough counters to compare +wire savings against compression latency on LAN and relay routes. + +## Acceptance criteria + +- [ ] A compression policy is explicit and configurable for LAN, relay, and + benchmark environments. +- [ ] Bodies that do not meet the configured savings threshold are sent raw. +- [ ] Compression and decompression time plus input/output bytes are reported. +- [ ] Prefill and decode policies can differ; decode latency is not regressed by + compressing small one-step activations. +- [ ] Tests cover incompressible, compressible, threshold, malformed, and + legacy-uncompressed bodies. + +## Blocked by + +- 01 — Baseline and profiling harness. + diff --git a/.scratch/distributed-inference-performance/issues/06-activation-framing-copies.md b/.scratch/distributed-inference-performance/issues/06-activation-framing-copies.md new file mode 100644 index 0000000..46e2258 --- /dev/null +++ b/.scratch/distributed-inference-performance/issues/06-activation-framing-copies.md @@ -0,0 +1,27 @@ +Status: ready-for-agent + +# 06 — Activation framing and copy reduction + +## What to build + +Profile and reduce avoidable allocations while activation data crosses a seam: +binary frame assembly, header JSON, base64 metadata, CPU/GPU conversion, and +response decompression. Preserve the current binary wire contract and use +zero-copy or pooled buffers only where ownership and lifetime are explicit. + +## Acceptance criteria + +- [ ] The benchmark identifies copy/allocation cost separately from model and + network time. +- [ ] Decode hidden-state conversion has no unnecessary float32 round trip. +- [ ] Binary framing avoids base64 for activation bodies and does not retain + buffers after a request completes. +- [ ] Position/attention metadata is validated and encoded efficiently without + changing semantic headers or cache positions. +- [ ] A focused test proves byte-for-byte wire compatibility and stable output + tokens before and after the optimization. + +## Blocked by + +- 01 — Baseline and profiling harness. + diff --git a/.scratch/distributed-inference-performance/issues/07-prefill-backpressure.md b/.scratch/distributed-inference-performance/issues/07-prefill-backpressure.md new file mode 100644 index 0000000..a955547 --- /dev/null +++ b/.scratch/distributed-inference-performance/issues/07-prefill-backpressure.md @@ -0,0 +1,27 @@ +Status: ready-for-agent + +# 07 — Bounded prefill chunk backpressure + +## What to build + +Make large prefill transfer bounded across every Activation Seam. Chunk prompt +activations, limit in-flight chunks, and propagate downstream congestion so a +slow Shard node cannot cause the head or relay bridge to buffer the entire +context in memory. Keep decode-step traffic sequential and low-latency. + +## Acceptance criteria + +- [ ] Configurable prefill chunk size and in-flight limit exist with safe + defaults. +- [ ] Peak per-hop buffered bytes are bounded by the configured limits. +- [ ] A slow downstream stub applies backpressure and does not lose or reorder + chunks within a Route Session. +- [ ] Cancellation and route failure release queued chunks and local buffers. +- [ ] Tests cover small prompts, multi-chunk prompts, slow consumers, retry/fail + closeout, and legacy single-chunk peers. + +## Blocked by + +- 01 — Baseline and profiling harness. +- 04 — Activation Seam telemetry and bounded progress reporting. + diff --git a/.scratch/distributed-inference-performance/issues/08-end-to-end-performance-gate.md b/.scratch/distributed-inference-performance/issues/08-end-to-end-performance-gate.md new file mode 100644 index 0000000..e407ef4 --- /dev/null +++ b/.scratch/distributed-inference-performance/issues/08-end-to-end-performance-gate.md @@ -0,0 +1,33 @@ +Status: ready-for-agent + +# 08 — End-to-end distributed performance gate + +## What to build + +Turn the benchmark and optimizations into a repeatable performance gate for a +small two-node route and a relay route. Compare stateless legacy mode, cached +decode, direct HTTP, and persistent relay. Fail only on stable regressions and +publish the measurements needed to decide whether further work belongs in +transport, serialization, queueing, or model execution. + +## Acceptance criteria + +- [ ] CI/local benchmark runs a deterministic fixed-token scenario without a + real model or external network. +- [ ] The report compares tokens/sec, p50/p95 token latency, seam latency, + bytes/token, connection count, compression CPU, and peak buffered bytes. +- [ ] Thresholds are documented and tolerant of normal host variance while + catching a meaningful regression. +- [ ] A real-model opt-in command records the same metrics for LAN validation. +- [ ] The gate verifies output token identity, Route Session stability, and + cleanup of sessions, sockets, queues, and telemetry state. + +## Blocked by + +- 02 — Persistent relay compatibility hardening. +- 03 — Direct and bridge HTTP keep-alive. +- 04 — Activation Seam telemetry and bounded progress reporting. +- 05 — Trace-driven activation compression. +- 06 — Activation framing and copy reduction. +- 07 — Bounded prefill chunk backpressure. + diff --git a/.scratch/distributed-inference-performance/prd.json b/.scratch/distributed-inference-performance/prd.json new file mode 100644 index 0000000..a9000e7 --- /dev/null +++ b/.scratch/distributed-inference-performance/prd.json @@ -0,0 +1,135 @@ +{ + "name": "Distributed inference performance", + "description": "Measure and reduce avoidable transport, HTTP, telemetry, compression, buffering, and copy overhead around Route Session cached decode.", + "branchName": "ralph/distributed-inference-performance", + "userStories": [ + { + "id": "DIP-001", + "title": "01 — Baseline and profiling harness", + "description": "Create a deterministic stub-backed direct and relay Route Session benchmark that reports per-token and per-seam timing, bytes, compression, queueing, and connection counts.", + "acceptanceCriteria": [ + "Fixed prompt/token scenario runs in direct and relay modes", + "Reports p50/p95 latency, payload bytes, compression ratio, connections, and queue wait", + "Distinguishes prefill/decode and cached/stateless modes", + "Produces machine-readable JSON and human-readable summary", + "Can assert connection count and output token identity" + ], + "priority": 1, + "passes": false, + "notes": "Source issue: .scratch/distributed-inference-performance/issues/01-baseline-profiling-harness.md", + "dependsOn": [] + }, + { + "id": "DIP-002", + "title": "02 — Persistent relay compatibility hardening", + "description": "Harden one persistent relay RPC connection per Route Session seam while preserving unique request IDs, legacy compatibility, cleanup, cancellation, and safe failure behavior.", + "acceptanceCriteria": [ + "One healthy relay connection serves sequential cached decode requests", + "Legacy one-request relays fail over without replaying uncertain mutations", + "Pending RPC state is cleaned on close, cancellation, timeout, and disconnect", + "Concurrent sessions do not share an unsafe socket", + "Tests cover binary, JSON, timeout, disconnect, cancellation, and cleanup" + ], + "priority": 2, + "passes": false, + "notes": "Source issue: .scratch/distributed-inference-performance/issues/02-relay-session-compatibility.md", + "dependsOn": ["DIP-001"] + }, + { + "id": "DIP-003", + "title": "03 — Direct and bridge HTTP keep-alive", + "description": "Amortize TCP setup for direct node hops and relay bridge loopback forwarding with bounded connection ownership and correct HTTP framing.", + "acceptanceCriteria": [ + "Direct cached decode reuses downstream HTTP connections", + "Bridge loopback forwarding reuses connections without blocking unrelated workers", + "HTTP/1.1 framing works for success, error, empty, stream, and cancellation", + "Stale connections are invalidated under the existing fallback policy", + "Benchmark shows healthy-session connection count independent of token count" + ], + "priority": 3, + "passes": false, + "notes": "Source issue: .scratch/distributed-inference-performance/issues/03-http-keepalive.md", + "dependsOn": ["DIP-001"] + }, + { + "id": "DIP-004", + "title": "04 — Activation Seam telemetry and bounded progress reporting", + "description": "Expose per-seam timing and byte counters while aggregating progress work so telemetry does not become per-token hot-loop overhead.", + "acceptanceCriteria": [ + "Telemetry includes seam latency and rolling tokens/sec", + "Stable session and unique activation IDs remain correlated", + "Counters are sampled or aggregated without network I/O in model execution", + "Decode logs summarize by session with actionable failures", + "Tests verify cadence and cleanup" + ], + "priority": 4, + "passes": false, + "notes": "Source issue: .scratch/distributed-inference-performance/issues/04-seam-telemetry.md", + "dependsOn": ["DIP-001"] + }, + { + "id": "DIP-005", + "title": "05 — Trace-driven activation compression", + "description": "Choose zstd based on measured savings and CPU cost, preserving a raw fast path for small or incompressible decode activations.", + "acceptanceCriteria": [ + "Compression policy is explicit and configurable per route condition", + "Bodies below savings threshold are sent raw", + "Compression timing and byte counters are reported", + "Prefill and decode can use different policies", + "Tests cover compressible, incompressible, threshold, malformed, and legacy bodies" + ], + "priority": 5, + "passes": false, + "notes": "Source issue: .scratch/distributed-inference-performance/issues/05-adaptive-compression.md", + "dependsOn": ["DIP-001"] + }, + { + "id": "DIP-006", + "title": "06 — Activation framing and copy reduction", + "description": "Measure and reduce avoidable binary framing, metadata, CPU/GPU conversion, and decompression allocations without changing the wire contract.", + "acceptanceCriteria": [ + "Benchmark attributes copy/allocation cost separately", + "Decode hidden state avoids unnecessary float32 conversion", + "Activation bodies remain binary and buffers have explicit ownership", + "Metadata encoding remains semantically compatible", + "Wire and token-output regression tests pass" + ], + "priority": 6, + "passes": false, + "notes": "Source issue: .scratch/distributed-inference-performance/issues/06-activation-framing-copies.md", + "dependsOn": ["DIP-001"] + }, + { + "id": "DIP-007", + "title": "07 — Bounded prefill chunk backpressure", + "description": "Bound large prefill transfer with configurable chunking and in-flight limits, propagating downstream congestion without reordering or unbounded buffering.", + "acceptanceCriteria": [ + "Chunk size and in-flight limit have safe defaults", + "Peak buffered bytes stay within configured bounds", + "Slow consumers apply backpressure and preserve order", + "Cancellation and route failure release queued buffers", + "Tests cover chunking, slow consumers, failure, and legacy peers" + ], + "priority": 7, + "passes": false, + "notes": "Source issue: .scratch/distributed-inference-performance/issues/07-prefill-backpressure.md", + "dependsOn": ["DIP-001", "DIP-004"] + }, + { + "id": "DIP-008", + "title": "08 — End-to-end distributed performance gate", + "description": "Make direct, persistent-relay, cached, and stateless benchmark comparisons repeatable and fail on meaningful transport regressions while verifying output and cleanup.", + "acceptanceCriteria": [ + "Deterministic stub benchmark runs locally and in CI", + "Report compares throughput, latency, bytes, connections, compression CPU, and buffers", + "Regression thresholds tolerate host variance and catch meaningful slowdowns", + "Opt-in real-model LAN command emits the same metrics", + "Gate verifies token identity, session stability, and resource cleanup" + ], + "priority": 8, + "passes": false, + "notes": "Source issue: .scratch/distributed-inference-performance/issues/08-end-to-end-performance-gate.md", + "dependsOn": ["DIP-002", "DIP-003", "DIP-004", "DIP-005", "DIP-006", "DIP-007"] + } + ] +} diff --git a/docs/adr/0014-relay-outbound-client.md b/docs/adr/0014-relay-outbound-client.md index 3930c75..6d2c28c 100644 --- a/docs/adr/0014-relay-outbound-client.md +++ b/docs/adr/0014-relay-outbound-client.md @@ -1,6 +1,6 @@ # ADR-0014: Relay outbound client for NAT/internet pipeline hops -## Status: Accepted +## Status: Accepted, amended 2026-07-10 ## Context @@ -25,20 +25,22 @@ of connection setup matters. ## Options considered -**A. Relay hop (WebSocket per hop, chosen)** -Node A opens a WebSocket to `wss://relay/rpc/{peer_id_B}`, sends the activation, -receives the response, closes. The relay's `_handle_rpc` forwards it to B's persistent -connection via the existing `relay-http-request` envelope mechanism. +**A. Relay hop (persistent per Route Session, chosen)** +Node A opens a WebSocket to `wss://relay/rpc/{peer_id_B}`, sends activation requests +sequentially for the Route Session, then closes it when generation ends. The relay's +`_handle_rpc` forwards each request to B's persistent connection via the existing +`relay-http-request` envelope mechanism. Pros: reuses the existing relay server unchanged. Each hop is independent; failures don't affect other requests. -Cons: WebSocket connection setup adds ~50–150 ms per hop on a fast relay. For -autoregressive inference (N tokens × M hops), this adds up. +The original implementation opened and closed this socket per token. It was amended +to retain one requester socket per downstream relay address for the generation, so +connection setup is amortized across all tokens. -**B. Persistent per-session tunnel** -Node A opens a persistent WebSocket to the relay for the duration of an inference session -and multiplexes all token hops over it. +**B. Multiplexed persistent tunnel** +Node A sends multiple concurrent Route Sessions over a shared WebSocket and demultiplexes +responses by request id. Pros: amortises connection setup across tokens. @@ -53,15 +55,16 @@ traffic through the tracker would saturate it. Rejected. ## Decision -Option A — per-hop WebSocket relay. Simple, reuses existing infrastructure, correct. -Option B is noted as a future optimization when activation-path latency becomes the -bottleneck. +Option A — one sequential WebSocket per relayed Activation Seam and Route Session. +Each activation still has a unique request id for response correlation, while +`X-Meshnet-Session` remains stable for KV state. Option B remains a possible +connection-count optimization for high-concurrency workloads. ## Protocol ``` -Node A opens WS → wss://relay/rpc/{peer_id_B} -Node A sends: +Node A opens WS once → wss://relay/rpc/{peer_id_B} +Node A sends repeatedly: { "request_id": "", "method": "POST", @@ -81,8 +84,8 @@ Response: # OR "body": "" ← for text (last-hop decode) } -Relay sends response JSON back to Node A. -Node A decodes body_base64, continues pipeline. +Relay sends each response back to Node A without closing the requester socket. +Node A continues the pipeline and closes the socket when generation ends. ``` ### Binary data through JSON: base64 @@ -115,6 +118,6 @@ The head node reads `relay_addr` from the injected `X-Meshnet-Route` header and - Nodes behind NAT (WSL2, 5G, home routers) can now participate in distributed pipeline inference without opening firewall ports - `relay_addr` is a stable registration field; nodes without a relay omit it and receive direct HTTP hops -- Per-hop WebSocket setup adds latency proportional to relay RTT; acceptable for prototype, optimize later with persistent tunnels +- WebSocket/TCP/TLS setup occurs once per relayed Activation Seam per Route Session, not once per generated token - Base64 encoding increases payload size by ~33%; acceptable for prototype - The relay server remains stateless and horizontally scalable; only the persistent per-peer `/ws` connections are stateful diff --git a/packages/node/meshnet_node/relay_bridge.py b/packages/node/meshnet_node/relay_bridge.py index 498c08a..ef90217 100644 --- a/packages/node/meshnet_node/relay_bridge.py +++ b/packages/node/meshnet_node/relay_bridge.py @@ -112,6 +112,8 @@ class RelayHttpBridge: self._connected = threading.Event() self._executor: ThreadPoolExecutor | None = None self._send_lock = threading.Lock() + self._decode_log_lock = threading.Lock() + self._decode_steps: dict[str, int] = {} self._ws = None @property @@ -230,8 +232,21 @@ class RelayHttpBridge: headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {} binary_mode = binary_body is not None + session = str(headers.get("X-Meshnet-Session") or "") + cache_mode = headers.get("X-Meshnet-Cache") req_suffix = f" request_id={request_id}" if request_id else "" - print(f" [node] relay {method} {path}{req_suffix}", flush=True) + if path == "/forward" and cache_mode == "decode" and session: + with self._decode_log_lock: + steps = self._decode_steps.get(session, 0) + 1 + self._decode_steps[session] = steps + if steps == 1 or steps % 32 == 0: + print( + f" [node] relay {method} {path} session={session[:8]} steps={steps}{req_suffix}", + flush=True, + ) + else: + session_suffix = f" session={session[:8]}" if session else "" + print(f" [node] relay {method} {path}{session_suffix}{req_suffix}", flush=True) if binary_mode: data = binary_body diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index b6cac85..45c5bea 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -100,6 +100,62 @@ def _write_progress_line(state: list[bool], message: str, *, final: bool = False sys.stdout.flush() +class _RelayHopClient: + """Persistent relay connection scoped to one generation handler.""" + + def __init__(self, relay_addr: str, timeout: float = 120.0) -> None: + self.relay_addr = relay_addr + self.timeout = timeout + self._ws = None + + def request( + self, + path: str, + body: bytes, + headers: dict[str, str], + ) -> tuple[int, dict[str, str], bytes]: + import websockets.sync.client as wsc # type: ignore[import] + + from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size + + if self._ws is None: + self._ws = wsc.connect( + self.relay_addr, + open_timeout=self.timeout, + max_size=ws_max_size(), + compression=None, + ) + request_id = f"{time.time_ns():x}" + frame = encode_binary_frame({ + "request_id": request_id, + "method": "POST", + "path": path, + "headers": headers, + }, body) + self._ws.send(frame) + raw = self._ws.recv(timeout=self.timeout) + if isinstance(raw, (bytes, bytearray)): + resp_header, resp_body = decode_binary_frame(bytes(raw)) + status = int(resp_header.get("status", 503)) + resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()} + return status, resp_headers, resp_body + resp = json.loads(raw) + status = int(resp.get("status", 503)) + resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()} + body_b64 = resp.get("body_base64") + resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode() + return status, resp_headers, resp_body + + def close(self) -> None: + if self._ws is not None: + try: + self._ws.close() + except Exception: + pass + finally: + self._ws = None + + def _relay_hop( relay_addr: str, path: str, @@ -107,7 +163,7 @@ def _relay_hop( headers: dict[str, str], timeout: float = 120.0, ) -> tuple[int, dict[str, str], bytes]: - """Send a single HTTP-shaped request through a relay RPC WebSocket. + """Send one request through a short-lived relay client (compatibility API). relay_addr is the wss://relay.../rpc/{peer_id} URL. The request and any binary response travel as binary frames (JSON header + raw body); relay @@ -115,33 +171,11 @@ def _relay_hop( Returns (status, response_headers_lower, response_body). Raises on connection failure so callers can fall back to direct. """ - import websockets.sync.client as wsc # type: ignore[import] - - from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size - - request_id = f"{time.time_ns():x}" - frame = encode_binary_frame({ - "request_id": request_id, - "method": "POST", - "path": path, - "headers": headers, - }, body) - with wsc.connect( - relay_addr, open_timeout=timeout, max_size=ws_max_size(), compression=None, - ) as ws: - ws.send(frame) - raw = ws.recv(timeout=timeout) - if isinstance(raw, (bytes, bytearray)): - resp_header, resp_body = decode_binary_frame(bytes(raw)) - status = int(resp_header.get("status", 503)) - resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()} - return status, resp_headers, resp_body - resp = json.loads(raw) - status = int(resp.get("status", 503)) - resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()} - body_b64 = resp.get("body_base64") - resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode() - return status, resp_headers, resp_body + client = _RelayHopClient(relay_addr, timeout) + try: + return client.request(path, body, headers) + finally: + client.close() # Below this, zstd overhead outweighs the win (per-token decode bodies are ~KBs). @@ -687,6 +721,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): progress_line = [False] last_token_id: int | None = None failure_reason: str | None = None + relay_clients: dict[str, _RelayHopClient] = {} def _prefill_step() -> tuple[str, int | None]: """Full-sequence prefill: initial step and cache-miss recovery.""" @@ -698,6 +733,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): return self._run_downstream_pipeline( payload, remaining_route, backend=backend, session=session_id, cache_mode="prefill" if use_kv else None, + relay_clients=relay_clients, ) for step in range(max_tokens): @@ -708,6 +744,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): token_str, token_id = self._run_downstream_pipeline( payload, remaining_route, backend=backend, session=session_id, cache_mode="decode", + relay_clients=relay_clients, ) except (KVCacheMiss, _PipelineCacheMiss) as miss: # Evicted/restarted node or head lost its own session: @@ -769,6 +806,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): backend.release_session(session_id) except Exception: pass + for relay_client in relay_clients.values(): + relay_client.close() if generated: elapsed = time.monotonic() - gen_started @@ -881,6 +920,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): backend: TorchModelShard | None = None, session: str | None = None, cache_mode: str | None = None, + relay_clients: dict[str, _RelayHopClient] | None = None, ) -> tuple[str, int | None]: """Forward an activation through the downstream route. @@ -953,9 +993,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): headers["X-Meshnet-Position-Ids"] = current_pos if relay_addr: try: - status, resp_headers, resp_body = _relay_hop( - relay_addr, "/forward", wire_body, headers, timeout=120.0, - ) + if relay_clients is None: + status, resp_headers, resp_body = _relay_hop( + relay_addr, "/forward", wire_body, headers, timeout=120.0, + ) + else: + relay_client = relay_clients.setdefault( + relay_addr, _RelayHopClient(relay_addr, timeout=120.0), + ) + status, resp_headers, resp_body = relay_client.request( + "/forward", wire_body, headers, + ) if status == 409 and _is_cache_miss_body(resp_body): raise _PipelineCacheMiss(node_url) if status >= 400: diff --git a/packages/relay/meshnet_relay/server.py b/packages/relay/meshnet_relay/server.py index 270c495..89355e2 100644 --- a/packages/relay/meshnet_relay/server.py +++ b/packages/relay/meshnet_relay/server.py @@ -273,7 +273,7 @@ class RelayServer: ) async def _handle_rpc(self, ws_requester, target_peer_id: str) -> None: - """Send one HTTP-shaped request to a connected peer and relay its response.""" + """Relay sequential HTTP-shaped requests over one requester connection.""" target = self._registry.get(target_peer_id) if target is None: await ws_requester.send(json.dumps({ @@ -284,64 +284,76 @@ class RelayServer: await ws_requester.close() return - try: - raw = await asyncio.wait_for(ws_requester.recv(), timeout=30.0) - if isinstance(raw, (bytes, bytearray)): - header, body = decode_binary_frame(bytes(raw)) - request_id = str(header.get("request_id") or uuid.uuid4()) - header["request_id"] = request_id - header["target_peer"] = target_peer_id - outbound: str | bytes = encode_binary_frame(header, body) - else: - payload = json.loads(raw) - request_id = str(payload.get("request_id") or uuid.uuid4()) - payload["request_id"] = request_id - payload["target_peer"] = target_peer_id - outbound = json.dumps({ - "topic": "relay-http-request", - "version": 1, - "from_peer": "relay", - "payload": payload, - }) - except Exception: - await ws_requester.close(1003, "invalid relay rpc request") - return + while True: + try: + raw = await ws_requester.recv() + except Exception: + return - queue: asyncio.Queue = asyncio.Queue() - self._pending_rpc[request_id] = queue - overall_timeout = 310.0 - idle_timeout = 120.0 - loop = asyncio.get_running_loop() - deadline = loop.time() + overall_timeout - try: - await target.ws.send(outbound) - # Forward frames until a terminal one: streamed responses (US-036) - # end with {"stream": true, "done": true}; a frame without "stream" - # is a complete legacy single response. A binary frame is always a - # complete single response. - while True: - remaining = deadline - loop.time() - if remaining <= 0: - raise asyncio.TimeoutError - frame = await asyncio.wait_for( - queue.get(), timeout=min(idle_timeout, remaining) - ) - if isinstance(frame, (bytes, bytearray)): - await ws_requester.send(frame) - break - await ws_requester.send(json.dumps(frame)) - if not frame.get("stream") or frame.get("done"): - break - except asyncio.TimeoutError: - await ws_requester.send(json.dumps({ - "request_id": request_id, - "status": 504, - "headers": {"Content-Type": "application/json"}, - "body": json.dumps({"error": "relay rpc timed out"}), - })) - finally: - self._pending_rpc.pop(request_id, None) - await ws_requester.close() + request_id: str | None = None + try: + if isinstance(raw, (bytes, bytearray)): + header, body = decode_binary_frame(bytes(raw)) + request_id = str(header.get("request_id") or uuid.uuid4()) + header["request_id"] = request_id + header["target_peer"] = target_peer_id + outbound: str | bytes = encode_binary_frame(header, body) + else: + payload = json.loads(raw) + request_id = str(payload.get("request_id") or uuid.uuid4()) + payload["request_id"] = request_id + payload["target_peer"] = target_peer_id + outbound = json.dumps({ + "topic": "relay-http-request", + "version": 1, + "from_peer": "relay", + "payload": payload, + }) + except Exception: + await ws_requester.close(1003, "invalid relay rpc request") + return + + queue: asyncio.Queue = asyncio.Queue() + self._pending_rpc[request_id] = queue + overall_timeout = 310.0 + idle_timeout = 120.0 + loop = asyncio.get_running_loop() + deadline = loop.time() + overall_timeout + target = self._registry.get(target_peer_id) + try: + if target is None: + await ws_requester.send(json.dumps({ + "request_id": request_id, + "status": 503, + "headers": {"Content-Type": "application/json"}, + "body": json.dumps({"error": f"peer {target_peer_id!r} disconnected"}), + })) + continue + await target.ws.send(outbound) + # Streamed responses end with done=true. Binary and legacy JSON + # responses are complete in one frame. + while True: + remaining = deadline - loop.time() + if remaining <= 0: + raise asyncio.TimeoutError + frame = await asyncio.wait_for( + queue.get(), timeout=min(idle_timeout, remaining) + ) + if isinstance(frame, (bytes, bytearray)): + await ws_requester.send(frame) + break + await ws_requester.send(json.dumps(frame)) + if not frame.get("stream") or frame.get("done"): + break + except asyncio.TimeoutError: + await ws_requester.send(json.dumps({ + "request_id": request_id, + "status": 504, + "headers": {"Content-Type": "application/json"}, + "body": json.dumps({"error": "relay rpc timed out"}), + })) + finally: + self._pending_rpc.pop(request_id, None) async def _broadcast(raw: str | bytes, peers: list) -> None: diff --git a/tests/test_gossip_and_relay.py b/tests/test_gossip_and_relay.py index ec01717..e20f6f3 100644 --- a/tests/test_gossip_and_relay.py +++ b/tests/test_gossip_and_relay.py @@ -372,6 +372,70 @@ def test_relay_rpc_round_trips_http_request_to_peer(): assert json.loads(response["body"]) == {"ok": True, "path": "/v1/chat/completions"} +def test_relay_rpc_reuses_connection_for_sequential_requests(monkeypatch): + """One route session should not repeat the WebSocket handshake per token.""" + import websockets.sync.client as wsc # type: ignore[import] + from meshnet_node.relay_bridge import decode_binary_frame, encode_binary_frame + from meshnet_node.torch_server import _RelayHopClient + from meshnet_relay.server import RelayServer + + relay = RelayServer(host="127.0.0.1", port=0) + port = relay.start() + ready = threading.Event() + + def run_peer(): + with wsc.connect(f"ws://127.0.0.1:{port}/ws") as ws: + ws.send(json.dumps({ + "topic": "peer-register", + "version": 1, + "from_peer": "persistent_peer", + "msg_id": "persistent-reg-001", + "payload": {"peer_id": "persistent_peer", "addr": ""}, + })) + ws.recv() + ready.set() + for _ in range(2): + raw = ws.recv(timeout=3) + header, body = decode_binary_frame(raw) + ws.send(encode_binary_frame({ + "request_id": header["request_id"], + "status": 200, + "headers": {"Content-Type": "application/octet-stream"}, + }, body)) + + peer_thread = threading.Thread(target=run_peer, daemon=True) + peer_thread.start() + assert ready.wait(timeout=5) + + real_connect = wsc.connect + connection_attempts = 0 + + def tracked_connect(*args, **kwargs): + nonlocal connection_attempts + connection_attempts += 1 + return real_connect(*args, **kwargs) + + monkeypatch.setattr(wsc, "connect", tracked_connect) + client = _RelayHopClient(f"ws://127.0.0.1:{port}/rpc/persistent_peer", timeout=5) + try: + responses = [ + client.request( + "/forward", + f"token-{index}".encode(), + {"X-Meshnet-Session": "route-session"}, + ) + for index in range(2) + ] + finally: + client.close() + relay.stop() + peer_thread.join(timeout=3) + + assert [status for status, _, _ in responses] == [200, 200] + assert [body for _, _, body in responses] == [b"token-0", b"token-1"] + assert connection_attempts == 1 + + def test_binary_relay_frame_codecs_interoperate(): """Node and relay ship the same binary frame format as separate copies.""" import os