@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 112 lines 3.81 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.chat_completion.batch_serving import OpenAIServingChatBatch
11 from vllm.entrypoints.openai.chat_completion.protocol import (
12 BatchChatCompletionRequest,
13 ChatCompletionRequest,
14 ChatCompletionResponse,
15 )
16 from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
17 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
18 from vllm.entrypoints.serve.utils.api_utils import (
19 load_aware_call,
20 validate_json_request,
21 with_cancellation,
22 )
23 from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
24 from vllm.entrypoints.serve.utils.sse_keep_alive import with_sse_keep_alive
25 from vllm.logger import init_logger
26
27 logger = init_logger(__name__)
28
29 router = APIRouter()
30 ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format"
31
32
33 def chat(request: Request) -> OpenAIServingChat | None:
34 return request.app.state.openai_serving_chat
35
36
37 def batch_chat(request: Request) -> OpenAIServingChatBatch | None:
38 return request.app.state.openai_serving_chat_batch
39
40
41 @router.post(
42 "/v1/chat/completions",
43 dependencies=[Depends(validate_json_request)],
44 responses={
45 HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
46 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
47 HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
48 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
49 HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
50 },
51 )
52 @with_cancellation
53 @load_aware_call
54 async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
55 metrics_header_format = raw_request.headers.get(
56 ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, ""
57 )
58 handler = chat(raw_request)
59 if handler is None:
60 raise NotImplementedError("The model does not support Chat Completions API")
61
62 generator = await handler.create_chat_completion(request, raw_request)
63
64 if isinstance(generator, ErrorResponse):
65 return JSONResponse(
66 content=generator.model_dump(), status_code=generator.error.code
67 )
68
69 elif isinstance(generator, ChatCompletionResponse):
70 return JSONResponse(
71 content=generator.model_dump(),
72 headers=metrics_header(metrics_header_format),
73 )
74
75 args = getattr(raw_request.app.state, "args", None)
76 keep_alive_interval = getattr(args, "sse_keep_alive_interval", 0)
77 return StreamingResponse(
78 content=with_sse_keep_alive(generator, float(keep_alive_interval)),
79 media_type="text/event-stream",
80 )
81
82
83 @router.post(
84 "/v1/chat/completions/batch",
85 dependencies=[Depends(validate_json_request)],
86 responses={
87 HTTPStatus.OK.value: {},
88 HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
89 HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
90 HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
91 HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
92 },
93 )
94 @with_cancellation
95 @load_aware_call
96 async def create_batch_chat_completion(
97 request: BatchChatCompletionRequest, raw_request: Request
98 ):
99 handler = batch_chat(raw_request)
100 if handler is None:
101 raise NotImplementedError("The model does not support Chat Completions API")
102
103 result = await handler.create_batch_chat_completion(request, raw_request)
104
105 if isinstance(result, ErrorResponse):
106 return JSONResponse(content=result.model_dump(), status_code=result.error.code)
107
108 return JSONResponse(content=result.model_dump())
109
110
111 def attach_router(app: FastAPI):
112 app.include_router(router)