diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index a405550c4c..ae48fc1a9c 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +from argparse import Namespace from typing import Final import pytest import schemathesis +from fastapi.testclient import TestClient from hypothesis import HealthCheck, settings +from jsonschema import Draft202012Validator from schemathesis import GenerationMode from schemathesis.checks import not_a_server_error from schemathesis.config import ( @@ -30,6 +33,76 @@ DEFAULT_TIMEOUT_SECONDS: Final[int] = 10 * _ROCM_TIMEOUT_MULTIPLIER LONG_TIMEOUT_SECONDS: Final[int] = 60 * _ROCM_TIMEOUT_MULTIPLIER +def test_openapi_exposes_completion_response_schemas() -> None: + """The published 200 contract of both completion APIs must reference the + real response models, for the JSON body and for each SSE data chunk. + + Uses the production app: FastAPI renames models used as both input and + output (-Input/-Output), which a router-only app cannot show. + """ + from vllm.entrypoints.launchers.app import build_app + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionResponse, + ChatCompletionStreamResponse, + ) + from vllm.entrypoints.openai.completion.protocol import ( + CompletionResponse, + CompletionStreamResponse, + ) + + args = Namespace( + disable_fastapi_docs=False, + enable_offline_docs=False, + root_path="", + allowed_origins=["*"], + allow_credentials=False, + allowed_methods=["*"], + allowed_headers=["*"], + api_key=None, + enable_request_id_headers=False, + enable_fault_tolerance=False, + middleware=[], + log_error_stack=False, + ) + # No lifespan: schema generation does not need an inference engine. + client = TestClient(build_app(args, supported_tasks=("generate",))) + document = client.get("/openapi.json").json() + message = {"role": "assistant", "content": "hi"} + for path, unary, unary_choice, chunk, chunk_choice in ( + ( + "/v1/chat/completions", + ChatCompletionResponse, + {"index": 0, "message": message}, + ChatCompletionStreamResponse, + {"index": 0, "delta": {"content": "hi"}}, + ), + ( + "/v1/completions", + CompletionResponse, + {"index": 0, "text": "hi"}, + CompletionStreamResponse, + {"index": 0, "text": "hi"}, + ), + ): + content = document["paths"][path]["post"]["responses"]["200"]["content"] + data = content["text/event-stream"]["schema"]["properties"]["data"] + assert data["anyOf"][0] == {"const": "[DONE]"} + for schema, model, choice in ( + (content["application/json"]["schema"], unary, unary_choice), + (data["anyOf"][1]["contentSchema"], chunk, chunk_choice), + ): + assert schema == {"$ref": f"#/components/schemas/{model.__name__}"} + validator = Draft202012Validator( + {**schema, "components": document["components"]} + ) + payload = model(model="m", choices=[choice], usage={}).model_dump( + mode="json" + ) + validator.validate(payload) + payload["choices"] = "not an array" + assert not validator.is_valid(payload) + + @pytest.fixture(scope="module") def server(): args = [ diff --git a/tests/entrypoints/serve/utils/test_api_utils.py b/tests/entrypoints/serve/utils/test_api_utils.py index a41f39187f..4dcf2dcffe 100644 --- a/tests/entrypoints/serve/utils/test_api_utils.py +++ b/tests/entrypoints/serve/utils/test_api_utils.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio from argparse import Namespace import pytest +from fastapi import FastAPI, Request from vllm.entrypoints.generate.base.protocol import StreamOptions from vllm.entrypoints.serve.utils import api_utils @@ -14,6 +16,55 @@ from vllm.entrypoints.serve.utils.api_utils import ( ) +@pytest.mark.asyncio +async def test_with_cancellation_disconnect_skips_response_validation(): + """A disconnect must cancel the handler without a response_model 500.""" + app = FastAPI() + started = asyncio.Event() + cancelled = asyncio.Event() + + @app.post("/", response_model=str) + @api_utils.with_cancellation + async def handler(body: dict, raw_request: Request): + started.set() + try: + await asyncio.Future() + finally: + cancelled.set() + + messages = [] + body_sent = False + + async def receive(): + nonlocal body_sent + if not body_sent: + body_sent = True + return {"type": "http.request", "body": b"{}", "more_body": False} + await started.wait() + return {"type": "http.disconnect"} + + async def send(message): + messages.append(message) + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/", + "raw_path": b"/", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + "client": ("127.0.0.1", 1234), + "server": ("test", 80), + } + await asyncio.wait_for(app(scope, receive, send), timeout=5) + assert messages[0]["status"] == 200 + assert messages[1]["body"] == b"" + await asyncio.wait_for(cancelled.wait(), timeout=5) + + @pytest.mark.parametrize( ("stream_options", "expected"), [ diff --git a/vllm/entrypoints/openai/chat_completion/api_router.py b/vllm/entrypoints/openai/chat_completion/api_router.py index aecf5bba44..0997e8d634 100644 --- a/vllm/entrypoints/openai/chat_completion/api_router.py +++ b/vllm/entrypoints/openai/chat_completion/api_router.py @@ -12,10 +12,12 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( BatchChatCompletionRequest, ChatCompletionRequest, ChatCompletionResponse, + ChatCompletionStreamResponse, ) from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.api_utils import ( + json_or_sse_responses, load_aware_call, validate_json_request, with_cancellation, @@ -41,13 +43,16 @@ def batch_chat(request: Request) -> OpenAIServingChatBatch | None: @router.post( "/v1/chat/completions", dependencies=[Depends(validate_json_request)], - responses={ - HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, - HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, - HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse}, - }, + **json_or_sse_responses( + ChatCompletionResponse, + ChatCompletionStreamResponse, + { + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse}, + }, + ), ) @with_cancellation @load_aware_call diff --git a/vllm/entrypoints/openai/completion/api_router.py b/vllm/entrypoints/openai/completion/api_router.py index fe909c7bae..6b100f1c78 100644 --- a/vllm/entrypoints/openai/completion/api_router.py +++ b/vllm/entrypoints/openai/completion/api_router.py @@ -10,10 +10,12 @@ from fastapi.responses import JSONResponse, StreamingResponse from vllm.entrypoints.openai.completion.protocol import ( CompletionRequest, CompletionResponse, + CompletionStreamResponse, ) from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.api_utils import ( + json_or_sse_responses, load_aware_call, validate_json_request, with_cancellation, @@ -35,12 +37,15 @@ def completion(request: Request) -> OpenAIServingCompletion | None: @router.post( "/v1/completions", dependencies=[Depends(validate_json_request)], - responses={ - HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, - HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, - }, + **json_or_sse_responses( + CompletionResponse, + CompletionStreamResponse, + { + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + }, + ), ) @with_cancellation @load_aware_call diff --git a/vllm/entrypoints/serve/utils/api_utils.py b/vllm/entrypoints/serve/utils/api_utils.py index 6b82a2afa0..14de8b558a 100644 --- a/vllm/entrypoints/serve/utils/api_utils.py +++ b/vllm/entrypoints/serve/utils/api_utils.py @@ -6,13 +6,15 @@ import dataclasses import functools import os from argparse import Namespace +from http import HTTPStatus from logging import Logger from string import Template from typing import Any from fastapi import Request from fastapi.exceptions import RequestValidationError -from fastapi.responses import JSONResponse, StreamingResponse +from fastapi.responses import JSONResponse, Response, StreamingResponse +from pydantic import BaseModel from starlette.background import BackgroundTask, BackgroundTasks from vllm import envs @@ -88,11 +90,80 @@ def with_cancellation(handler_func): if handler_task in done: return handler_task.result() - return None + # The client is gone; an empty response skips response_model validation. + return Response() return wrapper +_ERROR_REF = {"$ref": "#/components/schemas/ErrorResponse"} + + +def _sse_event_schema(chunk_ref: dict[str, str]) -> dict[str, Any]: + # One SSE event: `data` is a JSON chunk, a JSON error, or "[DONE]". This is + # the envelope FastAPI emits for EventSourceResponse routes, but published + # under `schema` because OpenAPI 3.1 has no `itemSchema`. + return { + "type": "object", + "required": ["data"], + "properties": { + "data": { + "type": "string", + "anyOf": [ + {"const": "[DONE]"}, + { + "contentMediaType": "application/json", + "contentSchema": chunk_ref, + }, + { + "contentMediaType": "application/json", + "contentSchema": _ERROR_REF, + }, + ], + }, + "event": {"type": "string"}, + "id": {"type": "string"}, + "retry": {"type": "integer", "minimum": 0}, + }, + } + + +def json_or_sse_responses( + unary_model: type[BaseModel], + chunk_model: type[BaseModel], + error_responses: dict[int, dict[str, Any]], +) -> dict[str, Any]: + """Route kwargs for an endpoint answering with ``unary_model`` as JSON, or + with SSE ``data`` events carrying ``chunk_model`` when ``stream=true``. + + FastAPI only registers models reachable from ``response_model`` or an + additional response's ``model``. The chunk model rides on the 200 entry, + which also overwrites the JSON schema, so ``openapi_extra`` restores it. + """ + unary_ref = {"$ref": f"#/components/schemas/{unary_model.__name__}"} + chunk_ref = {"$ref": f"#/components/schemas/{chunk_model.__name__}"} + return { + "response_model": unary_model, + "responses": { + HTTPStatus.OK.value: { + "model": chunk_model, + "description": ( + "JSON body when stream=false; SSE events when stream=true." + ), + "content": { + "text/event-stream": {"schema": _sse_event_schema(chunk_ref)} + }, + }, + **error_responses, + }, + "openapi_extra": { + "responses": { + "200": {"content": {"application/json": {"schema": unary_ref}}} + } + }, + } + + def decrement_server_load(request: Request): request.app.state.server_load_metrics -= 1