diff --git a/mstar/api_server/openai/_util.py b/mstar/api_server/openai/_util.py index 48845f6e2..7ae0f7f24 100644 --- a/mstar/api_server/openai/_util.py +++ b/mstar/api_server/openai/_util.py @@ -19,3 +19,9 @@ def rid(prefix: str) -> str: def sse(obj: dict) -> str: return f"data: {json.dumps(obj)}\n\n" + + +def error_type(status: int) -> str: + """The OpenAI error ``type`` for an HTTP status: a 4xx is the client's + (``invalid_request_error``), anything else is ours (``server_error``).""" + return "invalid_request_error" if 400 <= status < 500 else "server_error" diff --git a/mstar/api_server/openai/serving_chat.py b/mstar/api_server/openai/serving_chat.py index 9878a1286..8e007de3d 100644 --- a/mstar/api_server/openai/serving_chat.py +++ b/mstar/api_server/openai/serving_chat.py @@ -10,8 +10,10 @@ import base64 +from fastapi import HTTPException + from mstar.api_server import media_io -from mstar.api_server.openai._util import SSE_DONE, now, rid, sse +from mstar.api_server.openai._util import SSE_DONE, error_type, now, rid, sse async def create_chat_completion(api, model_name, adapter, req, raw_request=None): @@ -89,14 +91,36 @@ def chunk(delta, finish=None) -> str: "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], }) + def error(message: str, status: int) -> str: + return sse({"error": {"message": message, "type": error_type(status), "code": status}}) + yield chunk({"role": "assistant"}) - async for c in api.iter_result_chunks(request_id): - if c.modality == "text": - yield chunk({"content": c.data.decode("utf-8", "replace")}) - elif c.modality == "audio": - # Streaming audio deltas are base64 16-bit PCM at the model rate. - yield chunk({"audio": {"id": rid("audio"), "data": base64.b64encode(c.data).decode("ascii")}}) - elif c.modality == "image": - yield chunk({"content": media_io.png_to_data_url(c.data)}) - yield chunk({}, finish="stop") + failed = False + try: + async for c in api.iter_result_chunks(request_id): + if c.modality == "text": + yield chunk({"content": c.data.decode("utf-8", "replace")}) + elif c.modality == "audio": + # Streaming audio deltas are base64 16-bit PCM at the model rate. + yield chunk({"audio": {"id": rid("audio"), "data": base64.b64encode(c.data).decode("ascii")}}) + elif c.modality == "image": + yield chunk({"content": media_io.png_to_data_url(c.data)}) + elif c.modality == "error": + # The request failed after the stream opened (an engine error + # mid generation, a preprocess error); the HTTP status is + # committed, so the failure travels in-band the way the + # non-streaming path's error body does, not as a normal + # ``stop`` a client would take for a complete answer. The error + # is the iterator's last chunk, so the loop ends on its own; + # returning here instead would trip the iterator's abort on a + # request that is already gone. + failed = True + yield error(c.data.decode("utf-8", "replace"), int(c.metadata.get("status", 500))) + except HTTPException as exc: + # The delivery timeout raises out of the iterator (which aborts the + # request on its way out); report it the same way. + failed = True + yield error(str(exc.detail), exc.status_code) + if not failed: + yield chunk({}, finish="stop") yield SSE_DONE diff --git a/test/modular/test_openai_router.py b/test/modular/test_openai_router.py index 855747e41..3003c1d8d 100644 --- a/test/modular/test_openai_router.py +++ b/test/modular/test_openai_router.py @@ -44,6 +44,9 @@ def __init__(self, model_name="bagel"): self._chunks: dict = {} self.next_chunks: list = [] self.last_raw_request = None + # raised out of the stream after its chunks, like the delivery timeout + self.raise_after: Exception | None = None + self.aborted: list = [] def submit_request(self, **kw): self.last_submit = kw @@ -55,8 +58,18 @@ async def collect_results(self, request_id, raw_request=None): return self._chunks.get(request_id, []) async def iter_result_chunks(self, request_id): - for c in self._chunks.get(request_id, []): - yield c + # Same contract as the real one: a consumer that stops early, or an + # exception, aborts the request; a fully drained stream does not. + finished = False + try: + for c in self._chunks.get(request_id, []): + yield c + if self.raise_after is not None: + raise self.raise_after + finished = True + finally: + if not finished: + self.aborted.append(request_id) @pytest.fixture @@ -227,6 +240,71 @@ def test_chat_stream(client_and_stub): assert lines[-1]["choices"][0]["finish_reason"] == "stop" +def test_chat_stream_reports_a_failed_request_in_band(client_and_stub): + """A request that fails after the stream opened ends with an error event, + not a ``finish_reason: stop`` that reads as a complete answer.""" + client, stub = client_and_stub + stub.model_name = "bagel" + stub.next_chunks = [ + _Chunk("text", b"Paris"), + _Chunk("error", b"Error in worker: ValueError: q implies q_len_per_req=5", {"status": 500}), + ] + text = client.post( + "/v1/chat/completions", + json={"model": "bagel", "messages": [{"role": "user", "content": "go"}], "stream": True}, + ).text + events = [json.loads(l[6:]) for l in text.splitlines() if l.startswith("data: ") and "[DONE]" not in l] + assert events[1]["choices"][0]["delta"]["content"] == "Paris" + assert events[-1]["error"] == { + "message": "Error in worker: ValueError: q implies q_len_per_req=5", "type": "server_error", "code": 500, + } + assert not any(e.get("choices", [{}])[0].get("finish_reason") for e in events) + assert text.rstrip().endswith("data: [DONE]") + # The request is already gone by the time its error chunk arrives; the + # stream must not abort it on the way out. + assert stub.aborted == [] + + +def _stream_events(client, model="bagel"): + text = client.post( + "/v1/chat/completions", + json={"model": model, "messages": [{"role": "user", "content": "go"}], "stream": True}, + ).text + events = [json.loads(l[6:]) for l in text.splitlines() if l.startswith("data: ") and "[DONE]" not in l] + return text, events + + +def test_chat_stream_error_type_follows_the_status(client_and_stub): + """A 4xx from the data worker (a ValueError or TypeError in preprocess) + is the client's error, not a server_error.""" + client, stub = client_and_stub + stub.model_name = "bagel" + stub.next_chunks = [_Chunk("error", b"ValueError: unknown domain 'x'", {"status": 400})] + _, events = _stream_events(client) + assert events[-1]["error"] == { + "message": "ValueError: unknown domain 'x'", "type": "invalid_request_error", "code": 400, + } + assert stub.aborted == [] + + +def test_chat_stream_reports_a_timeout_in_band(client_and_stub): + """The delivery timeout raises an HTTPException out of the chunk iterator + mid-stream; the client gets an error event and [DONE], not a cut + connection, and the iterator's own abort still runs.""" + from fastapi import HTTPException + + client, stub = client_and_stub + stub.model_name = "bagel" + stub.next_chunks = [_Chunk("text", b"Par")] + stub.raise_after = HTTPException(status_code=500, detail="Request timed out") + text, events = _stream_events(client) + assert events[1]["choices"][0]["delta"]["content"] == "Par" + assert events[-1]["error"] == {"message": "Request timed out", "type": "server_error", "code": 500} + assert not any(e.get("choices", [{}])[0].get("finish_reason") for e in events) + assert text.rstrip().endswith("data: [DONE]") + assert stub.aborted == [stub.last_submit["request_id"]] + + def test_unsupported_model_404(client_and_stub): client, stub = client_and_stub stub.model_name = "pi05"