@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 229 lines 7.74 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import asyncio
4 import multiprocessing
5 import multiprocessing.forkserver as forkserver
6 import os
7 import signal
8 import socket
9 import tempfile
10 from argparse import Namespace
11 from collections.abc import AsyncIterator
12 from contextlib import asynccontextmanager
13 from typing import Any
14
15 import vllm.envs as envs
16 from vllm.engine.arg_utils import AsyncEngineArgs
17 from vllm.engine.protocol import EngineClient
18 from vllm.logger import configure_logging_from_args, init_logger
19 from vllm.reasoning import ReasoningParserManager
20 from vllm.tool_parsers import ToolParserManager
21 from vllm.usage.usage_lib import UsageContext
22 from vllm.utils.system_utils import decorate_logs
23
24 from ..app import build_app
25 from ..launcher import serve_http, setup_server
26 from ..utils.server_utils import get_uvicorn_log_config
27 from .app_state import init_app_state
28
29 prometheus_multiproc_dir: tempfile.TemporaryDirectory
30
31 # Cannot use __name__ (https://github.com/vllm-project/vllm/pull/4765)
32 logger = init_logger("vllm.entrypoints.launchers.api_server.entry")
33
34
35 @asynccontextmanager
36 async def build_async_engine_client(
37 args: Namespace,
38 *,
39 usage_context: UsageContext = UsageContext.OPENAI_API_SERVER,
40 client_config: dict[str, Any] | None = None,
41 ) -> AsyncIterator[EngineClient]:
42 if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
43 # The executor is expected to be mp.
44 # Pre-import heavy modules in the forkserver process
45 logger.debug("Setup forkserver with pre-imports")
46 multiprocessing.set_start_method("forkserver")
47 multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"])
48 forkserver.ensure_running()
49 logger.debug("Forkserver setup complete!")
50
51 # Context manager to handle engine_client lifecycle
52 # Ensures everything is shutdown and cleaned up on error/exit
53 engine_args = AsyncEngineArgs.from_cli_args(args)
54 from ..cli_args import propagate_flash_late_interaction
55
56 propagate_flash_late_interaction(args, engine_args)
57 if client_config:
58 engine_args._api_process_count = client_config.get("client_count", 1)
59 engine_args._api_process_rank = client_config.get("client_index", 0)
60
61 async with build_async_engine_client_from_engine_args(
62 engine_args,
63 usage_context=usage_context,
64 client_config=client_config,
65 ) as engine:
66 yield engine
67
68
69 @asynccontextmanager
70 async def build_async_engine_client_from_engine_args(
71 engine_args: AsyncEngineArgs,
72 *,
73 usage_context: UsageContext = UsageContext.OPENAI_API_SERVER,
74 client_config: dict[str, Any] | None = None,
75 ) -> AsyncIterator[EngineClient]:
76 """Create EngineClient, either:
77 - in-process using the AsyncLLMEngine Directly
78 - multiprocess using AsyncLLMEngine RPC
79
80 Returns the Client or None if the creation failed.
81 """
82 # Create the EngineConfig (determines if we can use V1).
83 vllm_config = engine_args.create_engine_config(usage_context=usage_context)
84
85 from vllm.v1.engine.async_llm import AsyncLLM
86
87 async_llm: AsyncLLM | None = None
88
89 # Don't mutate the input client_config
90 client_config = dict(client_config) if client_config else {}
91 client_count = client_config.pop("client_count", 1)
92 client_index = client_config.pop("client_index", 0)
93
94 try:
95 async_llm = AsyncLLM.from_vllm_config(
96 vllm_config=vllm_config,
97 usage_context=usage_context,
98 enable_log_requests=engine_args.enable_log_requests,
99 aggregate_engine_logging=engine_args.aggregate_engine_logging,
100 disable_log_stats=engine_args.disable_log_stats,
101 client_addresses=client_config,
102 client_count=client_count,
103 client_index=client_index,
104 )
105
106 # Don't keep the dummy data in memory
107 assert async_llm is not None
108 await async_llm.reset_mm_cache()
109
110 yield async_llm
111 finally:
112 if async_llm:
113 async_llm.shutdown(timeout=vllm_config.shutdown_timeout)
114
115
116 async def build_and_serve(
117 engine_client: EngineClient,
118 listen_address: str,
119 sock: socket.socket,
120 args: Namespace,
121 **uvicorn_kwargs,
122 ) -> asyncio.Task:
123 """Build FastAPI app, initialize state, and start serving.
124
125 Returns the shutdown task for the caller to await.
126 """
127 # Get uvicorn log config (from file or with endpoint filter)
128 log_config = get_uvicorn_log_config(args)
129 if log_config is not None:
130 uvicorn_kwargs["log_config"] = log_config
131
132 supported_tasks = await engine_client.get_supported_tasks()
133 model_config = engine_client.model_config
134
135 logger.info("Supported tasks: %s", supported_tasks)
136 app = build_app(args, supported_tasks, model_config)
137 await init_app_state(engine_client, app.state, args, supported_tasks)
138
139 logger.info("Starting vLLM server on %s", listen_address)
140
141 return await serve_http(
142 app,
143 sock=sock,
144 enable_ssl_refresh=args.enable_ssl_refresh,
145 host=args.host,
146 port=args.port,
147 log_level=args.uvicorn_log_level,
148 # NOTE: When the 'disable_uvicorn_access_log' value is True,
149 # no access log will be output.
150 access_log=not args.disable_uvicorn_access_log,
151 timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE,
152 ssl_keyfile=args.ssl_keyfile,
153 ssl_certfile=args.ssl_certfile,
154 ssl_ca_certs=args.ssl_ca_certs,
155 ssl_cert_reqs=args.ssl_cert_reqs,
156 ssl_ciphers=args.ssl_ciphers,
157 h11_max_incomplete_event_size=args.h11_max_incomplete_event_size,
158 h11_max_header_count=args.h11_max_header_count,
159 **uvicorn_kwargs,
160 )
161
162
163 async def run_server(args, **uvicorn_kwargs) -> None:
164 """Run a single-worker API server."""
165 decorate_logs("APIServer", skip_if_decorated=True)
166
167 # Interrupt initialization if SIGTERM arrives before uvicorn installs its
168 # own signal handlers. Once uvicorn is running it replaces this.
169 def _interrupt_init(*_) -> None:
170 raise KeyboardInterrupt("terminated")
171
172 signal.signal(signal.SIGTERM, _interrupt_init)
173
174 listen_address, sock = setup_server(args, reuse_port=False)
175 await run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
176
177
178 async def run_server_worker(
179 listen_address, sock, args, client_config=None, **uvicorn_kwargs
180 ) -> None:
181 """Run a single API server worker."""
182 if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
183 ToolParserManager.import_tool_parser(args.tool_parser_plugin)
184
185 if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3:
186 ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin)
187
188 async with build_async_engine_client(
189 args,
190 client_config=client_config,
191 ) as engine_client:
192 shutdown_task = await build_and_serve(
193 engine_client, listen_address, sock, args, **uvicorn_kwargs
194 )
195 # NB: Await server shutdown only after the backend context is exited
196 try:
197 await shutdown_task
198 finally:
199 sock.close()
200
201
202 def main():
203 import uvloop
204
205 from vllm.entrypoints.serve.utils.api_utils import cli_env_setup
206 from vllm.utils.argparse_utils import FlexibleArgumentParser
207
208 from ..cli_args import (
209 make_arg_parser,
210 validate_parsed_serve_args,
211 )
212
213 # NOTE(simon):
214 # This section should be in sync with vllm/entrypoints/cli/main.py for CLI
215 # entrypoints.
216 cli_env_setup()
217 parser = FlexibleArgumentParser(
218 description="vLLM OpenAI-Compatible RESTful API server."
219 )
220 parser = make_arg_parser(parser)
221 args = parser.parse_args()
222 configure_logging_from_args(args)
223 validate_parsed_serve_args(args)
224
225 uvloop.run(run_server(args))
226
227
228 if __name__ == "__main__":
229 main()