From 1d11099c9065a14cc03fc19a1b62b83e29888717 Mon Sep 17 00:00:00 2001 From: gnezim Date: Mon, 19 Oct 2026 19:14:19 +0300 Subject: [PATCH] fix: preserve edge streaming responses on upstream failure --- README.md | 2 + src/thunder_forge/cluster/edge.py | 34 ++++-- tests/test_edge_streaming.py | 182 ++++++++++++++++++++++++++++++ 3 files changed, 210 insertions(+), 8 deletions(-) create mode 100644 tests/test_edge_streaming.py diff --git a/README.md b/README.md index cbd1834..94c6918 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,8 @@ When `services.olla.version` is omitted inside an explicit `services.olla` block After changing model placement or node topology in `tfconfig.yaml`, run `make restart gateway-cache-01` (or full `make restart`) before `make smoke ` so Olla and TF edge reload the generated router config. +TF edge defaults to a 180-second upstream I/O timeout, leaving room for the generated Olla config's 120-second backend timeouts and connection attempts. This limits waits for network activity, not total generation time. If an upstream stream fails after HTTP headers have been sent, edge closes the incomplete response and records a 502 failure in its access log; it cannot change the status already sent to the client. Streaming bodies are captured only when `services.edge.inspector.enabled` is true. Inspector response capture is limited to the first 1 MiB, with `body_truncated: true` in the JSONL entry when clipped; clients still receive the full stream. + Bootstrap verifies minimal service readiness only: Olla `/internal/health`, TF edge auth boundary health, and direct oMLX `/health`. It does not require inference models or chat to be ready; use `make smoke` for model visibility, routing, and chat checks after the cluster has warmed up. Run system-daemon install targets from a real terminal, not the VS Code guarded terminal, because macOS sudo/su password prompts can be blocked by the editor guard: diff --git a/src/thunder_forge/cluster/edge.py b/src/thunder_forge/cluster/edge.py index 42ad9cd..8d1010a 100644 --- a/src/thunder_forge/cluster/edge.py +++ b/src/thunder_forge/cluster/edge.py @@ -61,6 +61,7 @@ EDGE_DEFAULT_PORT = DEFAULT_EDGE_PORT EDGE_LAUNCHD_LABEL_PREFIX = "com.thunder-forge.edge" EDGE_DEFAULT_MAX_BODY_BYTES = 16 * 1024 * 1024 +EDGE_MAX_STREAM_CAPTURE_BYTES = 1024 * 1024 _ENV_LINE_RE = re.compile( r"^(?P\s*)(?P[A-Za-z_][A-Za-z0-9_]*)(?P\s*=\s*)(?P.*?)(?P\s*)$" @@ -106,7 +107,8 @@ class EdgeProxyConfig: active_requests: dict[str, tuple[str, str]] = field(default_factory=dict) active_requests_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) model_catalog: list[EdgeModelCatalogEntry] = field(default_factory=list) - timeout: float = 60.0 + # Leave room for Olla's 120-second backend timeouts and connection attempts. + timeout: float = 180.0 status_timeout: float = 8.0 max_body_bytes: int = EDGE_DEFAULT_MAX_BODY_BYTES inspector: EdgeInspectorConfig | None = None @@ -1166,6 +1168,7 @@ def _write_edge_inspector_entry( status_code: int | None = None, headers: dict[str, str] | httpx.Headers | None = None, body: bytes | None = None, + body_truncated: bool = False, error: str | None = None, ) -> None: inspector = config.inspector @@ -1188,6 +1191,8 @@ def _write_edge_inspector_entry( payload["headers"] = _normalize_headers_for_dump(headers) if body is not None: payload["body"] = _body_dump_payload(body) + if body_truncated: + payload["body_truncated"] = True if error: payload["error"] = error @@ -1574,9 +1579,10 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None: olla_endpoint=olla_endpoint, ) logged = True - _unregister_active_request(config, upstream.request_id) normalized_base_url = config.olla_base_url.rstrip("/") + headers_sent = False + olla_endpoint = "" try: with httpx.Client(base_url=normalized_base_url, timeout=config.timeout, trust_env=False) as client: with client.stream( @@ -1585,7 +1591,8 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None: headers=upstream.forwarded_headers, content=body, ) as response: - captured_body = bytearray() + captured_body = bytearray() if inspector is not None and inspector.enabled else None + body_truncated = False olla_endpoint = response.headers.get("X-Olla-Endpoint", "") self.send_response(response.status_code) response_headers = _proxy_response_headers(response.headers) @@ -1594,6 +1601,7 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None: self.send_header(name, value) try: self.end_headers() + headers_sent = True except OSError as exc: if self._is_client_disconnect(exc): record_stream(response.status_code, olla_endpoint=olla_endpoint) @@ -1601,14 +1609,17 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None: raise for chunk in response.iter_bytes(): if chunk: - captured_body.extend(chunk) + if captured_body is not None: + remaining = EDGE_MAX_STREAM_CAPTURE_BYTES - len(captured_body) + captured_body.extend(chunk[:remaining]) + body_truncated |= len(chunk) > remaining try: self.wfile.write(chunk) self.wfile.flush() except OSError as exc: if self._is_client_disconnect(exc): record_stream(response.status_code, olla_endpoint=olla_endpoint) - if inspector is not None and inspector.enabled: + if captured_body is not None: _write_edge_inspector_entry( config, event_type="response", @@ -1618,10 +1629,11 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None: status_code=response.status_code, headers=response.headers, body=bytes(captured_body), + body_truncated=body_truncated, ) return raise - if inspector is not None and inspector.enabled: + if captured_body is not None: _write_edge_inspector_entry( config, event_type="response", @@ -1631,6 +1643,7 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None: status_code=response.status_code, headers=response.headers, body=bytes(captured_body), + body_truncated=body_truncated, ) record_stream(response.status_code, olla_endpoint=olla_endpoint) except httpx.HTTPError as exc: @@ -1643,8 +1656,13 @@ def record_stream(status_code: int, *, olla_endpoint: str = "") -> None: path=self.path, error=str(exc), ) - record_stream(502) - self._send_edge_response(_json_response(502, {"error": f"upstream_failed: {exc}"})) + record_stream(502, olla_endpoint=olla_endpoint) + if headers_sent: + self.close_connection = True + else: + self._send_edge_response(_json_response(502, {"error": f"upstream_failed: {exc}"})) + finally: + _unregister_active_request(config, upstream.request_id) def _send_edge_response(self, result: EdgeProxyResponse) -> None: self.send_response(result.status_code) diff --git a/tests/test_edge_streaming.py b/tests/test_edge_streaming.py new file mode 100644 index 0000000..41de36b --- /dev/null +++ b/tests/test_edge_streaming.py @@ -0,0 +1,182 @@ +"""Exercise the actual streaming HTTP handler over a local socket pair.""" + +import json +import socket +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from types import SimpleNamespace + +import httpx +import pytest + +from thunder_forge.cluster import edge + + +@pytest.fixture +def exchange(monkeypatch): + client_type = httpx.Client + + def run(config, upstream, *, timeout_scale=1): + handlers = [] + + def server(_address, handler): + handlers.append(handler) + return SimpleNamespace(serve_forever=lambda: None) + + monkeypatch.setattr(edge, "ThreadingHTTPServer", server) + def proxy_client(**kwargs): + kwargs["timeout"] *= timeout_scale + return client_type(**kwargs, transport=httpx.MockTransport(upstream) if upstream else None) + + monkeypatch.setattr(edge.httpx, "Client", proxy_client) + edge.serve_edge_proxy(host="127.0.0.1", port=0, config=config) + body = json.dumps({"model": "agent", "messages": [], "stream": True}).encode() + request = ( + b"POST /v1/chat/completions HTTP/1.0\r\n" + b"Authorization: Bearer fixture-key\r\n" + + f"Content-Length: {len(body)}\r\n\r\n".encode() + body + ) + errors = [] + client, connection = socket.socketpair() + + def handle(): + try: + handlers[0](connection, ("127.0.0.1", 1234), None) + except Exception as exc: + errors.append(exc) + finally: + connection.close() + + with client: + client.settimeout(5) + client.sendall(request) + client.shutdown(socket.SHUT_WR) + thread = threading.Thread(target=handle, daemon=True) + thread.start() + try: + with client.makefile("rb") as response: + wire = response.read() + finally: + thread.join(timeout=5) + assert not thread.is_alive() + assert not errors + return wire + + return run + + +def test_edge_allows_a_response_within_olla_timeout(exchange): + # Scale 90 seconds of header latency to 450 ms: Olla allows 120 seconds. + class Upstream(BaseHTTPRequestHandler): + def do_POST(self): + self.rfile.read(int(self.headers["Content-Length"])) + time.sleep(0.45) + try: + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + self.wfile.write(b"data: [DONE]\n\n") + except (BrokenPipeError, ConnectionResetError): + pass + + def log_message(self, *_args): + pass + + with ThreadingHTTPServer(("127.0.0.1", 0), Upstream) as upstream: + thread = threading.Thread(target=upstream.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + config = edge.EdgeProxyConfig( + olla_base_url=f"http://127.0.0.1:{upstream.server_port}", + clients_by_key={"fixture-key": edge.EdgeClient("fixture")}, + ) + try: + wire = exchange(config, None, timeout_scale=0.005) + finally: + upstream.shutdown() + thread.join(timeout=5) + assert wire.startswith(b"HTTP/1.0 200") + assert wire.endswith(b"data: [DONE]\n\n") + + +@pytest.mark.parametrize("failure", ["connect", "before_body", "after_chunk", None]) +def test_stream_failure_sends_only_one_http_response(exchange, failure): + chunk = b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + closed = [] + records = [] + + class Stream(httpx.SyncByteStream): + def __iter__(self): + if failure == "before_body": + raise httpx.ReadError("fixture read failure") + yield chunk + if failure == "after_chunk": + raise httpx.ReadError("fixture read failure") + yield b"data: [DONE]\n\n" + + def close(self): + closed.append(True) + + def upstream(request): + if failure == "connect": + raise httpx.ConnectError("fixture connection failure", request=request) + return httpx.Response(200, headers={ + "Content-Type": "text/event-stream", "X-Olla-Endpoint": "infer-01", + }, stream=Stream()) + + config = edge.EdgeProxyConfig( + olla_base_url="http://olla.test", clients_by_key={"fixture-key": edge.EdgeClient("fixture")}, + access_log_sink=records.append, + ) + wire = exchange(config, upstream) + assert wire.count(b"HTTP/1.0 ") == 1 + assert wire.startswith(b"HTTP/1.0 502" if failure == "connect" else b"HTTP/1.0 200") + body = wire.split(b"\r\n\r\n", 1)[1] + if failure == "connect": + assert json.loads(body)["error"].startswith("upstream_failed:") + else: + assert closed == [True] + assert body == (b"" if failure == "before_body" else chunk if failure else chunk + b"data: [DONE]\n\n") + assert config.active_requests == {} + assert len(records) == 1 + record = json.loads(records[0]) + assert record["status_code"] == (502 if failure else 200) + if failure != "connect": + assert record["olla_endpoint"] == "infer-01" + + +@pytest.mark.parametrize("inspect", [False, True]) +def test_stream_capture_is_optional_and_bounded(exchange, monkeypatch, tmp_path, inspect): + chunk = b"data: " + b"x" * 65536 + b"\n\n" + chunks = 20 + retained = [] + + class Capture(bytearray): + def extend(self, data): + super().extend(data) + retained.append(len(self)) + + class Stream(httpx.SyncByteStream): + def __iter__(self): + for _ in range(chunks): + yield chunk + + monkeypatch.setattr(edge, "bytearray", Capture, raising=False) + config = edge.EdgeProxyConfig( + olla_base_url="http://olla.test", clients_by_key={"fixture-key": edge.EdgeClient("fixture")}, + inspector=edge.EdgeInspectorConfig(enabled=inspect, output_dir=tmp_path), + ) + wire = exchange(config, lambda _request: httpx.Response(200, stream=Stream())) + assert wire.split(b"\r\n\r\n", 1)[1] == chunk * chunks + assert max(retained, default=0) <= (1024 * 1024 if inspect else 0) + if inspect: + entries = [ + json.loads(line) + for path in tmp_path.glob("*/requests.jsonl") + for line in path.read_text().splitlines() + ] + response = next(entry for entry in entries if entry["type"] == "response") + assert len(response["body"].encode()) == 1024 * 1024 + assert response["body_truncated"] is True + else: + assert not list(tmp_path.glob("*/requests.jsonl"))