issues, chat FPS; optimisations

2026-07-10 01:30:07 +03:00
parent 916f531e9d
commit f54ea100fb
17 changed files with 688 additions and 108 deletions
--- a/packages/node/meshnet_node/relay_bridge.py
+++ b/packages/node/meshnet_node/relay_bridge.py
@@ -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
--- a/packages/node/meshnet_node/torch_server.py
+++ b/packages/node/meshnet_node/torch_server.py
@@ -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:
--- a/packages/relay/meshnet_relay/server.py
+++ b/packages/relay/meshnet_relay/server.py
@@ -273,7 +273,7 @@ class RelayServer:
        )

    async def _handle_rpc(self, ws_requester, target_peer_id: str) -> None:
-        """Send one HTTP-shaped request to a connected peer and relay its response."""
+        """Relay sequential HTTP-shaped requests over one requester connection."""
        target = self._registry.get(target_peer_id)
        if target is None:
            await ws_requester.send(json.dumps({
@@ -284,64 +284,76 @@ class RelayServer:
            await ws_requester.close()
            return

-        try:
-            raw = await asyncio.wait_for(ws_requester.recv(), timeout=30.0)
-            if isinstance(raw, (bytes, bytearray)):
-                header, body = decode_binary_frame(bytes(raw))
-                request_id = str(header.get("request_id") or uuid.uuid4())
-                header["request_id"] = request_id
-                header["target_peer"] = target_peer_id
-                outbound: str | bytes = encode_binary_frame(header, body)
-            else:
-                payload = json.loads(raw)
-                request_id = str(payload.get("request_id") or uuid.uuid4())
-                payload["request_id"] = request_id
-                payload["target_peer"] = target_peer_id
-                outbound = json.dumps({
-                    "topic": "relay-http-request",
-                    "version": 1,
-                    "from_peer": "relay",
-                    "payload": payload,
-                })
-        except Exception:
-            await ws_requester.close(1003, "invalid relay rpc request")
-            return
+        while True:
+            try:
+                raw = await ws_requester.recv()
+            except Exception:
+                return

-        queue: asyncio.Queue = asyncio.Queue()
-        self._pending_rpc[request_id] = queue
-        overall_timeout = 310.0
-        idle_timeout = 120.0
-        loop = asyncio.get_running_loop()
-        deadline = loop.time() + overall_timeout
-        try:
-            await target.ws.send(outbound)
-            # Forward frames until a terminal one: streamed responses (US-036)
-            # end with {"stream": true, "done": true}; a frame without "stream"
-            # is a complete legacy single response. A binary frame is always a
-            # complete single response.
-            while True:
-                remaining = deadline - loop.time()
-                if remaining <= 0:
-                    raise asyncio.TimeoutError
-                frame = await asyncio.wait_for(
-                    queue.get(), timeout=min(idle_timeout, remaining)
-                )
-                if isinstance(frame, (bytes, bytearray)):
-                    await ws_requester.send(frame)
-                    break
-                await ws_requester.send(json.dumps(frame))
-                if not frame.get("stream") or frame.get("done"):
-                    break
-        except asyncio.TimeoutError:
-            await ws_requester.send(json.dumps({
-                "request_id": request_id,
-                "status": 504,
-                "headers": {"Content-Type": "application/json"},
-                "body": json.dumps({"error": "relay rpc timed out"}),
-            }))
-        finally:
-            self._pending_rpc.pop(request_id, None)
-            await ws_requester.close()
+            request_id: str | None = None
+            try:
+                if isinstance(raw, (bytes, bytearray)):
+                    header, body = decode_binary_frame(bytes(raw))
+                    request_id = str(header.get("request_id") or uuid.uuid4())
+                    header["request_id"] = request_id
+                    header["target_peer"] = target_peer_id
+                    outbound: str | bytes = encode_binary_frame(header, body)
+                else:
+                    payload = json.loads(raw)
+                    request_id = str(payload.get("request_id") or uuid.uuid4())
+                    payload["request_id"] = request_id
+                    payload["target_peer"] = target_peer_id
+                    outbound = json.dumps({
+                        "topic": "relay-http-request",
+                        "version": 1,
+                        "from_peer": "relay",
+                        "payload": payload,
+                    })
+            except Exception:
+                await ws_requester.close(1003, "invalid relay rpc request")
+                return
+
+            queue: asyncio.Queue = asyncio.Queue()
+            self._pending_rpc[request_id] = queue
+            overall_timeout = 310.0
+            idle_timeout = 120.0
+            loop = asyncio.get_running_loop()
+            deadline = loop.time() + overall_timeout
+            target = self._registry.get(target_peer_id)
+            try:
+                if target is None:
+                    await ws_requester.send(json.dumps({
+                        "request_id": request_id,
+                        "status": 503,
+                        "headers": {"Content-Type": "application/json"},
+                        "body": json.dumps({"error": f"peer {target_peer_id!r} disconnected"}),
+                    }))
+                    continue
+                await target.ws.send(outbound)
+                # Streamed responses end with done=true. Binary and legacy JSON
+                # responses are complete in one frame.
+                while True:
+                    remaining = deadline - loop.time()
+                    if remaining <= 0:
+                        raise asyncio.TimeoutError
+                    frame = await asyncio.wait_for(
+                        queue.get(), timeout=min(idle_timeout, remaining)
+                    )
+                    if isinstance(frame, (bytes, bytearray)):
+                        await ws_requester.send(frame)
+                        break
+                    await ws_requester.send(json.dumps(frame))
+                    if not frame.get("stream") or frame.get("done"):
+                        break
+            except asyncio.TimeoutError:
+                await ws_requester.send(json.dumps({
+                    "request_id": request_id,
+                    "status": 504,
+                    "headers": {"Content-Type": "application/json"},
+                    "body": json.dumps({"error": "relay rpc timed out"}),
+                }))
+            finally:
+                self._pending_rpc.pop(request_id, None)


 async def _broadcast(raw: str | bytes, peers: list) -> None: