dropp baes64 use binary

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