main
py 209 lines 6.62 KB
Raw
1 #!/usr/bin/env python3
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 # mypy: ignore-errors
5 """vLLM gRPC Server
6
7 Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer
8 from the smg-grpc-servicer package.
9
10 Usage:
11 python -m vllm.entrypoints.launchers.grpc_server --model <model_path>
12
13 Example:
14 python -m vllm.entrypoints.launchers.grpc_server \
15 --model meta-llama/Llama-2-7b-hf \
16 --host 0.0.0.0 \
17 --port 50051
18
19 """
20
21 import argparse
22 import asyncio
23 import signal
24 import sys
25 import time
26
27 try:
28 import grpc
29 from grpc_health.v1 import health_pb2_grpc
30 from grpc_reflection.v1alpha import reflection
31 from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
32 from smg_grpc_servicer.vllm.health_servicer import VllmHealthServicer
33 from smg_grpc_servicer.vllm.servicer import VllmEngineServicer
34 except ImportError as e:
35 raise ImportError(
36 "gRPC mode requires smg-grpc-servicer. "
37 "If not installed, run: pip install vllm[grpc]. "
38 "If already installed, there may be a broken import due to a "
39 "version mismatch — see the chained exception above for details."
40 ) from e
41
42 import uvloop
43
44 from vllm import envs
45 from vllm.engine.arg_utils import AsyncEngineArgs
46 from vllm.entrypoints.serve.utils.api_utils import log_version_and_model
47 from vllm.logger import configure_logging_from_args, init_logger
48 from vllm.usage.usage_lib import UsageContext
49 from vllm.utils.argparse_utils import FlexibleArgumentParser
50 from vllm.v1.engine.async_llm import AsyncLLM
51 from vllm.version import __version__ as VLLM_VERSION
52
53 logger = init_logger(__name__)
54
55
56 async def serve_grpc(args: argparse.Namespace):
57 """Main gRPC serving function.
58
59 Args:
60 args: Parsed command line arguments
61
62 """
63 log_version_and_model(logger, VLLM_VERSION, args.model)
64 logger.info("vLLM gRPC server args: %s", args)
65
66 start_time = time.time()
67
68 # Create engine args
69 engine_args = AsyncEngineArgs.from_cli_args(args)
70
71 # Build vLLM config
72 vllm_config = engine_args.create_engine_config(
73 usage_context=UsageContext.OPENAI_API_SERVER,
74 )
75
76 # Create AsyncLLM
77 async_llm = AsyncLLM.from_vllm_config(
78 vllm_config=vllm_config,
79 usage_context=UsageContext.OPENAI_API_SERVER,
80 enable_log_requests=args.enable_log_requests,
81 disable_log_stats=args.disable_log_stats,
82 )
83
84 # Create servicer
85 servicer = VllmEngineServicer(async_llm, start_time)
86
87 # Create gRPC server
88 server = grpc.aio.server(
89 options=[
90 ("grpc.max_send_message_length", -1),
91 ("grpc.max_receive_message_length", -1),
92 # A non-streaming Generate writes nothing while the engine
93 # decodes, so client keepalive and BDP-probe PINGs arriving in
94 # that window are strikes and eventually earn a
95 # GOAWAY(ENHANCE_YOUR_CALM, "too_many_pings") that fails every
96 # in-flight RPC. Disable the strike count to tolerate them.
97 ("grpc.http2.max_ping_strikes", 0),
98 # GRPC_ARG_HTTP2_MIN_RECV_PING_INTERVAL_WITHOUT_DATA_MS expands to
99 # a string with no "recv", and gRPC Core silently ignores unknown
100 # keys, so the misspelt form left the 300s default in force.
101 ("grpc.http2.min_ping_interval_without_data_ms", 10000),
102 # Without this, gRPC Core widens the floor to a fixed two hours
103 # once no streams are open, striking keepalives between requests.
104 ("grpc.keepalive_permit_without_calls", True),
105 ],
106 )
107
108 # Add servicer to server
109 vllm_engine_pb2_grpc.add_VllmEngineServicer_to_server(servicer, server)
110
111 # Add standard gRPC health service for Kubernetes probes
112 health_servicer = VllmHealthServicer(async_llm)
113 health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
114
115 # Enable reflection for grpcurl and other tools
116 service_names = (
117 vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name,
118 "grpc.health.v1.Health",
119 reflection.SERVICE_NAME,
120 )
121 reflection.enable_server_reflection(service_names, server)
122
123 # Bind to address
124 host = args.host or "0.0.0.0"
125 address = f"{host}:{args.port}"
126 server.add_insecure_port(address)
127
128 try:
129 # Start server
130 await server.start()
131 logger.info("vLLM gRPC server started on %s", address)
132 logger.info("Server is ready to accept requests")
133
134 # Start periodic stats logging (mirrors the HTTP server's lifespan task)
135 if not args.disable_log_stats:
136
137 async def _force_log():
138 while True:
139 await asyncio.sleep(envs.VLLM_LOG_STATS_INTERVAL)
140 await async_llm.do_log_stats()
141
142 stats_task = asyncio.create_task(_force_log())
143 else:
144 stats_task = None
145
146 # Handle shutdown signals
147 loop = asyncio.get_running_loop()
148 stop_event = asyncio.Event()
149
150 def signal_handler():
151 logger.info("Received shutdown signal")
152 stop_event.set()
153
154 for sig in (signal.SIGTERM, signal.SIGINT):
155 loop.add_signal_handler(sig, signal_handler)
156
157 try:
158 await stop_event.wait()
159 except KeyboardInterrupt:
160 logger.info("Interrupted by user")
161 finally:
162 logger.info("Shutting down vLLM gRPC server...")
163 if stats_task is not None:
164 stats_task.cancel()
165 try:
166 health_servicer.set_not_serving()
167 except Exception: # broad: must not prevent server.stop() / shutdown()
168 logger.warning("Failed to set health status to NOT_SERVING", exc_info=True)
169 await server.stop(grace=5.0)
170 logger.info("gRPC server stopped")
171 async_llm.shutdown()
172 logger.info("AsyncLLM engine stopped")
173 logger.info("Shutdown complete")
174
175
176 def main():
177 """Main entry point for python -m vllm.entrypoints.launchers.grpc_server."""
178 parser = FlexibleArgumentParser(
179 description="vLLM gRPC Server",
180 )
181
182 # Server args
183 parser.add_argument(
184 "--host",
185 type=str,
186 default="0.0.0.0",
187 help="Host to bind gRPC server to",
188 )
189 parser.add_argument(
190 "--port",
191 type=int,
192 default=50051,
193 help="Port to bind gRPC server to",
194 )
195 parser = AsyncEngineArgs.add_cli_args(parser)
196
197 args = parser.parse_args()
198 configure_logging_from_args(args)
199
200 # Run server
201 try:
202 uvloop.run(serve_grpc(args))
203 except Exception as e:
204 logger.exception("Server failed: %s", e)
205 sys.exit(1)
206
207
208 if __name__ == "__main__":
209 main()