[03/06] Sampling Parameters | API Server Filtering

mstar/api_server/entrypoint.py

@@ -3,6 +3,7 @@
 import asyncio
 import base64
 import collections
+import contextlib
 import itertools
 import json
 import logging
@@ -23,6 +24,8 @@
 import msgpack
 import uvicorn
 from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile, WebSocket, WebSocketDisconnect
+from fastapi.encoders import jsonable_encoder
+from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 from starlette.concurrency import run_in_threadpool
@@ -62,6 +65,7 @@
 from mstar.utils.client_errors import client_error_text, scrub_paths
 from mstar.utils.exitcode import describe_exitcode
 from mstar.utils.gc_freeze import freeze_after_setup
+from mstar.utils.generation_kwargs import normalize_generation_kwargs
 from mstar.utils.logging_config import quiet_noisy_loggers
 from mstar.utils.orphan import watch_parent
 from mstar.utils.socket_prefix import per_port_socket_prefix
@@ -334,6 +338,18 @@ def _chunk_to_ndjson_payload(chunk: ResultChunk) -> str:
     }) + "\n"
 
 
+async def _disconnected(raw_request: Request) -> None:
+    """Return once the client has gone; the body is already read."""
+    while (await raw_request.receive())["type"] != "http.disconnect":
+        pass
+
+
+async def _empty_stream():
+    """The body for a client that left before its first chunk."""
+    return
+    yield
+
+
 def _chunk_to_binary_frame(chunk: ResultChunk) -> tuple[bytes, bytes]:
     """Serialize one result chunk as a header line plus its untouched payload.
 
@@ -589,6 +605,16 @@ def enable_metrics(self) -> None:
     # Submitting a request
     # ----------------------------------------------------------
 
+    def _report_ignored_params(self, model_kwargs: dict) -> None:
+        """Log the keys the model does not read; the request is served regardless."""
+        model = getattr(self, "model", None)
+        declared = model.request_kwargs() if model is not None else None
+        if declared is None:
+            return
+        ignored = sorted(set(model_kwargs) - declared - set(model.ENGINE_KWARGS) - set(model.LENGTH_KWARGS))
+        if ignored:
+            logger.warning("Ignoring model_kwargs this model does not read: %s", ignored)
+
     def submit_request(
         self,
         *,
@@ -644,6 +670,10 @@ def submit_request(
                 f"Output modality {names} requires streaming=True; raw frame "
                 "chunks cannot be returned as an aggregated response."
             )
+        # here, before anything is queued: null is unset, types and ranges are
+        # checked once for every endpoint, and an over-wide seed cannot be packed
+        model_kwargs = normalize_generation_kwargs(model_kwargs)
+        self._report_ignored_params(model_kwargs)
 
         if file_streams:
             try:
@@ -1287,7 +1317,7 @@ async def _iter_result_chunk_batches(
             if not finished:
                 self.abort_submission(key, outcome="timeout" if timed_out else "abort")
 
-    def async_stream_results(self, request_id: str, binary: bool = False):
+    def async_stream_results(self, request_id: str, binary: bool = False, chunks=None):
         """Yield the serialized body of ``/generate`` one piece at a time.
 
         ``binary`` selects the length-framed form negotiated through ``Accept``.
@@ -1298,12 +1328,58 @@ def async_stream_results(self, request_id: str, binary: bool = False):
         than an ``async def`` delegating to it: the branch is per-request, not
         per-chunk, and this keeps both bodies flat. It also binds the
         submission now, before Starlette starts iterating the body.
