more inference fixes

packages/node/meshnet_node/torch_server.py

@@ -686,6 +686,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
         last_gen_log = gen_started
         progress_line = [False]
         last_token_id: int | None = None
+        failure_reason: str | None = None
 
         def _prefill_step() -> tuple[str, int | None]:
             """Full-sequence prefill: initial step and cache-miss recovery."""
@@ -721,12 +722,15 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                     token_str, token_id = _prefill_step()
             except _PipelineCacheMiss as exc:
                 print(f"  [node] unexpected cache miss on prefill: {exc}", flush=True)
+                failure_reason = f"cache miss on prefill: {exc}"
                 break
             except Exception as exc:
                 print(f"  [node] distributed encode error: {exc}", flush=True)
+                failure_reason = f"distributed encode error: {exc}"
                 break
             # Stop on error responses or EOS.
             if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
+                failure_reason = token_str
                 break
             if token_id is not None and token_id in eos_ids:
                 break
@@ -778,6 +782,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             )
 
         result_text = "".join(generated)
+        # A failure before the first token is an upstream error, not an empty
+        # completion — tell the client instead of returning a blank 200.
+        if failure_reason and not generated:
+            if stream_emit is not None:
+                stream_emit(None, error=failure_reason)
+                return
+            self._send_json(502, {
+                "error": {"message": failure_reason, "type": "upstream_error"},
+            })
+            return
         if stream_emit is not None:
             stream_emit(None)
             return
@@ -1037,7 +1051,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
             "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
         }))
 
-        def emit_token(token_text: str | None) -> None:
+        def emit_token(token_text: str | None, *, error: str | None = None) -> None:
+            if error is not None:
+                # OpenAI-style mid-stream error frame; clients surface it
+                # instead of showing an empty completion.
+                _emit(json.dumps({"error": {"message": error, "type": "upstream_error"}}))
+                try:
+                    self.wfile.write(b"data: [DONE]\n\n")
+                    self.wfile.flush()
+                except (BrokenPipeError, ConnectionResetError):
+                    pass
+                return
             if token_text is None:
                 _emit(json.dumps({
                     "id": chunk_id, "object": "chat.completion.chunk", "created": created,

tests/test_kv_cache_distributed.py

@@ -9,6 +9,7 @@ on a real two-shard Qwen2.5-0.5B split.
 
 import json
 import os
+import urllib.error
 import urllib.request
 
 import pytest
@@ -278,6 +279,71 @@ def test_session_is_stable_and_decode_payloads_are_single_token():
     assert head_backend.released == [session_id]
 
 
+class _BrokenTailBackend(_CachedTailBackend):
+    """Tail whose forward always fails (e.g. missing compiler, OOM)."""
+
+    def forward_bytes(self, *args, **kwargs):
+        raise RuntimeError("Failed to find C compiler (test)")
+
+
+def test_pipeline_failure_before_first_token_returns_502():
+    """A dead hop must surface as an error, not an empty 200 completion."""
+    head = TorchNodeServer(backend=_CachedHeadBackend(), tracker_mode=True)
+    tail = TorchNodeServer(backend=_BrokenTailBackend([]))
+    head_port = head.start()
+    tail_port = tail.start()
+    try:
+        try:
+            _chat_once(head_port, tail_port, max_tokens=3)
+        except urllib.error.HTTPError as exc:
+            assert exc.code == 502
+            body = json.loads(exc.read())
+            assert "pipeline error" in body["error"]["message"]
+            assert "C compiler" in body["error"]["message"]
+        else:
+            raise AssertionError("expected HTTP 502 from a failed pipeline")
+    finally:
+        head.stop()
+        tail.stop()
+
+
+def test_pipeline_failure_in_stream_emits_error_frame():
+    """Streaming requests get an OpenAI-style error frame before [DONE]."""
+    head = TorchNodeServer(backend=_CachedHeadBackend(), tracker_mode=True)
+    tail = TorchNodeServer(backend=_BrokenTailBackend([]))
+    head_port = head.start()
+    tail_port = tail.start()
+    try:
+        payload = json.dumps({
+            "model": "fake-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "max_tokens": 3,
+            "stream": True,
+        }).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": 6},
+                ]),
+            },
+            method="POST",
+        )
+        with urllib.request.urlopen(req, timeout=10) as resp:
+            assert resp.status == 200
+            events = resp.read().decode().strip().split("\n\n")
+    finally:
+        head.stop()
+        tail.stop()
+
+    assert events[-1] == "data: [DONE]"
+    error_frame = json.loads(events[-2][len("data: "):])
+    assert error_frame["error"]["type"] == "upstream_error"
+    assert "C compiler" in error_frame["error"]["message"]
+
+
 def test_large_prefill_activation_survives_zstd_compressed_hop():
     """A prefill body above _COMPRESS_MIN_BYTES travels the hop zstd-compressed.