inference fixes
This commit is contained in:
@@ -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
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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 []:
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
Reference in New Issue
Block a user