+
+        ``chunks`` is the request's ``ResultChunk`` iterator when the caller has
+        already read from it (``open_result_stream``); None opens a fresh one.
         """
-        chunks = self.iter_result_chunks(request_id)
+        chunks = chunks if chunks is not None else self.iter_result_chunks(request_id)
         if binary:
             return self._stream_binary(chunks)
         return self._stream_ndjson(chunks)
 
+    async def open_result_stream(
+        self, request_id: str, raw_request: Request | None = None, *, batches: bool = False,
+    ):
+        """The request's result iterator, after its first item has arrived.
+
+        A request the engine refuses raises ``HTTPException`` with its status
+        here, before the caller commits a 200; otherwise the iterator replays
+        the first item. ``batches`` selects ``iter_result_chunk_batches``. A
+        client that disconnects first is aborted and gets an empty iterator.
+        """
+        items = (self.iter_result_chunk_batches if batches else self.iter_result_chunks)(request_id)
+        first_read = asyncio.ensure_future(anext(items, None))
+        gone = asyncio.ensure_future(_disconnected(raw_request)) if raw_request is not None else None
+        try:
+            await asyncio.wait({first_read, gone} - {None}, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            # a pending receive() would take the disconnect StreamingResponse listens for
+            for task in (gone, first_read):
+                if task is not None and not task.done():
+                    # cancelling the read unwinds the iterator, which aborts the request
+                    task.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await task
+        if first_read.cancelled():
+            return _empty_stream()
+        first = first_read.result()
+        head = first[0] if batches and first else first
+        if head is not None and head.modality == "error":
+            # read to the end, so the finished request is not aborted
+            async for _ in items:
+                pass
+            raise HTTPException(
+                status_code=int((head.metadata or {}).get("status", 500)),
+                detail=bytes(head.data).decode("utf-8", "replace"),
+            )
+
+        async def replay():
+            if first is not None:
+                yield first
+            async for item in items:
+                yield item
+        return replay()
+
     async def _stream_ndjson(self, chunks):
         async for chunk in chunks:
             line = self._chunk_to_ndjson(chunk)
@@ -1675,6 +1751,20 @@ async def _one(request: WarmupRequest, request_id: str) -> None:
     metrics=lambda: getattr(api_server, "metrics", None),
     routes=lambda: (*app.router.routes, *openai_router.routes),
 )
+
+
+@app.exception_handler(RequestValidationError)
+async def request_validation_error(request: Request, exc: RequestValidationError):
+    """A body that fails its schema is a 400 like every other refused input,
+    in the envelope of the route family it was sent to."""
+    if request.url.path.startswith("/v1/"):
+        return JSONResponse(
+            status_code=400,
+            content={"error": {
+                "message": str(exc.errors()), "type": "invalid_request_error", "code": 400,
+            }},
+        )
+    return JSONResponse(status_code=400, content={"detail": jsonable_encoder(exc.errors())})
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -2023,7 +2113,9 @@ def release_turn() -> None:
                 text, _ = client_error_text(
                     exc, f"generate/ws request {request_id} failed", log=logger, type_name=False,
                 )
-                payload = {"request_id": request_id, "error": text}
+                # the status /generate would have answered with
+                status = 400 if isinstance(exc, (ValueError, TypeError)) else 500
+                payload = {"request_id": request_id, "error": text, "status": status}
             try:
                 await send(payload, binary)
             except Exception:  # noqa: BLE001
@@ -2136,7 +2228,7 @@ def on_late(error: str) -> None:
             try:
                 check_request_id(request_id)
             except ValueError as exc:
-                reply({"request_id": request_id, "error": str(exc)}, binary)
+                reply({"request_id": request_id, "error": str(exc), "status": 400}, binary)
                 continue
             if _ws_is_append(message):
                 if request_id not in live:
@@ -2306,24 +2398,31 @@ async def generate(
 
     try:
         await api_server.wait_for_request_id(request_id)
-        request_id = api_server.submit_request(
-            text=text,
-            file_paths=file_paths or None,
-            input_modalities=in_mods,
-            output_modalities=out_mods,
-            model_kwargs=parsed_kwargs,
-            prompt_parts=parts or None,
-            streaming=streaming,
-            request_id=request_id,
-        )
+        try:
+            request_id = api_server.submit_request(
+                text=text,
+                file_paths=file_paths or None,
+                input_modalities=in_mods,
+                output_modalities=out_mods,
+                model_kwargs=parsed_kwargs,
+                prompt_parts=parts or None,
+                streaming=streaming,
+                request_id=request_id,
+            )
+        except json.JSONDecodeError:
+            raise   # a downstream parse failure, not a bad request (a 500 below)
+        except ValueError as e:   # submit_request's refusals are the client's fault
+            raise HTTPException(status_code=400, detail=str(e)) from e
 
         if streaming:
             # Substring match, not RFC 7231 q-value parsing: the value is a
             # private vendor type that appears in no other media range, and a
             # client that does not ask for it keeps the historical NDJSON body.
             binary = BINARY_STREAM_MEDIA_TYPE in request.headers.get("accept", "")
+            # wait for the first chunk: a refused request gets its real status, not a 200
+            chunks = await api_server.open_result_stream(request_id, request)
             return StreamingResponse(
-                api_server.async_stream_results(request_id, binary=binary),
+                api_server.async_stream_results(request_id, binary=binary, chunks=chunks),
                 media_type=BINARY_STREAM_MEDIA_TYPE if binary else NDJSON_STREAM_MEDIA_TYPE,
                 headers={"Cache-Control": "no-cache", "Vary": "Accept"},
             )