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