Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

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.