From e30272e83f4745e83774c539f06778d732c7ae31 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Thu, 9 Jul 2026 22:40:43 +0200 Subject: [PATCH] dropp baes64 use binary --- packages/node/meshnet_node/relay_bridge.py | 78 +++++++++++++++++++--- packages/node/meshnet_node/torch_server.py | 51 +++++++++++--- packages/relay/meshnet_relay/server.py | 68 ++++++++++++++++--- tests/test_gossip_and_relay.py | 41 ++++++++++++ tests/test_kv_cache_distributed.py | 37 ++++++++++ 5 files changed, 244 insertions(+), 31 deletions(-) diff --git a/packages/node/meshnet_node/relay_bridge.py b/packages/node/meshnet_node/relay_bridge.py index 02e01ad..498c08a 100644 --- a/packages/node/meshnet_node/relay_bridge.py +++ b/packages/node/meshnet_node/relay_bridge.py @@ -19,7 +19,7 @@ log = logging.getLogger(__name__) DEFAULT_MAX_CONCURRENCY = 8 -# Activation tensors ride the relay as base64 inside one JSON frame, so the +# Activation tensors ride the relay as one WebSocket frame per hop, so the # websockets default of 1 MiB rejects any real prefill (close code 1009). DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024 @@ -36,6 +36,26 @@ def ws_max_size() -> int | None: return None if value <= 0 else value +# Binary relay frame: JSON header + raw body in one WebSocket binary message, +# so activation bodies travel as bytes instead of base64 inside JSON. Same wire +# format as meshnet_relay.server — duplicated because node and relay ship as +# independent distributions. +BINARY_FRAME_MAGIC = b"MRF1" + + +def encode_binary_frame(header: dict, body: bytes) -> bytes: + header_bytes = json.dumps(header, separators=(",", ":")).encode() + return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body + + +def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]: + if len(frame) < 8 or frame[:4] != BINARY_FRAME_MAGIC: + raise ValueError("not a meshnet binary relay frame") + header_len = int.from_bytes(frame[4:8], "big") + header = json.loads(frame[8:8 + header_len].decode()) + return header, bytes(frame[8 + header_len:]) + + @dataclass(frozen=True) class RelayBridgeInfo: peer_id: str @@ -125,7 +145,9 @@ class RelayHttpBridge: while self._running: try: - with wsc.connect(self.relay_url, open_timeout=5, max_size=ws_max_size()) as ws: + with wsc.connect( + self.relay_url, open_timeout=5, max_size=ws_max_size(), compression=None, + ) as ws: self._ws = ws self._connected.set() ws.send(json.dumps(_make_envelope( @@ -138,6 +160,17 @@ class RelayHttpBridge: raw = ws.recv(timeout=1) except TimeoutError: continue + if isinstance(raw, (bytes, bytearray)): + try: + payload, body = decode_binary_frame(bytes(raw)) + except (ValueError, json.JSONDecodeError): + continue + if payload.get("target_peer") not in {None, self.peer_id}: + continue + if self._executor is None: + break + self._executor.submit(self._process_request, payload, body) + continue try: envelope = json.loads(raw) except (TypeError, json.JSONDecodeError): @@ -176,23 +209,41 @@ class RelayHttpBridge: log.debug("relay bridge send failed (request orphaned): %s", exc) return False - def _process_request(self, payload: dict) -> None: + def _send_binary_response_frame(self, header: dict, body: bytes) -> bool: + """Send one binary response frame; False if the socket is gone.""" + ws = self._ws + if ws is None: + return False + frame = encode_binary_frame(header, body) + try: + with self._send_lock: + ws.send(frame) + return True + except Exception as exc: + log.debug("relay bridge binary send failed (request orphaned): %s", exc) + return False + + def _process_request(self, payload: dict, binary_body: bytes | None = None) -> None: request_id = str(payload.get("request_id") or "") method = str(payload.get("method") or "POST").upper() path = str(payload.get("path") or "/") headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {} + binary_mode = binary_body is not None req_suffix = f" request_id={request_id}" if request_id else "" print(f" [node] relay {method} {path}{req_suffix}", flush=True) - # body_base64 carries binary data (e.g. bfloat16 activation tensors) safely. - # Fallback to text "body" for backward-compat with non-binary requests. - body_b64 = payload.get("body_base64") - if body_b64: - data = base64.b64decode(body_b64) + if binary_mode: + data = binary_body else: - body_text = payload.get("body") or "" - data = body_text.encode() if isinstance(body_text, str) else bytes(body_text) + # Legacy JSON request: body_base64 carries binary data, text "body" + # covers non-binary requests. + body_b64 = payload.get("body_base64") + if body_b64: + data = base64.b64decode(body_b64) + else: + body_text = payload.get("body") or "" + data = body_text.encode() if isinstance(body_text, str) else bytes(body_text) url = f"{self.local_base_url}{path}" req = urllib.request.Request(url, data=data, headers=headers, method=method) @@ -205,6 +256,13 @@ class RelayHttpBridge: return resp_bytes = resp.read() # Forward all X-Meshnet-* headers so the caller can reconstruct the activation. + if binary_mode: + self._send_binary_response_frame({ + "request_id": request_id, + "status": resp.status, + "headers": resp_headers, + }, resp_bytes) + return is_binary = "octet-stream" in content_type result: dict = { "request_id": request_id, diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 4dcfdac..3bcbb9c 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -109,25 +109,33 @@ def _relay_hop( ) -> tuple[int, dict[str, str], bytes]: """Send a single HTTP-shaped request through a relay RPC WebSocket. - relay_addr is the wss://relay.../rpc/{peer_id} URL. + 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 + error responses and legacy peers still answer with JSON text frames. 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 ws_max_size + from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size request_id = f"{time.time_ns():x}" - payload = json.dumps({ + frame = encode_binary_frame({ "request_id": request_id, "method": "POST", "path": path, "headers": headers, - "body_base64": base64.b64encode(body).decode(), - }) - with wsc.connect(relay_addr, open_timeout=timeout, max_size=ws_max_size()) as ws: - ws.send(payload) + }, 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()} @@ -136,6 +144,20 @@ def _relay_hop( return status, resp_headers, resp_body +# Below this, zstd overhead outweighs the win (per-token decode bodies are ~KBs). +_COMPRESS_MIN_BYTES = 64 * 1024 + + +def _maybe_compress_activation(body: bytes) -> tuple[bytes, str | None]: + """zstd-compress large activation bodies; returns (wire_body, encoding).""" + if len(body) < _COMPRESS_MIN_BYTES: + return body, None + try: + return _compress_body(body, "zstd"), "zstd" + except Exception: + return body, None + + def _is_cache_miss_body(body: bytes) -> bool: try: return json.loads(body).get("error") == "cache_miss" @@ -892,6 +914,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): + (f" relay={relay_addr}" if relay_addr else ""), flush=True, ) + wire_body, wire_encoding = _maybe_compress_activation(current_body) headers: dict[str, str] = { "Content-Type": "application/octet-stream", "X-Meshnet-Wire": _WIRE_VERSION, @@ -903,6 +926,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "X-Meshnet-Hop-Index": str(hop_index), "X-Meshnet-Start-Layer": str(start_layer), } + if wire_encoding: + headers["X-Meshnet-Encoding"] = wire_encoding if cache_mode: headers["X-Meshnet-Cache"] = cache_mode past_len = getattr(payload, "past_len", None) @@ -915,7 +940,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if relay_addr: try: status, resp_headers, resp_body = _relay_hop( - relay_addr, "/forward", current_body, headers, timeout=120.0, + relay_addr, "/forward", wire_body, headers, timeout=120.0, ) if status == 409 and _is_cache_miss_body(resp_body): raise _PipelineCacheMiss(node_url) @@ -938,7 +963,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if not relay_addr: req = urllib.request.Request( f"{node_url}/forward", - data=current_body, + data=wire_body, headers=headers, method="POST", ) @@ -970,7 +995,13 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): # Binary activation — update and forward to next node shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape)) current_shape = _parse_shape(shape_header) - current_body = resp_body + try: + current_body = _decompress_body( + resp_body, resp_headers.get("x-meshnet-encoding") + ) + except ValueError as exc: + print(f" [node] pipeline hop {hop_index} bad response encoding: {exc}", flush=True) + return f"pipeline error at {node_url}: {exc}", None current_attn = resp_headers.get("x-meshnet-attn-mask") current_pos = resp_headers.get("x-meshnet-position-ids") return "", None diff --git a/packages/relay/meshnet_relay/server.py b/packages/relay/meshnet_relay/server.py index 26369b1..270c495 100644 --- a/packages/relay/meshnet_relay/server.py +++ b/packages/relay/meshnet_relay/server.py @@ -41,6 +41,26 @@ def ws_max_size() -> int | None: return None if value <= 0 else value +# Binary relay frame: JSON header (request/response metadata) + raw body in one +# WebSocket binary message. Activation bodies stay bytes end to end — no base64 +# inflation, no JSON re-encode of megabytes per hop. Text JSON frames remain the +# control plane (gossip, peer-register, streamed SSE responses). +BINARY_FRAME_MAGIC = b"MRF1" + + +def encode_binary_frame(header: dict, body: bytes) -> bytes: + header_bytes = json.dumps(header, separators=(",", ":")).encode() + return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body + + +def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]: + if len(frame) < 8 or frame[:4] != BINARY_FRAME_MAGIC: + raise ValueError("not a meshnet binary relay frame") + header_len = int.from_bytes(frame[4:8], "big") + header = json.loads(frame[8:8 + header_len].decode()) + return header, bytes(frame[8 + header_len:]) + + class RelayServer: """Async WebSocket relay server that runs in a background thread. @@ -118,6 +138,9 @@ class RelayServer: self.port, ssl=ssl_ctx, max_size=ws_max_size(), + # Bulk payloads are zstd-compressed at the pipeline layer already; + # per-message deflate would recompress them on every hop for nothing. + compression=None, ) # Record actual port after bind for sock in server.sockets or []: @@ -162,6 +185,17 @@ class RelayServer: try: async for raw in ws: + # Binary frames are relay-http-response bodies from a bridge — + # route them to the waiting rpc requester as-is, never fan out. + if isinstance(raw, (bytes, bytearray)): + try: + header, _ = decode_binary_frame(bytes(raw)) + except (ValueError, json.JSONDecodeError): + continue + queue = self._pending_rpc.get(header.get("request_id")) + if queue is not None: + queue.put_nowait(bytes(raw)) + continue try: envelope = json.loads(raw) except (json.JSONDecodeError, TypeError): @@ -252,14 +286,27 @@ class RelayServer: try: raw = await asyncio.wait_for(ws_requester.recv(), timeout=30.0) - payload = json.loads(raw) + 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 - request_id = str(payload.get("request_id") or uuid.uuid4()) - payload["request_id"] = request_id - payload["target_peer"] = target_peer_id queue: asyncio.Queue = asyncio.Queue() self._pending_rpc[request_id] = queue overall_timeout = 310.0 @@ -267,15 +314,11 @@ class RelayServer: loop = asyncio.get_running_loop() deadline = loop.time() + overall_timeout try: - await target.ws.send(json.dumps({ - "topic": "relay-http-request", - "version": 1, - "from_peer": "relay", - "payload": payload, - })) + 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. + # is a complete legacy single response. A binary frame is always a + # complete single response. while True: remaining = deadline - loop.time() if remaining <= 0: @@ -283,6 +326,9 @@ class RelayServer: 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 diff --git a/tests/test_gossip_and_relay.py b/tests/test_gossip_and_relay.py index b42a557..ec01717 100644 --- a/tests/test_gossip_and_relay.py +++ b/tests/test_gossip_and_relay.py @@ -372,12 +372,53 @@ def test_relay_rpc_round_trips_http_request_to_peer(): assert json.loads(response["body"]) == {"ok": True, "path": "/v1/chat/completions"} +def test_binary_relay_frame_codecs_interoperate(): + """Node and relay ship the same binary frame format as separate copies.""" + import os + + from meshnet_node import relay_bridge + from meshnet_relay import server as relay_server + + header = {"request_id": "r-1", "status": 200, "headers": {"X": "y"}} + body = os.urandom(4096) + for encoder, decoder in ( + (relay_bridge.encode_binary_frame, relay_server.decode_binary_frame), + (relay_server.encode_binary_frame, relay_bridge.decode_binary_frame), + ): + assert decoder(encoder(header, body)) == (header, body) + + try: + relay_server.decode_binary_frame(b"not a frame") + except ValueError: + pass + else: + raise AssertionError("garbage bytes must not decode as a binary frame") + + +def test_activation_compression_round_trips_and_skips_small_bodies(): + """Pipeline hops zstd-compress large activations; tiny decode bodies pass raw.""" + import os + + from meshnet_node.server import _decompress_body + from meshnet_node.torch_server import _maybe_compress_activation + + small = os.urandom(1024) + assert _maybe_compress_activation(small) == (small, None) + + large = os.urandom(96 * 1024) * 2 # repetition so zstd actually shrinks it + wire, encoding = _maybe_compress_activation(large) + assert encoding == "zstd" + assert len(wire) < len(large) + assert _decompress_body(wire, encoding) == large + + def test_relay_rpc_carries_activation_sized_frames(): """A >1 MiB activation body must survive the full relay round trip. Regression: the websockets library caps frames at 1 MiB by default, so prefill activations forwarded via /rpc/ died with close code 1009 at every hop (requester → relay, relay → bridge, bridge → relay → requester). + The body now travels as binary frames — raw bytes, no base64. """ import http.server import os diff --git a/tests/test_kv_cache_distributed.py b/tests/test_kv_cache_distributed.py index 66b25ca..22b5827 100644 --- a/tests/test_kv_cache_distributed.py +++ b/tests/test_kv_cache_distributed.py @@ -278,6 +278,43 @@ def test_session_is_stable_and_decode_payloads_are_single_token(): assert head_backend.released == [session_id] +def test_large_prefill_activation_survives_zstd_compressed_hop(): + """A prefill body above _COMPRESS_MIN_BYTES travels the hop zstd-compressed. + + The head compresses and sets X-Meshnet-Encoding; the tail's /forward must + decompress before shape validation, so a passing generation proves the + compressed round trip (a mishandled encoding fails validation with 400). + """ + + class _BigHeadBackend(_CachedHeadBackend): + def encode_prompt(self, prompt, session_id=None): + self.prefills.append(session_id) + if session_id: + self._seq[session_id] = 2048 + return TensorPayload( + body=b"\x00" * (1 * 2048 * 32 * 2), # 128 KiB, above the zstd threshold + shape=[1, 2048, 32], + attention_mask_header=None, + position_ids_header=None, + ) + + head_backend = _BigHeadBackend() + tail_backend = _CachedTailBackend([(" a", 1), (" b", 2)]) + head = TorchNodeServer(backend=head_backend, tracker_mode=True) + tail = TorchNodeServer(backend=tail_backend) + head_port = head.start() + tail_port = tail.start() + try: + content = _chat_once(head_port, tail_port, max_tokens=2) + finally: + head.stop() + tail.stop() + + assert content == " a b" + assert tail_backend.calls[0]["mode"] == "prefill" + assert tail_backend.calls[0]["shape"] == [1, 2048, 32] + + def test_eos_token_id_stops_generation(): head_backend = _CachedHeadBackend() tail_backend = _CachedTailBackend([(" a", 1), ("", 99)])