fixing streaming

2026-07-07 16:06:05 +02:00
parent 6fa69aecaa
commit 3eb7c6b93e
5 changed files with 210 additions and 83 deletions
--- 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`.
--- 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,
--- 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)}: <span class="num">${count}</span> active`)
-    .join(" &middot; ");
-  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)),
+      `<span class="num">${esc(tps(p.tokens_per_sec))}</span>`,
+      `<span class="num">${esc(String(p.tokens ?? 0))}</span>`,
+      `<span class="num">${esc(String(maxQueue))}</span>`,
+      p.stream ? "stream" : "json",
+    ]);
+  }
+  const active = Object.entries(activeByModel)
+    .map(([model, count]) => `${esc(model)}: <span class="num">${count}</span> active`)
+    .join(" &middot; ");
+  const liveSummary = active
+    ? `${active}${queuedEstimate ? ` &middot; queued estimate: <span class="num">${queuedEstimate}</span>` : ""}`
+    : "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) {
      `<span class="num">${esc(String(f.elapsed_seconds ?? "?"))}</span>`,
      f.stream ? "stream" : "json",
    ];
-  });
-  $("inference-history").innerHTML =
-    `<div class="dim" style="margin-bottom:6px">${active || "no active requests"}</div>` +
-    (rows.length ? table(["time", "model", "request", "tps", "tokens", "sec", "mode"], rows)
-                 : '<div class="empty">no completed inference requests</div>');
-}
+  });
+  $("inference-history").innerHTML =
+    `<div class="dim" style="margin-bottom:6px">${liveSummary}</div>` +
+    (activeRows.length ? table(["started", "model", "request", "live tps", "tokens", "queue", "mode"], activeRows.reverse()) : "") +
+    (rows.length ? table(["time", "model", "request", "tps", "tokens", "sec", "mode"], rows)
+                 : '<div class="empty">no completed inference requests</div>');
+}

 function renderConsole(data) {
  const events = (data && data.events) || [];
--- 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,
--- 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")