issues, chat FPS; optimisations

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: