@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 184 lines 5.71 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """NOTE: This API server is used only for demonstrating usage of AsyncEngine
4 and simple performance benchmarks. It is not intended for production use.
5 For production use, we recommend using our OpenAI compatible server.
6 We are also not going to accept PRs modifying this file.
7 """
8
9 import asyncio
10 import json
11 import ssl
12 from argparse import Namespace
13 from collections.abc import AsyncGenerator
14 from typing import Any
15
16 from fastapi import FastAPI, Request
17 from fastapi.responses import JSONResponse, Response, StreamingResponse
18
19 import vllm.envs as envs
20 from vllm.engine.arg_utils import AsyncEngineArgs
21 from vllm.engine.async_llm_engine import AsyncLLMEngine
22 from vllm.entrypoints.launchers.launcher import serve_http
23 from vllm.entrypoints.serve.utils.api_utils import with_cancellation
24 from vllm.logger import init_logger
25 from vllm.sampling_params import SamplingParams
26 from vllm.usage.usage_lib import UsageContext
27 from vllm.utils import random_uuid
28 from vllm.utils.argparse_utils import FlexibleArgumentParser
29 from vllm.utils.system_utils import set_ulimit
30 from vllm.version import __version__ as VLLM_VERSION
31
32 logger = init_logger("api_server")
33
34 app = FastAPI()
35 engine = None
36
37
38 @app.get("/health")
39 async def health() -> Response:
40 """Health check."""
41 return Response(status_code=200)
42
43
44 @app.post("/generate")
45 async def generate(request: Request) -> Response:
46 """Generate completion for the request.
47
48 The request should be a JSON object with the following fields:
49 - prompt: the prompt to use for the generation.
50 - stream: whether to stream the results or not.
51 - other fields: the sampling parameters (See `SamplingParams` for details).
52 """
53 request_dict = await request.json()
54 return await _generate(request_dict, raw_request=request)
55
56
57 @with_cancellation
58 async def _generate(request_dict: dict, raw_request: Request) -> Response:
59 prompt = request_dict.pop("prompt")
60 stream = request_dict.pop("stream", False)
61 # Since SamplingParams is created fresh per request, safe to skip clone
62 sampling_params = SamplingParams(**request_dict, skip_clone=True)
63 request_id = random_uuid()
64
65 assert engine is not None
66 results_generator = engine.generate(prompt, sampling_params, request_id)
67
68 # Streaming case
69 async def stream_results() -> AsyncGenerator[bytes, None]:
70 async for request_output in results_generator:
71 prompt = request_output.prompt
72 assert prompt is not None
73 text_outputs = [prompt + output.text for output in request_output.outputs]
74 ret = {"text": text_outputs}
75 yield (json.dumps(ret) + "\n").encode("utf-8")
76
77 if stream:
78 return StreamingResponse(stream_results())
79
80 # Non-streaming case
81 final_output = None
82 try:
83 async for request_output in results_generator:
84 final_output = request_output
85 except asyncio.CancelledError:
86 return Response(status_code=499)
87
88 assert final_output is not None
89 prompt = final_output.prompt
90 assert prompt is not None
91 text_outputs = [prompt + output.text for output in final_output.outputs]
92 ret = {"text": text_outputs}
93 return JSONResponse(ret)
94
95
96 def build_app(args: Namespace) -> FastAPI:
97 global app
98
99 app.root_path = args.root_path
100 return app
101
102
103 async def init_app(
104 args: Namespace,
105 llm_engine: AsyncLLMEngine | None = None,
106 ) -> FastAPI:
107 app = build_app(args)
108
109 global engine
110
111 engine_args = AsyncEngineArgs.from_cli_args(args)
112 engine = (
113 llm_engine
114 if llm_engine is not None
115 else AsyncLLMEngine.from_engine_args(
116 engine_args, usage_context=UsageContext.API_SERVER
117 )
118 )
119 app.state.engine_client = engine
120 app.state.args = args
121 return app
122
123
124 async def run_server(
125 args: Namespace, llm_engine: AsyncLLMEngine | None = None, **uvicorn_kwargs: Any
126 ) -> None:
127 logger.info("vLLM API server version %s", VLLM_VERSION)
128 logger.info("args: %s", args)
129
130 set_ulimit()
131
132 app = await init_app(args, llm_engine)
133 assert engine is not None
134
135 shutdown_task = await serve_http(
136 app,
137 sock=None,
138 enable_ssl_refresh=args.enable_ssl_refresh,
139 host=args.host,
140 port=args.port,
141 log_level=args.log_level,
142 timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE,
143 ssl_keyfile=args.ssl_keyfile,
144 ssl_certfile=args.ssl_certfile,
145 ssl_ca_certs=args.ssl_ca_certs,
146 ssl_cert_reqs=args.ssl_cert_reqs,
147 **uvicorn_kwargs,
148 )
149
150 await shutdown_task
151
152
153 if __name__ == "__main__":
154 parser = FlexibleArgumentParser()
155 parser.add_argument("--host", type=str, default=None)
156 parser.add_argument("--port", type=parser.check_port, default=8000)
157 parser.add_argument("--ssl-keyfile", type=str, default=None)
158 parser.add_argument("--ssl-certfile", type=str, default=None)
159 parser.add_argument(
160 "--ssl-ca-certs", type=str, default=None, help="The CA certificates file"
161 )
162 parser.add_argument(
163 "--enable-ssl-refresh",
164 action="store_true",
165 default=False,
166 help="Refresh SSL Context when SSL certificate files change",
167 )
168 parser.add_argument(
169 "--ssl-cert-reqs",
170 type=int,
171 default=int(ssl.CERT_NONE),
172 help="Whether client certificate is required (see stdlib ssl module's)",
173 )
174 parser.add_argument(
175 "--root-path",
176 type=str,
177 default=None,
178 help="FastAPI root_path when app is behind a path based routing proxy",
179 )
180 parser.add_argument("--log-level", type=str, default="debug")
181 parser = AsyncEngineArgs.add_cli_args(parser)
182 args = parser.parse_args()
183
184 asyncio.run(run_server(args))