main
py 76 lines 2.54 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 from http import HTTPStatus
6
7 from fastapi import APIRouter, Depends, FastAPI, Request
8 from fastapi.responses import JSONResponse, StreamingResponse
9
10 from vllm.entrypoints.openai.completion.protocol import (
11 CompletionRequest,
12 CompletionResponse,
13 )
14 from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
15 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
16 from vllm.entrypoints.serve.utils.api_utils import (
17 load_aware_call,
18 validate_json_request,
19 with_cancellation,
20 )
21 from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
22 from vllm.entrypoints.serve.utils.sse_keep_alive import with_sse_keep_alive
23 from vllm.logger import init_logger
24
25 logger = init_logger(__name__)
26
27 router = APIRouter()
28 ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format"
29
30
31 def completion(request: Request) -> OpenAIServingCompletion | None:
32 return request.app.state.openai_serving_completion
33
34
35 @router.post(
36 "/v1/completions",
37 dependencies=[Depends(validate_json_request)],
38 responses={
39 HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
40 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
41 HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
42 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
43 },
44 )
45 @with_cancellation
46 @load_aware_call
47 async def create_completion(request: CompletionRequest, raw_request: Request):
48 metrics_header_format = raw_request.headers.get(
49 ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, ""
50 )
51 handler = completion(raw_request)
52 if handler is None:
53 raise NotImplementedError("The model does not support Completions API")
54
55 generator = await handler.create_completion(request, raw_request)
56
57 if isinstance(generator, ErrorResponse):
58 return JSONResponse(
59 content=generator.model_dump(), status_code=generator.error.code
60 )
61 elif isinstance(generator, CompletionResponse):
62 return JSONResponse(
63 content=generator.model_dump(),
64 headers=metrics_header(metrics_header_format),
65 )
66
67 args = getattr(raw_request.app.state, "args", None)
68 keep_alive_interval = getattr(args, "sse_keep_alive_interval", 0)
69 return StreamingResponse(
70 content=with_sse_keep_alive(generator, float(keep_alive_interval)),
71 media_type="text/event-stream",
72 )
73
74
75 def attach_router(app: FastAPI):
76 app.include_router(router)