inference fixes

This commit is contained in:
Dobromir Popov
2026-07-09 20:46:29 +02:00
parent 2b000ce9c3
commit 3d264a500a
6 changed files with 120 additions and 4 deletions

View File

@@ -218,6 +218,7 @@ Useful optional variables:
PUBLIC_RELAY_URL=wss://ai.neuron.d-popov.com/ws PUBLIC_RELAY_URL=wss://ai.neuron.d-popov.com/ws
HEARTBEAT_TIMEOUT=30 HEARTBEAT_TIMEOUT=30
ENABLE_BILLING_DB=1 ENABLE_BILLING_DB=1
MESHNET_WS_MAX_BYTES=268435456 # relay WebSocket frame cap (default 256 MiB; <=0 = unlimited)
STARTING_CREDIT=1 STARTING_CREDIT=1
DEVNET_TOPUP=1 DEVNET_TOPUP=1
``` ```

View File

@@ -19,6 +19,22 @@ log = logging.getLogger(__name__)
DEFAULT_MAX_CONCURRENCY = 8 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) @dataclass(frozen=True)
class RelayBridgeInfo: class RelayBridgeInfo:
@@ -109,7 +125,7 @@ class RelayHttpBridge:
while self._running: while self._running:
try: 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._ws = ws
self._connected.set() self._connected.set()
ws.send(json.dumps(_make_envelope( ws.send(json.dumps(_make_envelope(

View File

@@ -115,6 +115,8 @@ def _relay_hop(
""" """
import websockets.sync.client as wsc # type: ignore[import] import websockets.sync.client as wsc # type: ignore[import]
from .relay_bridge import ws_max_size
request_id = f"{time.time_ns():x}" request_id = f"{time.time_ns():x}"
payload = json.dumps({ payload = json.dumps({
"request_id": request_id, "request_id": request_id,
@@ -123,7 +125,7 @@ def _relay_hop(
"headers": headers, "headers": headers,
"body_base64": base64.b64encode(body).decode(), "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) ws.send(payload)
raw = ws.recv(timeout=timeout) raw = ws.recv(timeout=timeout)
resp = json.loads(raw) resp = json.loads(raw)

View File

@@ -15,6 +15,7 @@ from __future__ import annotations
import asyncio import asyncio
import json import json
import logging import logging
import os
import threading import threading
import uuid import uuid
from pathlib import Path from pathlib import Path
@@ -23,6 +24,22 @@ from .peer_registry import PeerRegistry
log = logging.getLogger(__name__) 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: class RelayServer:
"""Async WebSocket relay server that runs in a background thread. """Async WebSocket relay server that runs in a background thread.
@@ -100,6 +117,7 @@ class RelayServer:
self.host, self.host,
self.port, self.port,
ssl=ssl_ctx, ssl=ssl_ctx,
max_size=ws_max_size(),
) )
# Record actual port after bind # Record actual port after bind
for sock in server.sockets or []: for sock in server.sockets or []:

View File

@@ -1508,6 +1508,23 @@ def _assign_redundant_managed_nodes(
_emit_shard_change_directives(node, model, previous_range, preset) _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( def _relay_http_request_frames(
relay_addr: str, relay_addr: str,
path: str, path: str,
@@ -1537,7 +1554,7 @@ def _relay_http_request_frames(
request_id = str(uuid.uuid4()) request_id = str(uuid.uuid4())
deadline = time.monotonic() + timeout deadline = time.monotonic() + timeout
try: 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_holder is not None:
if ws_lock is not None: if ws_lock is not None:
with ws_lock: with ws_lock:

View File

@@ -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"} 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/<peer> 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): def test_node_relay_bridge_reconnects_after_failed_connection(monkeypatch):
"""Node-side relay bridge keeps retrying its outbound WebSocket connection.""" """Node-side relay bridge keeps retrying its outbound WebSocket connection."""
import websockets.sync.client as wsc # type: ignore[import] 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): def recv(self, timeout=1):
raise TimeoutError raise TimeoutError
def fake_connect(url, open_timeout=5): def fake_connect(url, open_timeout=5, **kwargs):
attempts.append((url, open_timeout)) attempts.append((url, open_timeout))
if len(attempts) == 1: if len(attempts) == 1:
raise OSError("temporary relay outage") raise OSError("temporary relay outage")