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

misc

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 18:49:32 +03:00
parent e9a094b620
commit b1f08c45cd
6 changed files with 541 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

344
packages/tracker/meshnet_tracker/server.py
View File

@@ -35,6 +35,7 @@ import urllib.parse
import urllib.request
import uuid
from collections import deque
from dataclasses import dataclass, field
from importlib.resources import files
from pathlib import Path
from typing import Any
@@ -51,6 +52,7 @@ from .raft import RaftNode
_CONSOLE_LIMIT = 300
_PROXY_PROGRESS_LOG_INTERVAL = 5.0
def _preset_price_keys(name: str, preset: dict) -> set[str]:
@@ -1413,6 +1415,10 @@ def _relay_http_request_frames(
headers: dict[str, str],
timeout: float = 310.0,
idle_timeout: float = 120.0,
*,
cancel_event: threading.Event | None = None,
ws_holder: list[Any] | None = None,
ws_lock: threading.Lock | None = None,
):
"""Send an HTTP-shaped request through a relay RPC WebSocket, yielding
response frames until a terminal one (US-036).
@@ -1430,6 +1436,14 @@ def _relay_http_request_frames(
deadline = time.monotonic() + timeout
try:
with wsc.connect(relay_addr, open_timeout=10, close_timeout=5) as ws:
if ws_holder is not None:
if ws_lock is not None:
with ws_lock:
ws_holder.clear()
ws_holder.append(ws)
else:
ws_holder.clear()
ws_holder.append(ws)
ws.send(json.dumps({
"request_id": request_id,
"method": "POST",
@@ -1438,6 +1452,8 @@ def _relay_http_request_frames(
"body": body.decode(errors="replace"),
}))
while True:
if cancel_event is not None and cancel_event.is_set():
return
remaining = deadline - time.monotonic()
if remaining <= 0:
return
@@ -2076,7 +2092,15 @@ def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -
return None
def _tracker_log(server: "_TrackerHTTPServer", level: str, message: str, **fields: Any) -> None:
def _tracker_log(
server: "_TrackerHTTPServer",
level: str,
message: str,
*,
stdout: bool = True,
update_console_key: str | None = None,
**fields: Any,
) -> None:
event = {
"ts": time.time(),
"level": level,
@@ -2088,10 +2112,88 @@ def _tracker_log(server: "_TrackerHTTPServer", level: str, message: str, **field
},
}
with server.console_lock:
server.console_events.append(event)
extras = " ".join(f"{key}={value}" for key, value in event["fields"].items())
suffix = f" {extras}" if extras else ""
print(f"[tracker] {level}: {message}{suffix}", flush=True)
if update_console_key is not None:
updated = False
for existing in reversed(server.console_events):
if (
existing.get("message") == message
and existing.get("fields", {}).get("request_id") == update_console_key
):
existing["ts"] = event["ts"]
existing["fields"] = event["fields"]
updated = True
break
if not updated:
server.console_events.append(event)
else:
server.console_events.append(event)
if stdout:
extras = " ".join(f"{key}={value}" for key, value in event["fields"].items())
suffix = f" {extras}" if extras else ""
print(f"[tracker] {level}: {message}{suffix}", flush=True)
@dataclass
class _ActiveProxyContext:
request_id: str
cancel_event: threading.Event = field(default_factory=threading.Event)
upstream: Any | None = None
upstream_lock: threading.Lock = field(default_factory=threading.Lock)
relay_ws: Any | None = None
relay_ws_lock: threading.Lock = field(default_factory=threading.Lock)
def _register_active_proxy(server: "_TrackerHTTPServer", request_id: str) -> _ActiveProxyContext:
ctx = _ActiveProxyContext(request_id=request_id)
with server.active_proxies_lock:
server.active_proxies[request_id] = ctx
return ctx
def _unregister_active_proxy(server: "_TrackerHTTPServer", request_id: str) -> None:
with server.active_proxies_lock:
server.active_proxies.pop(request_id, None)
def _request_proxy_cancel(server: "_TrackerHTTPServer", request_id: str) -> bool:
with server.active_proxies_lock:
ctx = server.active_proxies.get(request_id)
if ctx is None:
return False
ctx.cancel_event.set()
def _close_resources() -> None:
with ctx.upstream_lock:
upstream = ctx.upstream
if upstream is not None:
try:
upstream.close()
except Exception:
pass
with ctx.relay_ws_lock:
relay_ws = ctx.relay_ws
if relay_ws is not None:
try:
relay_ws.close()
except Exception:
pass
threading.Thread(target=_close_resources, daemon=True).start()
return True
def _set_upstream_read_timeout(upstream: Any, timeout: float) -> None:
fp = getattr(upstream, "fp", None)
raw = getattr(fp, "raw", None) if fp is not None else None
sock = getattr(raw, "_sock", None) if raw is not None else None
if sock is not None:
sock.settimeout(timeout)
def _clear_proxy_progress_log_state(server: "_TrackerHTTPServer", request_id: str) -> None:
state = getattr(server, "_proxy_progress_log_state", None)
if state is not None:
state.pop(request_id, None)
def _tracker_log_proxy_progress(
@@ -2108,10 +2210,21 @@ def _tracker_log_proxy_progress(
) -> None:
elapsed = time.monotonic() - started
effective_elapsed = max(elapsed, 1e-6)
now = time.monotonic()
state = getattr(server, "_proxy_progress_log_state", None)
if state is None:
state = {}
server._proxy_progress_log_state = state
last_stdout = state.get(request_id)
stdout = last_stdout is None or (now - last_stdout) >= _PROXY_PROGRESS_LOG_INTERVAL
if stdout:
state[request_id] = now
_tracker_log(
server,
"info",
"proxy progress",
stdout=stdout,
update_console_key=request_id,
request_id=request_id,
model=model,
route_model=route_model,
@@ -2209,6 +2322,8 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
self.models_dir = models_dir
self.console_events = deque(maxlen=_CONSOLE_LIMIT)
self.console_lock = threading.Lock()
self.active_proxies: dict[str, _ActiveProxyContext] = {}
self.active_proxies_lock = threading.Lock()
class _TrackerHandler(http.server.BaseHTTPRequestHandler):
@@ -2364,6 +2479,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat":
self._handle_heartbeat(parts[3])
return
# /v1/proxy/requests/<request_id>/cancel
if (
len(parts) == 6
and parts[1] == "v1"
and parts[2] == "proxy"
and parts[3] == "requests"
and parts[5] == "cancel"
and parts[4]
):
self._handle_proxy_request_cancel(urllib.parse.unquote(parts[4]))
return
self.send_response(404)
self.end_headers()
@@ -2886,6 +3012,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if inflight_recorded:
_record_proxy_inflight(server, inflight_nodes, -1)
inflight_recorded = False
_unregister_active_proxy(server, request_id)
proxy_ctx = _register_active_proxy(server, request_id)
_tracker_log(
server,
@@ -2930,13 +3059,34 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
direct_endpoint=target_url,
)
started = time.monotonic()
relay_ws_holder: list[Any] = []
frames = _relay_http_request_frames(
node.relay_addr,
path="/v1/chat/completions",
body=raw_body,
headers=relay_headers,
cancel_event=proxy_ctx.cancel_event,
ws_holder=relay_ws_holder,
ws_lock=proxy_ctx.relay_ws_lock,
)
first = next(frames, None)
with proxy_ctx.relay_ws_lock:
proxy_ctx.relay_ws = relay_ws_holder[0] if relay_ws_holder else None
if proxy_ctx.cancel_event.is_set():
if self._finalize_proxy_cancel(
proxy_ctx=proxy_ctx,
server=server,
request_id=request_id,
started=started,
model=model,
route_model=route_model,
route_nodes=route_nodes,
api_key=api_key,
node_work=node_work,
body=body,
finish_proxy_inflight=finish_proxy_inflight,
):
return
if first is not None and first.get("stream"):
# Streamed response (US-036): forward SSE chunks as they arrive
# and run the same token accounting as the direct stream path.
@@ -2945,6 +3095,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
model, route_model, route_nodes, api_key, node_work,
request_body=body,
request_id=request_id,
proxy_ctx=proxy_ctx,
finish_proxy_inflight=finish_proxy_inflight,
)
finish_proxy_inflight()
return
@@ -2959,6 +3111,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
in_tokens, out_tokens = 0, 0
tokens = in_tokens + out_tokens
self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes)
_clear_proxy_progress_log_state(server, request_id)
_tracker_log(
server,
"info",
@@ -2989,7 +3142,69 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
try:
started = time.monotonic()
upstream = urllib.request.urlopen(req, timeout=300.0)
upstream_result: list[Any] = []
connect_errors: list[BaseException] = []
def _connect_upstream() -> None:
try:
upstream_result.append(urllib.request.urlopen(req, timeout=300.0))
except BaseException as exc:
connect_errors.append(exc)
connect_thread = threading.Thread(target=_connect_upstream, daemon=True)
connect_thread.start()
while connect_thread.is_alive():
if proxy_ctx.cancel_event.is_set():
connect_thread.join(timeout=310.0)
if upstream_result:
try:
upstream_result[0].close()
except Exception:
pass
if self._finalize_proxy_cancel(
proxy_ctx=proxy_ctx,
server=server,
request_id=request_id,
started=started,
model=model,
route_model=route_model,
route_nodes=route_nodes,
api_key=api_key,
node_work=node_work,
body=body,
finish_proxy_inflight=finish_proxy_inflight,
):
return
connect_thread.join(0.2)
if proxy_ctx.cancel_event.is_set():
if upstream_result:
try:
upstream_result[0].close()
except Exception:
pass
if self._finalize_proxy_cancel(
proxy_ctx=proxy_ctx,
server=server,
request_id=request_id,
started=started,
model=model,
route_model=route_model,
route_nodes=route_nodes,
api_key=api_key,
node_work=node_work,
body=body,
finish_proxy_inflight=finish_proxy_inflight,
):
return
if connect_errors:
raise connect_errors[0]
upstream = upstream_result[0]
with proxy_ctx.upstream_lock:
proxy_ctx.upstream = upstream
_set_upstream_read_timeout(upstream, 0.5)
_tracker_log(server, "info", "proxy connected", request_id=request_id, target_url=target_url)
except urllib.error.HTTPError as exc:
# Relay error status + body from node
@@ -3002,9 +3217,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self.wfile.write(err_body)
except BrokenPipeError:
pass
_clear_proxy_progress_log_state(server, request_id)
finish_proxy_inflight()
return
except Exception as exc:
_clear_proxy_progress_log_state(server, request_id)
if node.relay_addr:
_tracker_log(
server,
@@ -3045,8 +3262,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
client_gone = False
try:
while True:
line = upstream.readline()
if proxy_ctx.cancel_event.is_set():
break
try:
line = upstream.readline()
except TimeoutError:
continue
if not line:
if proxy_ctx.cancel_event.is_set():
break
break
if not client_gone:
try:
@@ -3071,6 +3295,22 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
stream_usage = usage
except (BrokenPipeError, ConnectionResetError):
pass
if self._finalize_proxy_cancel(
proxy_ctx=proxy_ctx,
server=server,
request_id=request_id,
started=started,
model=model,
route_model=route_model,
route_nodes=route_nodes,
api_key=api_key,
node_work=node_work,
body=body,
finish_proxy_inflight=finish_proxy_inflight,
observed_stream_tokens=observed_stream_tokens,
stream_usage=stream_usage,
):
return
elapsed = time.monotonic() - started
# Bill even on client disconnect — the nodes did the work.
# Observed stream chunks are authoritative for the upper bound;
@@ -3082,6 +3322,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
model, route_model, in_tokens + out_tokens, elapsed, route_nodes
)
tokens = in_tokens + out_tokens
_clear_proxy_progress_log_state(server, request_id)
_tracker_log(
server,
"info",
@@ -3113,6 +3354,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
in_tokens, out_tokens = 0, 0
tokens = in_tokens + out_tokens
self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes)
_clear_proxy_progress_log_state(server, request_id)
_tracker_log(
server,
"info",
@@ -3272,6 +3514,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
node_work: list,
request_body: dict,
request_id: str,
*,
proxy_ctx: _ActiveProxyContext | None = None,
finish_proxy_inflight: Any = None,
) -> None:
"""Forward a streamed relay response (US-036) to the client as SSE,
billing with the same accounting as the direct stream path."""
@@ -3285,6 +3530,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
observed_stream_tokens = 0
client_gone = False
for frame in itertools.chain([first], frames):
if proxy_ctx is not None and proxy_ctx.cancel_event.is_set():
break
chunk = frame.get("chunk") or ""
if not chunk:
continue
@@ -3312,6 +3559,26 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
)
if usage is not None:
stream_usage = usage
if (
proxy_ctx is not None
and finish_proxy_inflight is not None
and self._finalize_proxy_cancel(
proxy_ctx=proxy_ctx,
server=server,
request_id=request_id,
started=started,
model=model,
route_model=route_model,
route_nodes=route_nodes,
api_key=api_key,
node_work=node_work,
body=request_body,
finish_proxy_inflight=finish_proxy_inflight,
observed_stream_tokens=observed_stream_tokens,
stream_usage=stream_usage,
)
):
return
elapsed = time.monotonic() - started
in_tokens, out_tokens = _stream_billable_split(
observed_stream_tokens, stream_usage, request_body
@@ -3320,6 +3587,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
model, route_model, in_tokens + out_tokens, elapsed, route_nodes
)
tokens = in_tokens + out_tokens
_clear_proxy_progress_log_state(server, request_id)
_tracker_log(
server,
"info",
@@ -3765,6 +4033,68 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
events = [dict(event) for event in server.console_events]
self._send_json(200, {"events": events})
def _handle_proxy_request_cancel(self, request_id: str) -> None:
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if server.accounts is not None and not self._require_role("admin"):
return
if not _request_proxy_cancel(server, request_id):
self._send_json(404, {"error": f"no active proxy for request {request_id!r}"})
return
self._send_json(200, {"status": "canceled", "request_id": request_id})
def _finalize_proxy_cancel(
self,
*,
proxy_ctx: _ActiveProxyContext,
server: "_TrackerHTTPServer",
request_id: str,
started: float,
model: str,
route_model: str,
route_nodes: list,
api_key: str | None,
node_work: list,
body: dict,
finish_proxy_inflight: Any,
observed_stream_tokens: int = 0,
stream_usage: dict | None = None,
) -> bool:
if not proxy_ctx.cancel_event.is_set():
return False
elapsed = time.monotonic() - started
_clear_proxy_progress_log_state(server, request_id)
tokens = observed_stream_tokens
if observed_stream_tokens > 0:
in_tokens, out_tokens = _stream_billable_split(
observed_stream_tokens, stream_usage, body,
)
tokens = in_tokens + out_tokens
self._record_observed_throughput(
model, route_model, tokens, elapsed, route_nodes,
)
_tracker_log(
server,
"info",
"proxy canceled",
request_id=request_id,
model=model,
route_model=route_model,
tokens=tokens,
elapsed_seconds=round(elapsed, 4),
tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 and tokens > 0 else 0.0,
route=_node_route_summary(route_nodes),
)
if observed_stream_tokens > 0:
in_tokens, out_tokens = _stream_billable_split(
observed_stream_tokens, stream_usage, body,
)
self._bill_completed(
api_key, model, in_tokens + out_tokens, node_work,
input_tokens=in_tokens, output_tokens=out_tokens,
)
finish_proxy_inflight()
return True
def _handle_registry_wallets(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
if not self._require_role("admin"):