From 3d264a500aad4731eedbd107019f52b5a906ce04 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Thu, 9 Jul 2026 20:46:29 +0200 Subject: [PATCH] inference fixes --- deploy/portainer/README.md | 1 + packages/node/meshnet_node/relay_bridge.py | 18 +++++- packages/node/meshnet_node/torch_server.py | 4 +- packages/relay/meshnet_relay/server.py | 18 ++++++ packages/tracker/meshnet_tracker/server.py | 19 ++++++- tests/test_gossip_and_relay.py | 64 +++++++++++++++++++++- 6 files changed, 120 insertions(+), 4 deletions(-) diff --git a/deploy/portainer/README.md b/deploy/portainer/README.md index 653cd11..7077c41 100644 --- a/deploy/portainer/README.md +++ b/deploy/portainer/README.md @@ -218,6 +218,7 @@ Useful optional variables: PUBLIC_RELAY_URL=wss://ai.neuron.d-popov.com/ws HEARTBEAT_TIMEOUT=30 ENABLE_BILLING_DB=1 +MESHNET_WS_MAX_BYTES=268435456 # relay WebSocket frame cap (default 256 MiB; <=0 = unlimited) STARTING_CREDIT=1 DEVNET_TOPUP=1 ``` diff --git a/packages/node/meshnet_node/relay_bridge.py b/packages/node/meshnet_node/relay_bridge.py index af68e43..02e01ad 100644 --- a/packages/node/meshnet_node/relay_bridge.py +++ b/packages/node/meshnet_node/relay_bridge.py @@ -19,6 +19,22 @@ log = logging.getLogger(__name__) DEFAULT_MAX_CONCURRENCY = 8 +# Activation tensors ride the relay as base64 inside one JSON frame, so the +# websockets default of 1 MiB rejects any real prefill (close code 1009). +DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024 + + +def ws_max_size() -> int | None: + """Max inbound WebSocket frame size; MESHNET_WS_MAX_BYTES<=0 means unlimited.""" + raw = os.environ.get("MESHNET_WS_MAX_BYTES", "").strip() + if not raw: + return DEFAULT_WS_MAX_BYTES + try: + value = int(raw) + except ValueError: + return DEFAULT_WS_MAX_BYTES + return None if value <= 0 else value + @dataclass(frozen=True) class RelayBridgeInfo: @@ -109,7 +125,7 @@ class RelayHttpBridge: while self._running: try: - with wsc.connect(self.relay_url, open_timeout=5) as ws: + with wsc.connect(self.relay_url, open_timeout=5, max_size=ws_max_size()) as ws: self._ws = ws self._connected.set() ws.send(json.dumps(_make_envelope( diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index d03315e..4dcfdac 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -115,6 +115,8 @@ def _relay_hop( """ import websockets.sync.client as wsc # type: ignore[import] + from .relay_bridge import ws_max_size + request_id = f"{time.time_ns():x}" payload = json.dumps({ "request_id": request_id, @@ -123,7 +125,7 @@ def _relay_hop( "headers": headers, "body_base64": base64.b64encode(body).decode(), }) - with wsc.connect(relay_addr, open_timeout=timeout) as ws: + with wsc.connect(relay_addr, open_timeout=timeout, max_size=ws_max_size()) as ws: ws.send(payload) raw = ws.recv(timeout=timeout) resp = json.loads(raw) diff --git a/packages/relay/meshnet_relay/server.py b/packages/relay/meshnet_relay/server.py index c1f7ba0..26369b1 100644 --- a/packages/relay/meshnet_relay/server.py +++ b/packages/relay/meshnet_relay/server.py @@ -15,6 +15,7 @@ from __future__ import annotations import asyncio import json import logging +import os import threading import uuid from pathlib import Path @@ -23,6 +24,22 @@ from .peer_registry import PeerRegistry log = logging.getLogger(__name__) +# Activation tensors ride the relay as base64 inside one JSON frame, so the +# websockets default of 1 MiB rejects any real prefill (close code 1009). +DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024 + + +def ws_max_size() -> int | None: + """Max inbound WebSocket frame size; MESHNET_WS_MAX_BYTES<=0 means unlimited.""" + raw = os.environ.get("MESHNET_WS_MAX_BYTES", "").strip() + if not raw: + return DEFAULT_WS_MAX_BYTES + try: + value = int(raw) + except ValueError: + return DEFAULT_WS_MAX_BYTES + return None if value <= 0 else value + class RelayServer: """Async WebSocket relay server that runs in a background thread. @@ -100,6 +117,7 @@ class RelayServer: self.host, self.port, ssl=ssl_ctx, + max_size=ws_max_size(), ) # Record actual port after bind for sock in server.sockets or []: diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 64b2b7d..94aa0b1 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -1508,6 +1508,23 @@ def _assign_redundant_managed_nodes( _emit_shard_change_directives(node, model, previous_range, preset) +# Activation tensors ride the relay as base64 inside one JSON frame, so the +# websockets default of 1 MiB rejects any real prefill (close code 1009). +DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024 + + +def _ws_max_size() -> int | None: + """Max inbound WebSocket frame size; MESHNET_WS_MAX_BYTES<=0 means unlimited.""" + raw = os.environ.get("MESHNET_WS_MAX_BYTES", "").strip() + if not raw: + return DEFAULT_WS_MAX_BYTES + try: + value = int(raw) + except ValueError: + return DEFAULT_WS_MAX_BYTES + return None if value <= 0 else value + + def _relay_http_request_frames( relay_addr: str, path: str, @@ -1537,7 +1554,7 @@ def _relay_http_request_frames( request_id = str(uuid.uuid4()) deadline = time.monotonic() + timeout try: - with wsc.connect(relay_addr, open_timeout=10, close_timeout=5) as ws: + with wsc.connect(relay_addr, open_timeout=10, close_timeout=5, max_size=_ws_max_size()) as ws: if ws_holder is not None: if ws_lock is not None: with ws_lock: diff --git a/tests/test_gossip_and_relay.py b/tests/test_gossip_and_relay.py index bb6b1ad..b42a557 100644 --- a/tests/test_gossip_and_relay.py +++ b/tests/test_gossip_and_relay.py @@ -372,6 +372,68 @@ 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_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). + """ + import http.server + import os + + from meshnet_node.relay_bridge import RelayHttpBridge + from meshnet_node.torch_server import _relay_hop + from meshnet_relay.server import RelayServer + + body = os.urandom(2 * 1024 * 1024) # 2 MiB raw, ~2.7 MiB once base64-wrapped + + class EchoHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self): + data = self.rfile.read(int(self.headers.get("Content-Length", 0))) + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, *args): + pass + + local = http.server.HTTPServer(("127.0.0.1", 0), EchoHandler) + local_thread = threading.Thread(target=local.serve_forever, daemon=True) + local_thread.start() + + relay = RelayServer(host="127.0.0.1", port=0) + port = relay.start() + + bridge = RelayHttpBridge( + relay_url=f"ws://127.0.0.1:{port}/ws", + peer_id="big_peer", + local_base_url=f"http://127.0.0.1:{local.server_address[1]}", + advertised_addr="", + ) + try: + bridge.start() + assert bridge.wait_connected(timeout=5.0) + time.sleep(0.2) # let peer-register land in the relay registry + status, headers, resp_body = _relay_hop( + f"ws://127.0.0.1:{port}/rpc/big_peer", + "/forward", + body, + {"Content-Type": "application/octet-stream"}, + timeout=30.0, + ) + finally: + bridge.stop() + relay.stop() + local.shutdown() + local_thread.join(timeout=3) + + assert status == 200 + assert resp_body == body + + def test_node_relay_bridge_reconnects_after_failed_connection(monkeypatch): """Node-side relay bridge keeps retrying its outbound WebSocket connection.""" import websockets.sync.client as wsc # type: ignore[import] @@ -392,7 +454,7 @@ def test_node_relay_bridge_reconnects_after_failed_connection(monkeypatch): def recv(self, timeout=1): raise TimeoutError - def fake_connect(url, open_timeout=5): + def fake_connect(url, open_timeout=5, **kwargs): attempts.append((url, open_timeout)) if len(attempts) == 1: raise OSError("temporary relay outage")