popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

issues, chat FPS; optimisations

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-10 01:30:07 +03:00
parent 916f531e9d
commit f54ea100fb
17 changed files with 688 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
packages/node/meshnet_node/relay_bridge.py
View File

@@ -112,6 +112,8 @@ class RelayHttpBridge:
self._connected = threading.Event()
self._executor: ThreadPoolExecutor | None = None
self._send_lock = threading.Lock()
self._decode_log_lock = threading.Lock()
self._decode_steps: dict[str, int] = {}
self._ws = None
@property
@@ -230,8 +232,21 @@ class RelayHttpBridge:
headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {}
binary_mode = binary_body is not None
session = str(headers.get("X-Meshnet-Session") or "")
cache_mode = headers.get("X-Meshnet-Cache")
req_suffix = f" request_id={request_id}" if request_id else ""
print(f" [node] relay {method} {path}{req_suffix}", flush=True)
if path == "/forward" and cache_mode == "decode" and session:
with self._decode_log_lock:
steps = self._decode_steps.get(session, 0) + 1
self._decode_steps[session] = steps
if steps == 1 or steps % 32 == 0:
print(
f" [node] relay {method} {path} session={session[:8]} steps={steps}{req_suffix}",
flush=True,
)
else:
session_suffix = f" session={session[:8]}" if session else ""
print(f" [node] relay {method} {path}{session_suffix}{req_suffix}", flush=True)
if binary_mode:
data = binary_body

110
packages/node/meshnet_node/torch_server.py
View File

@@ -100,6 +100,62 @@ def _write_progress_line(state: list[bool], message: str, *, final: bool = False
sys.stdout.flush()
class _RelayHopClient:
"""Persistent relay connection scoped to one generation handler."""
def __init__(self, relay_addr: str, timeout: float = 120.0) -> None:
self.relay_addr = relay_addr
self.timeout = timeout
self._ws = None
def request(
self,
path: str,
body: bytes,
headers: dict[str, str],
) -> tuple[int, dict[str, str], bytes]:
import websockets.sync.client as wsc # type: ignore[import]
from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size
if self._ws is None:
self._ws = wsc.connect(
self.relay_addr,
open_timeout=self.timeout,
max_size=ws_max_size(),
compression=None,
)
request_id = f"{time.time_ns():x}"
frame = encode_binary_frame({
"request_id": request_id,
"method": "POST",
"path": path,
"headers": headers,
}, body)
self._ws.send(frame)
raw = self._ws.recv(timeout=self.timeout)
if isinstance(raw, (bytes, bytearray)):
resp_header, resp_body = decode_binary_frame(bytes(raw))
status = int(resp_header.get("status", 503))
resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()}
return status, resp_headers, resp_body
resp = json.loads(raw)
status = int(resp.get("status", 503))
resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()}
body_b64 = resp.get("body_base64")
resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode()
return status, resp_headers, resp_body
def close(self) -> None:
if self._ws is not None:
try:
self._ws.close()
except Exception:
pass
finally:
self._ws = None
def _relay_hop(
relay_addr: str,
path: str,
@@ -107,7 +163,7 @@ def _relay_hop(
headers: dict[str, str],
timeout: float = 120.0,
) -> tuple[int, dict[str, str], bytes]:
"""Send a single HTTP-shaped request through a relay RPC WebSocket.
"""Send one request through a short-lived relay client (compatibility API).
relay_addr is the wss://relay.../rpc/{peer_id} URL. The request and any
binary response travel as binary frames (JSON header + raw body); relay
@@ -115,33 +171,11 @@ def _relay_hop(
Returns (status, response_headers_lower, response_body).
Raises on connection failure so callers can fall back to direct.
"""
import websockets.sync.client as wsc # type: ignore[import]
from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size
request_id = f"{time.time_ns():x}"
frame = encode_binary_frame({
"request_id": request_id,
"method": "POST",
"path": path,
"headers": headers,
}, body)
with wsc.connect(
relay_addr, open_timeout=timeout, max_size=ws_max_size(), compression=None,
) as ws:
ws.send(frame)
raw = ws.recv(timeout=timeout)
if isinstance(raw, (bytes, bytearray)):
resp_header, resp_body = decode_binary_frame(bytes(raw))
status = int(resp_header.get("status", 503))
resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()}
return status, resp_headers, resp_body
resp = json.loads(raw)
status = int(resp.get("status", 503))
resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()}
body_b64 = resp.get("body_base64")
resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode()
return status, resp_headers, resp_body
client = _RelayHopClient(relay_addr, timeout)
try:
return client.request(path, body, headers)
finally:
client.close()
# Below this, zstd overhead outweighs the win (per-token decode bodies are ~KBs).
@@ -687,6 +721,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
progress_line = [False]
last_token_id: int | None = None
failure_reason: str | None = None
relay_clients: dict[str, _RelayHopClient] = {}
def _prefill_step() -> tuple[str, int | None]:
"""Full-sequence prefill: initial step and cache-miss recovery."""
@@ -698,6 +733,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
return self._run_downstream_pipeline(
payload, remaining_route, backend=backend,
session=session_id, cache_mode="prefill" if use_kv else None,
relay_clients=relay_clients,
)
for step in range(max_tokens):
@@ -708,6 +744,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
token_str, token_id = self._run_downstream_pipeline(
payload, remaining_route, backend=backend,
session=session_id, cache_mode="decode",
relay_clients=relay_clients,
)
except (KVCacheMiss, _PipelineCacheMiss) as miss:
# Evicted/restarted node or head lost its own session:
@@ -769,6 +806,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
backend.release_session(session_id)
except Exception:
pass
for relay_client in relay_clients.values():
relay_client.close()
if generated:
elapsed = time.monotonic() - gen_started
@@ -881,6 +920,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
backend: TorchModelShard | None = None,
session: str | None = None,
cache_mode: str | None = None,
relay_clients: dict[str, _RelayHopClient] | None = None,
) -> tuple[str, int | None]:
"""Forward an activation through the downstream route.
@@ -953,9 +993,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
headers["X-Meshnet-Position-Ids"] = current_pos
if relay_addr:
try:
status, resp_headers, resp_body = _relay_hop(
relay_addr, "/forward", wire_body, headers, timeout=120.0,
)
if relay_clients is None:
status, resp_headers, resp_body = _relay_hop(
relay_addr, "/forward", wire_body, headers, timeout=120.0,
)
else:
relay_client = relay_clients.setdefault(
relay_addr, _RelayHopClient(relay_addr, timeout=120.0),
)
status, resp_headers, resp_body = relay_client.request(
"/forward", wire_body, headers,
)
if status == 409 and _is_cache_miss_body(resp_body):
raise _PipelineCacheMiss(node_url)
if status >= 400: