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

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

View File

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

View File

@@ -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 []:

View File

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