diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 61e950a..c31cdc3 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -205,6 +205,21 @@ def _max_assignable_layers( return min(total_layers, int((budget_bytes * 0.8) // layer_bytes)) +def _format_shard_label( + shard_start: int, + shard_end: int, + total_layers: int | None = None, + *, + model_name: str | None = None, +) -> str: + layer_count = shard_end - shard_start + 1 + if isinstance(total_layers, int) and total_layers > 0: + return f"layers {shard_start}–{shard_end} ({layer_count} of {total_layers})" + if model_name: + return f"layers {shard_start}–{shard_end} ({model_name})" + return f"layers {shard_start}–{shard_end}" + + def _shard_budget_line( memory_mb: int, memory_source: str, @@ -698,11 +713,7 @@ def run_startup( _node_start_time = time.monotonic() actual_port = node.start() total_layers = getattr(getattr(node, "backend", None), "total_layers", None) - if isinstance(total_layers, int) and total_layers > 0: - layer_count = shard_end - shard_start + 1 - shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}" - else: - shard_label = f"layers {shard_start}–{shard_end}" + shard_label = _format_shard_label(shard_start, shard_end, total_layers) public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) endpoint = f"http://{public_host}:{actual_port}" local_base_url = f"http://127.0.0.1:{actual_port}" @@ -891,14 +902,17 @@ def run_startup( except Exception as exc: setattr(node, "tracker_node_id", None) print(f" Warning: tracker registration failed: {exc}", flush=True) - shard_count = assigned_shard_end - assigned_shard_start + 1 + shard_label = _format_shard_label( + assigned_shard_start, + assigned_shard_end, + assigned_num_layers, + ) print( f"\n{'=' * 32}\n" f"meshnet-node ready (auto-joined)\n" f" Wallet: {address}\n" f" Model ID: {assigned_hf_repo}\n" - f" Shard: layers {assigned_shard_start}–{assigned_shard_end} " - f"({shard_count} of {assigned_num_layers})\n" + f" Shard: {shard_label}\n" f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" @@ -940,7 +954,16 @@ def run_startup( peers: list[dict] = assignment.get("peers", []) model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", []) assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization) - print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True) + model_layers_end = assignment.get("model_layers_end") + assigned_total_layers = ( + int(model_layers_end) + 1 + if model_layers_end is not None + else None + ) + print( + f" Shard: {_format_shard_label(shard_start, shard_end, assigned_total_layers, model_name=assigned_model)}", + flush=True, + ) # 4. Download shard print("Downloading shard...", flush=True) @@ -1021,12 +1044,18 @@ def run_startup( hw_str = device.upper() if gpu_name: hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)" + shard_label = _format_shard_label( + shard_start, + shard_end, + assigned_total_layers, + model_name=assigned_model, + ) print( f"\n{'=' * 32}\n" f"meshnet-node ready\n" f" Wallet: {address}\n" - f" Shard: layers {shard_start}-{shard_end} ({assigned_model})\n" - f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n" + f" Shard: {shard_label}\n" + f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n" f" Endpoint: {endpoint}\n" f" Node ID: {node_id}\n" f" Hardware: {hw_str}\n" diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index e8cf923..398ad6e 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -368,7 +368,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if stream: stream_emit = self._start_openai_stream(model_name) - for _ in range(max_tokens): + _GENERATION_LOG_INTERVAL = 5.0 + gen_started = time.monotonic() + last_gen_log = gen_started + + for step in range(max_tokens): try: payload = backend.encode_prompt(current_text) except Exception as exc: @@ -386,6 +390,21 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if stream_emit is not None: stream_emit(token_str) current_text = current_text + token_str + now = time.monotonic() + if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL: + print( + f" [node] generating step={step + 1}/{max_tokens} " + f"tokens={len(generated)} elapsed_s={now - gen_started:.1f}", + flush=True, + ) + last_gen_log = now + + if generated: + print( + f" [node] generation complete tokens={len(generated)} " + f"elapsed_s={time.monotonic() - gen_started:.1f}", + flush=True, + ) result_text = "".join(generated) if stream_emit is not None: diff --git a/packages/tracker/meshnet_tracker/dashboard.html b/packages/tracker/meshnet_tracker/dashboard.html index decdad3..18297ee 100644 --- a/packages/tracker/meshnet_tracker/dashboard.html +++ b/packages/tracker/meshnet_tracker/dashboard.html @@ -88,6 +88,9 @@ .status-processing { color:var(--accent); } .status-failed { color:var(--bad); } .status-complete { color:var(--ok); } + .status-canceled { color:var(--dim); } + button.btn-cancel { color:var(--dim); padding:0 5px; min-width:1.4em; line-height:1.2; } + button.btn-cancel:hover { color:var(--bad); border-color:var(--bad); } @@ -350,6 +353,13 @@ function buildCallWallStates(events) { rec.elapsed = f.elapsed_seconds; rec.stream = f.stream; rec.terminal = e; + } else if (msg === "proxy canceled") { + rec.status = "canceled"; + rec.model = rec.model || f.model || f.route_model || "?"; + rec.tokens = f.tokens; + rec.tps = f.tokens_per_sec; + rec.elapsed = f.elapsed_seconds; + rec.terminal = e; } else if (msg === "proxy failed" || msg === "direct proxy failed after relay") { rec.status = "failed"; rec.model = rec.model || f.model || f.route_model || "?"; @@ -379,7 +389,7 @@ function renderCallWall(consoleData, stats) { const terminal = []; for (const rec of states.values()) { if (rec.status === "pending" || rec.status === "processing") active.push(rec); - else if (rec.status === "complete" || rec.status === "failed") terminal.push(rec); + else if (rec.status === "complete" || rec.status === "failed" || rec.status === "canceled") terminal.push(rec); } active.sort((a, b) => (a.started || 0) - (b.started || 0)); terminal.sort((a, b) => (b.terminal && b.terminal.ts) - (a.terminal && a.terminal.ts)); @@ -401,10 +411,13 @@ function renderCallWall(consoleData, stats) { ``; if (active.length) { - html += table(["status", "age", "model", "request", "live tps", "tokens", "queue", "route / note"], active.map(rec => { + const canCancelProxies = isAdmin || !isLoggedIn; + const headers = ["status", "age", "model", "request", "live tps", "tokens", "queue", "route / note"]; + if (canCancelProxies) headers.push(""); + html += table(headers, active.map(rec => { const statusCls = rec.status === "pending" ? "status-pending" : "status-processing"; const note = rec.warn || (rec.route ? short(String(rec.route), 28) : ""); - return [ + const row = [ `${esc(rec.status)}`, `${esc(callWallAgeSeconds(rec, nowSec).toFixed(1))}s`, esc(short(rec.model || "?", 28)), @@ -414,6 +427,12 @@ function renderCallWall(consoleData, stats) { `${esc(String(callWallMaxQueue(rec)))}`, esc(note), ]; + if (canCancelProxies) { + row.push( + ``, + ); + } + return row; })); } else { html += '
no in-flight requests
'; @@ -422,10 +441,16 @@ function renderCallWall(consoleData, stats) { const historyRows = terminal.slice(0, 40).map(rec => { const e = rec.terminal || {}; const f = e.fields || {}; - const statusCls = rec.status === "failed" ? "status-failed" : "status-complete"; + const statusCls = rec.status === "failed" + ? "status-failed" + : rec.status === "canceled" + ? "status-canceled" + : "status-complete"; const detail = rec.status === "failed" ? esc(short(rec.error || "?", 40)) - : (f.stream ? "stream" : "json"); + : rec.status === "canceled" + ? "canceled" + : (f.stream ? "stream" : "json"); return [ new Date((e.ts || 0) * 1000).toLocaleTimeString(), `${esc(rec.status)}`, @@ -437,7 +462,7 @@ function renderCallWall(consoleData, stats) { detail, ]; }); - html += '
recent completed / failed
'; + html += '
recent completed / failed / canceled
'; html += historyRows.length ? table(["time", "status", "model", "request", "tps", "tokens", "sec", "detail"], historyRows) : '
no completed requests yet
'; @@ -956,6 +981,23 @@ async function renderAdminPanel() { $("admin").innerHTML = table(["account", "role", "keys", "balance (USDT)", "created"], rows); } +async function cancelProxyRequest(requestId) { + const r = await apiCall( + `/v1/proxy/requests/${encodeURIComponent(requestId)}/cancel`, + "POST", + {}, + ); + if (r.ok) refresh(); +} + +$("call-wall").addEventListener("click", (event) => { + const button = event.target.closest("[data-cancel-request]"); + if (!button) return; + event.preventDefault(); + const requestId = button.getAttribute("data-cancel-request"); + if (requestId) cancelProxyRequest(requestId); +}); + async function refresh() { $("self-url").textContent = location.host; const [raft, map, stats, models, consoleData, adminData] = await Promise.all([ diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index db330cd..1d7fbb7 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -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//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"): diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 269d0c4..1c4d57b 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -1116,7 +1116,7 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca assert captured_registration["vram_bytes"] == 6144 * 1024 * 1024 assert captured_registration["max_loaded_shards"] == 2 output = capsys.readouterr().out - assert "Shard: layers 0–23; 24 of 24" in output + assert "Shard: layers 0–23 (24 of 24)" in output assert "Node ID: node-test-123" in output diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index ef934a6..7e17cae 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -5,6 +5,7 @@ import json import threading import time import urllib.error +import urllib.parse import urllib.request import pytest @@ -501,6 +502,100 @@ def test_tracker_logs_stream_progress_before_request_completes(): node_thread.join(timeout=1.0) +def test_tracker_dashboard_can_cancel_inflight_proxy(): + chunk_sent = threading.Event() + release = threading.Event() + + class StreamingChatHandler(http.server.BaseHTTPRequestHandler): + def log_message(self, fmt, *args): + pass + + def do_POST(self): + if self.path != "/v1/chat/completions": + self.send_response(404) + self.end_headers() + return + self.rfile.read(int(self.headers.get("Content-Length", 0))) + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.end_headers() + payload = json.dumps({ + "choices": [{"delta": {"content": "hello world"}}], + }).encode() + self.wfile.write(b"data: " + payload + b"\n\n") + self.wfile.flush() + chunk_sent.set() + release.wait(timeout=3.0) + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + + node = http.server.HTTPServer(("127.0.0.1", 0), StreamingChatHandler) + node_thread = threading.Thread(target=node.serve_forever, daemon=True) + node_thread.start() + tracker = TrackerServer(heartbeat_timeout=60.0) + tracker_port = tracker.start() + response = None + request_id = None + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": f"http://127.0.0.1:{node.server_address[1]}", + "model": "cancel-proxy-model", "num_layers": 1, + "shard_start": 0, "shard_end": 0, + "hardware_profile": {}, "score": 1.0}, + ) + req = urllib.request.Request( + f"http://127.0.0.1:{tracker_port}/v1/chat/completions", + data=json.dumps({ + "model": "cancel-proxy-model", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + response = urllib.request.urlopen(req, timeout=3.0) + first_line = response.readline() + assert first_line.startswith(b"data:") + assert chunk_sent.wait(timeout=1.0) + + console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console") + selected = [ + event for event in console["events"] + if event["message"] == "proxy route selected" + ] + assert selected + request_id = selected[-1]["fields"]["request_id"] + + cancel = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/proxy/requests/{urllib.parse.quote(request_id, safe='')}/cancel", + {}, + ) + assert cancel["status"] == "canceled" + + deadline = time.time() + 5.0 + canceled_events = [] + while time.time() < deadline: + console = _get_json(f"http://127.0.0.1:{tracker_port}/v1/console") + canceled_events = [ + event for event in console["events"] + if event["message"] == "proxy canceled" + and event["fields"].get("request_id") == request_id + ] + if canceled_events: + break + time.sleep(0.05) + assert canceled_events + finally: + release.set() + if response is not None: + response.close() + tracker.stop() + node.shutdown() + node.server_close() + node_thread.join(timeout=1.0) + + def test_tracker_routes_hf_model_alias_from_quickstart(): """The documented qwen2.5-0.5b alias resolves a full HF repo registration.""" tracker = TrackerServer()