diff --git a/.claude/memory/project-status.md b/.claude/memory/project-status.md index b261e02..8da1b6f 100644 --- a/.claude/memory/project-status.md +++ b/.claude/memory/project-status.md @@ -45,3 +45,4 @@ Historical handoff note: `/mnt/c/Users/popov/Downloads/neuron-tai-alpha-handoff- - Qwen3.6-35B-A3B reserve-based split is expected: an 79 GB CPU node may be assigned layers 0-36, and a second node fills 37-39. Do not "fix" this by bypassing the 20% assignment reserve unless the shard-planning policy changes. - Route hardening: tracker chat proxy and `/v1/route` diagnostics now use alias-aware preset node matching for split Qwen3.6 routes; dashboard derives grouped inference history from proxy route/complete console events and shows observed TPS after completion. - Live proxy hardening: model lookup trims outer whitespace before alias matching (`qwen3.6-35b-a3b ` resolves), and tracker route logs/dashboard queue depth combine heartbeat queue with tracker-local proxy in-flight counts so Postman-style bursts no longer show every selected route as queue `0`. +- Split-shard streaming hardening: Qwen3.6-style distributed generation now emits SSE chunks token-by-token from the head node instead of buffering all generated text until completion. Tracker direct/relay stream proxy logs `proxy progress` with live tokens/TPS, dashboard Inference history shows currently processing requests with live TPS/tokens/queue, and relay stream completion no longer references an undefined `session_id`. diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index a78b818..50a4985 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -342,6 +342,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): generated: list[str] = [] current_text = prompt_text + stream_emit = None + if stream: + stream_emit = self._start_openai_stream(model_name) + for _ in range(max_tokens): try: payload = backend.encode_prompt(current_text) @@ -357,9 +361,14 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if eos_token and token_str == eos_token: break generated.append(token_str) + if stream_emit is not None: + stream_emit(token_str) current_text = current_text + token_str result_text = "".join(generated) + if stream_emit is not None: + stream_emit(None) + return self._send_openai_response(result_text, model_name, stream, messages) def _get_remaining_route(self, model: str) -> list[dict]: @@ -526,6 +535,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): def _stream_openai_response(self, token_iter, model: str) -> None: """Stream tokens from an iterator as SSE chunks.""" + emit = self._start_openai_stream(model) + for token_text in token_iter: + if not token_text: + continue + emit(token_text) + emit(None) + + def _start_openai_stream(self, model: str): + """Open an OpenAI-compatible SSE response and return a token emitter.""" chunk_id = "chatcmpl-node" created = int(time.time()) self.send_response(200) @@ -537,7 +555,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): try: self.wfile.write(f"data: {data}\n\n".encode()) self.wfile.flush() - except BrokenPipeError: + except (BrokenPipeError, ConnectionResetError): pass _emit(json.dumps({ @@ -545,24 +563,27 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "model": model, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}], })) - for token_text in token_iter: - if not token_text: - continue + + def emit_token(token_text: str | None) -> None: + if token_text is None: + _emit(json.dumps({ + "id": chunk_id, "object": "chat.completion.chunk", "created": created, + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + })) + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + pass + return _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}], })) - _emit(json.dumps({ - "id": chunk_id, "object": "chat.completion.chunk", "created": created, - "model": model, - "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], - })) - try: - self.wfile.write(b"data: [DONE]\n\n") - self.wfile.flush() - except BrokenPipeError: - pass + + return emit_token def _send_openai_response( self, diff --git a/packages/tracker/meshnet_tracker/dashboard.html b/packages/tracker/meshnet_tracker/dashboard.html index 35c126a..2180367 100644 --- a/packages/tracker/meshnet_tracker/dashboard.html +++ b/packages/tracker/meshnet_tracker/dashboard.html @@ -324,6 +324,7 @@ function buildCallWallStates(events) { rec.started = e.ts; rec.model = f.model || f.route_model || "?"; rec.route = f.route || f.nodes; + rec.nodes = f.nodes; rec.stream = f.stream; } else if (msg === "proxy via relay" || msg === "proxy connected") { rec.status = "processing"; @@ -362,6 +363,12 @@ function callWallAgeSeconds(rec, nowSec) { return Math.max(0, nowSec - start); } +function callWallMaxQueue(rec) { + const nodes = rec.nodes || []; + const nodeQueues = Array.isArray(nodes) ? nodes.map(n => Number(n.queue_depth || 0)) : []; + return nodeQueues.length ? Math.max(...nodeQueues) : 0; +} + function renderCallWall(consoleData, stats) { const events = (consoleData && consoleData.events) || []; const nowSec = Date.now() / 1000; @@ -379,17 +386,20 @@ function renderCallWall(consoleData, stats) { const pending = active.filter(r => r.status === "pending").length; const processing = active.filter(r => r.status === "processing").length; const failedRecent = terminal.filter(r => r.status === "failed").length; + let queuedEstimate = 0; + for (const rec of active) queuedEstimate += Math.max(0, callWallMaxQueue(rec) - 1); let html = `
` + `hive tps (1h): ${esc(tps(hive.totalTps))} · samples: ${hive.samples} · ` + `active: ${processing} processing · ` + `${pending} pending` + + (queuedEstimate ? ` · queued estimate: ${queuedEstimate}` : "") + (failedRecent ? ` · ${failedRecent} recent failures` : "") + `
`; if (active.length) { - html += table(["status", "age", "model", "request", "tps", "tokens", "route / note"], active.map(rec => { + html += table(["status", "age", "model", "request", "live tps", "tokens", "queue", "route / note"], active.map(rec => { const statusCls = rec.status === "pending" ? "status-pending" : "status-processing"; const note = rec.warn || (rec.route ? short(String(rec.route), 28) : ""); return [ @@ -399,6 +409,7 @@ function renderCallWall(consoleData, stats) { esc(short(rec.id, 18)), `${esc(tps(rec.tps))}`, `${esc(String(rec.tokens ?? "—"))}`, + `${esc(String(callWallMaxQueue(rec)))}`, esc(note), ]; })); diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 18b4d71..da11d72 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -544,6 +544,7 @@ class _NodeEntry: "relay_addr", "cert_fingerprint", "peer_id", # heartbeat stats (reported by node, cumulative) "total_requests", "failed_requests", "queue_depth", "proxy_inflight", "uptime_seconds", + "current_requests", "status", # "ready" | "loading" "heartbeats_expected", "heartbeats_received", # dynamic reassignment queued by the tracker @@ -608,6 +609,7 @@ class _NodeEntry: self.failed_requests: int = 0 self.queue_depth: int = 0 self.proxy_inflight: int = 0 + self.current_requests: list[dict] = [] self.uptime_seconds: float = 0.0 self.status: str = "ready" self.heartbeats_expected: int = 0 @@ -634,6 +636,40 @@ def _effective_queue_depth(node: "_NodeEntry") -> int: return max(node.queue_depth, node.proxy_inflight) +_CURRENT_REQUEST_FIELDS = frozenset({ + "request_id", "model", "kind", "tokens", "tokens_per_sec", + "elapsed_seconds", "routing_complete", +}) + + +def _normalize_current_requests(items: object, *, limit: int = 32) -> list[dict]: + """Sanitize node-reported in-flight request snapshots from heartbeats.""" + if not isinstance(items, list): + return [] + out: list[dict] = [] + for item in items[:limit]: + if not isinstance(item, dict): + continue + request_id = item.get("request_id") + if not request_id: + continue + rec: dict = {"request_id": str(request_id)} + for key in _CURRENT_REQUEST_FIELDS: + if key == "request_id" or key not in item: + continue + value = item[key] + if key in {"tokens"}: + rec[key] = int(value) + elif key in {"tokens_per_sec", "elapsed_seconds"}: + rec[key] = float(value) + elif key == "routing_complete": + rec[key] = bool(value) + else: + rec[key] = str(value) + out.append(rec) + return out + + def _record_proxy_inflight( server: "_TrackerHTTPServer", nodes: list["_NodeEntry"], @@ -712,6 +748,7 @@ def _node_route_summary(nodes: list["_NodeEntry"]) -> list[dict]: "queue_depth": _effective_queue_depth(node), "heartbeat_queue_depth": node.queue_depth, "proxy_inflight": node.proxy_inflight, + "current_requests": list(node.current_requests), } for node in nodes ] @@ -992,6 +1029,7 @@ def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict: "queue_depth": _effective_queue_depth(node), "heartbeat_queue_depth": node.queue_depth, "proxy_inflight": node.proxy_inflight, + "current_requests": list(node.current_requests), "total_requests": node.total_requests, "heartbeat_success_rate": hb_rate, "inference_success_rate": inf_rate, @@ -1194,10 +1232,11 @@ def _billable_non_stream_split(payload: dict, request_body: dict) -> tuple[int, Completion stays capped by the request's max-tokens bound, as before. """ usage = _usage_split(payload) + prompt_estimate = _estimate_prompt_tokens(request_body) or 0 prompt = (usage or {}).get("prompt") completion = (usage or {}).get("completion") if prompt is None: - prompt = _estimate_prompt_tokens(request_body) or 0 + prompt = prompt_estimate if completion is None: total = (usage or {}).get("total") if total is not None: @@ -1205,8 +1244,9 @@ def _billable_non_stream_split(payload: dict, request_body: dict) -> tuple[int, else: completion = _observed_non_stream_completion_tokens(payload) limit = _requested_completion_token_limit(request_body) - if limit is not None: + if limit is not None and completion > limit: completion = min(completion, limit) + prompt = max(prompt, prompt_estimate) return max(0, prompt), max(0, completion) @@ -1776,6 +1816,7 @@ def _tracker_log_proxy_progress( relay: bool = False, ) -> None: elapsed = time.monotonic() - started + effective_elapsed = max(elapsed, 1e-6) _tracker_log( server, "info", @@ -1787,7 +1828,7 @@ def _tracker_log_proxy_progress( relay=relay or None, tokens=tokens, elapsed_seconds=round(elapsed, 4), - tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, + tokens_per_sec=round(tokens / effective_elapsed, 4) if tokens > 0 else 0.0, route=_node_route_summary(route_nodes), ) @@ -2476,6 +2517,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): node = max(candidates, key=lambda n: _effective_throughput(n, model)) target_url = f"{node.endpoint}/v1/chat/completions" request_id = str(body.get("id") or f"req-{time.time_ns():x}") + body["id"] = request_id + raw_body = json.dumps(body).encode() # Pre-resolve the downstream route so the first-shard node skips its own # tracker query. We already hold the full registry picture — no need for @@ -2608,6 +2651,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): first, frames, started, model, route_model, route_nodes, api_key, node_work, request_body=body, + request_id=request_id, ) finish_proxy_inflight() return @@ -2705,18 +2749,34 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.end_headers() stream_usage: dict | None = None observed_stream_tokens = 0 + client_gone = False try: while True: line = upstream.readline() if not line: break - self.wfile.write(line) - self.wfile.flush() + if not client_gone: + try: + self.wfile.write(line) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + # Keep draining upstream so completed node work is still billed. + client_gone = True observed, usage = _stream_line_tokens(line) observed_stream_tokens += observed + if observed: + _tracker_log_proxy_progress( + server, + request_id=request_id, + model=model, + route_model=route_model, + tokens=observed_stream_tokens, + started=started, + route_nodes=route_nodes, + ) if usage is not None: stream_usage = usage - except BrokenPipeError: + except (BrokenPipeError, ConnectionResetError): pass elapsed = time.monotonic() - started # Bill even on client disconnect — the nodes did the work. @@ -2786,7 +2846,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.end_headers() try: self.wfile.write(resp_body) - except BrokenPipeError: + except (BrokenPipeError, ConnectionResetError): pass finish_proxy_inflight() @@ -2805,8 +2865,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): can refine this later without changing the external stats shape. """ server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.stats is None or total_tokens <= 0 or elapsed_seconds <= 0: + if server.stats is None or total_tokens <= 0: return + elapsed_seconds = max(elapsed_seconds, 1e-6) models = [m for m in (requested_model, route_model) if m] if len(models) == 2 and models[0] == models[1]: models = [models[0]] @@ -2917,6 +2978,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): api_key: str | None, node_work: list, request_body: dict, + request_id: str, ) -> None: """Forward a streamed relay response (US-036) to the client as SSE, billing with the same accounting as the direct stream path.""" @@ -2925,6 +2987,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.send_header("Content-Type", headers.get("Content-Type", "text/event-stream; charset=utf-8")) self.send_header("Cache-Control", "no-cache") self.end_headers() + server: _TrackerHTTPServer = self.server # type: ignore[assignment] stream_usage: dict | None = None observed_stream_tokens = 0 client_gone = False @@ -2943,6 +3006,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): for line in data.splitlines(): observed, usage = _stream_line_tokens(line) observed_stream_tokens += observed + if observed: + _tracker_log_proxy_progress( + server, + request_id=request_id, + model=model, + route_model=route_model, + tokens=observed_stream_tokens, + started=started, + route_nodes=route_nodes, + relay=True, + ) if usage is not None: stream_usage = usage elapsed = time.monotonic() - started @@ -2953,12 +3027,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): model, route_model, in_tokens + out_tokens, elapsed, route_nodes ) tokens = in_tokens + out_tokens - server: _TrackerHTTPServer = self.server # type: ignore[assignment] _tracker_log( server, "info", "proxy complete", - request_id=session_id, + request_id=request_id, model=model, route_model=route_model, status=200, @@ -3279,6 +3352,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): entry.failed_requests = int(body["failed_requests"]) if "queue_depth" in body: entry.queue_depth = int(body["queue_depth"]) + if "current_requests" in body: + entry.current_requests = _normalize_current_requests(body["current_requests"]) if "uptime_seconds" in body: entry.uptime_seconds = float(body["uptime_seconds"]) if "status" in body and body["status"] in ("ready", "loading"): diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index c03050e..4e4928b 100644 --- a/tests/test_real_model_backend.py +++ b/tests/test_real_model_backend.py @@ -4,6 +4,8 @@ import json import os from pathlib import Path import sys +import threading +import time import types import urllib.request @@ -98,7 +100,7 @@ class _FakePipelineHeadBackend(_FakeBackend): tokenizer = _FakeChatTokenizer() def encode_prompt(self, prompt: str) -> TensorPayload: - assert prompt == "debug prompt" + assert prompt.startswith("debug prompt") return TensorPayload( body=b"\x00" * (1 * 6 * 8 * 2), shape=[1, 6, 8], @@ -117,6 +119,19 @@ class _FakePipelineTailBackend(_FakeTailBackend): return " token" +class _BlockingStreamingTailBackend(_FakeTailBackend): + def __init__(self, second_token_release: threading.Event) -> None: + self._release = second_token_release + self.calls = 0 + + def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): + self.calls += 1 + if self.calls == 1: + return " first" + self._release.wait(timeout=3.0) + return " second" + + def test_quantization_flag_validation(): assert validate_quantization("bfloat16") == "bfloat16" assert validate_quantization("int8") == "int8" @@ -303,6 +318,56 @@ def test_pipeline_hop_logs_are_enabled_with_debug(capsys): assert " [node] pipeline hop 0 returned text=' token'" in out +def test_split_shard_chat_streams_each_generated_token_incrementally(): + release_second = threading.Event() + head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) + tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second)) + head_port = head.start() + tail_port = tail.start() + response = None + try: + payload = json.dumps({ + "model": "fake-model", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + "max_tokens": 2, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{head_port}/v1/chat/completions", + data=payload, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": json.dumps([ + {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, + ]), + }, + method="POST", + ) + response = urllib.request.urlopen(req, timeout=5) + + first_token_line = "" + deadline = time.time() + 2.0 + while time.time() < deadline: + line = response.readline().decode() + if '"content": " first"' in line: + first_token_line = line + break + + assert first_token_line + assert not release_second.is_set() + release_second.set() + rest = response.read().decode() + finally: + release_second.set() + if response is not None: + response.close() + head.stop() + tail.stop() + + assert '"content": " second"' in rest + assert "data: [DONE]" in rest + + def test_int_tensor_header_serializes_torch_tensors(): torch = pytest.importorskip("torch")