dropp baes64 use binary

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,

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