popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

dropp baes64 use binary

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-09 22:40:43 +02:00
parent 3d264a500a
commit e30272e83f
5 changed files with 244 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
packages/node/meshnet_node/relay_bridge.py
View File

@@ -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,

51
packages/node/meshnet_node/torch_server.py
View File

@@ -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

68
packages/relay/meshnet_relay/server.py
View File

@@ -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