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 e50dbd0..2d4d345 100644 --- a/packages/tracker/meshnet_tracker/dashboard.html +++ b/packages/tracker/meshnet_tracker/dashboard.html @@ -227,31 +227,53 @@ function renderThroughput(stats) { $("throughput").innerHTML = table(["node", "model", "tps (1h)", "samples"], rows); } -function renderInferenceHistory(data) { - const events = (data && data.events) || []; - const started = new Map(); - const completed = []; - for (const e of events) { - const f = e.fields || {}; - const id = f.request_id; - if (!id) continue; - if (e.message === "proxy route selected") { - started.set(id, e); - } else if (e.message === "proxy complete" || e.message === "proxy failed" || e.message === "direct proxy failed after relay") { - completed.push(e); - started.delete(id); - } - } - const activeByModel = {}; - for (const e of started.values()) { - const f = e.fields || {}; - const model = f.model || f.route_model || "?"; - activeByModel[model] = (activeByModel[model] || 0) + 1; - } - const active = Object.entries(activeByModel) - .map(([model, count]) => `${esc(model)}: ${count} active`) - .join(" · "); - const rows = completed.slice(-20).reverse().map(e => { +function renderInferenceHistory(data) { + const events = (data && data.events) || []; + const started = new Map(); + const progress = new Map(); + const completed = []; + for (const e of events) { + const f = e.fields || {}; + const id = f.request_id; + if (!id) continue; + if (e.message === "proxy route selected") { + started.set(id, e); + } else if (e.message === "proxy progress") { + progress.set(id, e); + } else if (e.message === "proxy complete" || e.message === "proxy failed" || e.message === "direct proxy failed after relay") { + completed.push(e); + started.delete(id); + progress.delete(id); + } + } + const activeByModel = {}; + let queuedEstimate = 0; + const activeRows = []; + for (const e of started.values()) { + const f = e.fields || {}; + const model = f.model || f.route_model || "?"; + activeByModel[model] = (activeByModel[model] || 0) + 1; + const p = (progress.get(f.request_id) || {}).fields || {}; + const nodeQueues = Array.isArray(f.nodes) ? f.nodes.map(n => Number(n.queue_depth || 0)) : []; + const maxQueue = nodeQueues.length ? Math.max(...nodeQueues) : 0; + queuedEstimate += Math.max(0, maxQueue - 1); + activeRows.push([ + new Date((e.ts || 0) * 1000).toLocaleTimeString(), + esc(short(model, 28)), + esc(short(f.request_id || "?", 18)), + `${esc(tps(p.tokens_per_sec))}`, + `${esc(String(p.tokens ?? 0))}`, + `${esc(String(maxQueue))}`, + p.stream ? "stream" : "json", + ]); + } + const active = Object.entries(activeByModel) + .map(([model, count]) => `${esc(model)}: ${count} active`) + .join(" · "); + const liveSummary = active + ? `${active}${queuedEstimate ? ` · queued estimate: ${queuedEstimate}` : ""}` + : "no active requests"; + const rows = completed.slice(-20).reverse().map(e => { const f = e.fields || {}; return [ new Date((e.ts || 0) * 1000).toLocaleTimeString(), @@ -262,12 +284,13 @@ function renderInferenceHistory(data) { `${esc(String(f.elapsed_seconds ?? "?"))}`, f.stream ? "stream" : "json", ]; - }); - $("inference-history").innerHTML = - `
${active || "no active requests"}
` + - (rows.length ? table(["time", "model", "request", "tps", "tokens", "sec", "mode"], rows) - : '
no completed inference requests
'); -} + }); + $("inference-history").innerHTML = + `
${liveSummary}
` + + (activeRows.length ? table(["started", "model", "request", "live tps", "tokens", "queue", "mode"], activeRows.reverse()) : "") + + (rows.length ? table(["time", "model", "request", "tps", "tokens", "sec", "mode"], rows) + : '
no completed inference requests
'); +} function renderConsole(data) { const events = (data && data.events) || []; diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 2e528c7..07c8039 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -1193,20 +1193,22 @@ def _billable_non_stream_split(payload: dict, request_body: dict) -> tuple[int, Prefers the response usage block; falls back to content estimates. Completion stays capped by the request's max-tokens bound, as before. """ - usage = _usage_split(payload) - prompt = (usage or {}).get("prompt") - completion = (usage or {}).get("completion") - if prompt is None: - prompt = _estimate_prompt_tokens(request_body) or 0 - if completion is None: - total = (usage or {}).get("total") - if total is not None: - completion = max(0, total - prompt) - else: - completion = _observed_non_stream_completion_tokens(payload) - limit = _requested_completion_token_limit(request_body) - if limit is not None: - completion = min(completion, limit) + 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 = prompt_estimate + if completion is None: + total = (usage or {}).get("total") + if total is not None: + completion = max(0, total - prompt) + else: + completion = _observed_non_stream_completion_tokens(payload) + limit = _requested_completion_token_limit(request_body) + 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 +1778,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 +1790,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), ) @@ -2608,6 +2611,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 @@ -2802,7 +2806,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.end_headers() try: self.wfile.write(resp_body) - except BrokenPipeError: + except (BrokenPipeError, ConnectionResetError): pass finish_proxy_inflight() @@ -2819,11 +2823,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): The tracker sees end-to-end request duration, not per-hop timings, so each hop gets the same route-level observation for now. Per-hop telemetry 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: - return - models = [m for m in (requested_model, route_model) if m] + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + 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]] for node in route_nodes: @@ -2930,19 +2935,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): model: str, route_model: str, route_nodes: list, - api_key: str | None, - node_work: list, - request_body: dict, - ) -> None: + 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.""" headers = first.get("headers") if isinstance(first.get("headers"), dict) else {} self.send_response(int(first.get("status", 200))) 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() - stream_usage: dict | None = None - observed_stream_tokens = 0 + 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 for frame in itertools.chain([first], frames): chunk = frame.get("chunk") or "" @@ -2956,11 +2963,22 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): except BrokenPipeError: # Keep draining frames — the nodes did the work; bill it. client_gone = True - for line in data.splitlines(): - observed, usage = _stream_line_tokens(line) - observed_stream_tokens += observed - if usage is not None: - stream_usage = usage + 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 in_tokens, out_tokens = _stream_billable_split( observed_stream_tokens, stream_usage, request_body @@ -2969,12 +2987,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( + _tracker_log( server, "info", "proxy complete", - request_id=session_id, + request_id=request_id, model=model, route_model=route_model, status=200, diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index de3b740..37c08bd 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 @@ -94,7 +96,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], @@ -113,6 +115,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" @@ -299,6 +314,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")