diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 3bcbb9c..b6cac85 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/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, diff --git a/tests/test_kv_cache_distributed.py b/tests/test_kv_cache_distributed.py index 22b5827..c90c054 100644 --- a/tests/test_kv_cache_distributed.py +++ b/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.