@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 507ddbe992
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 23:00 UTC

feat(grpc): extract gRPC servicer into smg-grpc-servicer package, add --grpc flag to vllm serve (#36169)

Signed-off-by: Chang Su <chang.s.su@oracle.com> Co-authored-by: Nick Hill <nhill@redhat.com>

Chang Su committed Mar 10, 2026 at 03:29 UTC 507ddbe9927f421a1d574b283d1611044859a30d
13 files changed +57 -1242
pyproject.toml
-5
@@ -9,7 +9,6 @@ requires = [
9 "torch == 2.10.0",
10 "wheel",
11 "jinja2",
12 - "grpcio-tools==1.78.0",
12 ]
13 build-backend = "setuptools.build_meta"
14
@@ -57,10 +56,6 @@ include = ["vllm*"]
56 "vllm/third_party/**" = ["ALL"]
57 "vllm/version.py" = ["F401"]
58 "vllm/_version.py" = ["ALL"]
60 -# Exclude generated protobuf files
61 -"vllm/grpc/*_pb2.py" = ["ALL"]
62 -"vllm/grpc/*_pb2_grpc.py" = ["ALL"]
63 -"vllm/grpc/*_pb2.pyi" = ["ALL"]
59
60 [tool.ruff.lint]
61 select = [
requirements/build.txt
-1
@@ -10,4 +10,3 @@ jinja2>=3.1.6
10 regex
11 build
12 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*
13 -grpcio-tools==1.78.0 # Required for grpc entrypoints
requirements/common.txt
-2
@@ -51,8 +51,6 @@ openai-harmony >= 0.0.3 # Required for gpt-oss
51 anthropic >= 0.71.0
52 model-hosting-container-standards >= 0.1.13, < 1.0.0
53 mcp
54 -grpcio
55 -grpcio-reflection
54 opentelemetry-sdk >= 1.27.0
55 opentelemetry-api >= 1.27.0
56 opentelemetry-exporter-otlp >= 1.27.0
requirements/rocm.txt
-1
@@ -4,7 +4,6 @@
4 # The version of gRPC libraries should be consistent with each other
5 grpcio==1.78.0
6 grpcio-reflection==1.78.0
7 -grpcio-tools==1.78.0
7
8 numba == 0.61.2 # Required for N-gram speculative decoding
9
requirements/test.in
-1
@@ -51,7 +51,6 @@ tritonclient>=2.51.0
51 # The version of gRPC libraries should be consistent with each other
52 grpcio==1.78.0
53 grpcio-reflection==1.78.0
54 -grpcio-tools==1.78.0
54
55 arctic-inference == 0.1.1 # Required for suffix decoding test
56 numba == 0.61.2 # Required for N-gram speculative decoding
requirements/test.txt
-5
@@ -289,13 +289,10 @@ grpcio==1.78.0
289 # via
290 # -r requirements/test.in
291 # grpcio-reflection
292 - # grpcio-tools
292 # ray
293 # tensorboard
294 grpcio-reflection==1.78.0
295 # via -r requirements/test.in
297 -grpcio-tools==1.78.0
298 - # via -r requirements/test.in
296 h11==0.14.0
297 # via
298 # httpcore
@@ -765,7 +762,6 @@ protobuf==6.33.2
762 # google-api-core
763 # googleapis-common-protos
764 # grpcio-reflection
768 - # grpcio-tools
765 # opentelemetry-proto
766 # proto-plus
767 # ray
@@ -1045,7 +1041,6 @@ sentry-sdk==2.52.0
1041 # via wandb
1042 setuptools==77.0.3
1043 # via
1048 - # grpcio-tools
1044 # lightning-utilities
1045 # pytablewriter
1046 # tensorboard
setup.py
+3 -83
@@ -18,8 +18,6 @@ import torch
18 from packaging.version import Version, parse
19 from setuptools import Extension, setup
20 from setuptools.command.build_ext import build_ext
21 -from setuptools.command.build_py import build_py
22 -from setuptools.command.develop import develop
21 from setuptools_scm import get_version
22 from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
23
@@ -81,81 +79,6 @@ def is_freethreaded():
79 return bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
80
81
84 -def compile_grpc_protos():
85 - """Compile gRPC protobuf definitions during build.
86 -
87 - This generates *_pb2.py, *_pb2_grpc.py, and *_pb2.pyi files from
88 - the vllm_engine.proto definition.
89 - """
90 - try:
91 - from grpc_tools import protoc
92 - except ImportError:
93 - logger.warning(
94 - "grpcio-tools not installed, skipping gRPC proto compilation. "
95 - "gRPC server functionality will not be available."
96 - )
97 - return False
98 -
99 - proto_file = ROOT_DIR / "vllm" / "grpc" / "vllm_engine.proto"
100 - if not proto_file.exists():
101 - logger.warning("Proto file not found at %s, skipping compilation", proto_file)
102 - return False
103 -
104 - logger.info("Compiling gRPC protobuf: %s", proto_file)
105 -
106 - result = protoc.main(
107 - [
108 - "grpc_tools.protoc",
109 - f"--proto_path={ROOT_DIR}",
110 - f"--python_out={ROOT_DIR}",
111 - f"--grpc_python_out={ROOT_DIR}",
112 - f"--pyi_out={ROOT_DIR}",
113 - str(proto_file),
114 - ]
115 - )
116 -
117 - if result != 0:
118 - logger.error("protoc failed with exit code %s", result)
119 - return False
120 -
121 - # Add SPDX headers and mypy ignore to generated files
122 - spdx_header = (
123 - "# SPDX-License-Identifier: Apache-2.0\n"
124 - "# SPDX-FileCopyrightText: Copyright contributors to the vLLM project\n"
125 - "# mypy: ignore-errors\n"
126 - )
127 -
128 - grpc_dir = ROOT_DIR / "vllm" / "grpc"
129 - for generated_file in [
130 - grpc_dir / "vllm_engine_pb2.py",
131 - grpc_dir / "vllm_engine_pb2_grpc.py",
132 - grpc_dir / "vllm_engine_pb2.pyi",
133 - ]:
134 - if generated_file.exists():
135 - content = generated_file.read_text()
136 - if not content.startswith("# SPDX-License-Identifier"):
137 - generated_file.write_text(spdx_header + content)
138 -
139 - logger.info("gRPC protobuf compilation successful")
140 - return True
141 -
142 -
143 -class BuildPyAndGenerateGrpc(build_py):
144 - """Build Python modules and generate gRPC stubs from proto files."""
145 -
146 - def run(self):
147 - compile_grpc_protos()
148 - super().run()
149 -
150 -
151 -class DevelopAndGenerateGrpc(develop):
152 - """Develop mode that also generates gRPC stubs from proto files."""
153 -
154 - def run(self):
155 - compile_grpc_protos()
156 - super().run()
157 -
158 -
82 class CMakeExtension(Extension):
83 def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
84 super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa)
@@ -1028,17 +951,12 @@ if _no_device():
951 ext_modules = []
952
953 if not ext_modules:
1031 - cmdclass = {
1032 - "build_py": BuildPyAndGenerateGrpc,
1033 - "develop": DevelopAndGenerateGrpc,
1034 - }
954 + cmdclass = {}
955 else:
956 cmdclass = {
957 "build_ext": precompiled_build_ext
958 if envs.VLLM_USE_PRECOMPILED
959 else cmake_build_ext,
1040 - "build_py": BuildPyAndGenerateGrpc,
1041 - "develop": DevelopAndGenerateGrpc,
960 }
961
962 setup(
@@ -1064,6 +982,8 @@ setup(
982 "petit-kernel": ["petit-kernel"],
983 # Optional deps for Helion kernel development
984 "helion": ["helion"],
985 + # Optional deps for gRPC server (vllm serve --grpc)
986 + "grpc": ["smg-grpc-servicer >= 0.4.2"],
987 # Optional deps for OpenTelemetry tracing
988 "otel": [
989 "opentelemetry-sdk>=1.26.0",
tests/entrypoints/test_grpc_server.py deleted
-428
@@ -1,428 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -End-to-end tests for the vLLM gRPC server.
5 -"""
6 -
7 -import asyncio
8 -import socket
9 -import subprocess
10 -import sys
11 -import time
12 -
13 -import grpc
14 -import pytest
15 -import pytest_asyncio
16 -
17 -from vllm.grpc import vllm_engine_pb2, vllm_engine_pb2_grpc
18 -
19 -# Use a small model for fast testing
20 -MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
21 -
22 -
23 -def find_free_port() -> int:
24 - """Find a free port on localhost."""
25 - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
26 - s.bind(("", 0))
27 - s.listen(1)
28 - port = s.getsockname()[1]
29 - return port
30 -
31 -
32 -async def wait_for_server(port: int, timeout: float = 60.0) -> bool:
33 - """Wait for the gRPC server to be ready by trying health checks."""
34 - start_time = time.time()
35 - print("waiting for server to start...")
36 - while time.time() - start_time < timeout:
37 - try:
38 - channel = grpc.aio.insecure_channel(f"localhost:{port}")
39 - stub = vllm_engine_pb2_grpc.VllmEngineStub(channel)
40 - request = vllm_engine_pb2.HealthCheckRequest()
41 - response = await stub.HealthCheck(request, timeout=5.0)
42 - await channel.close()
43 - if response.healthy:
44 - print("server returned healthy=True")
45 - return True
46 - except Exception:
47 - await asyncio.sleep(0.5)
48 - return False
49 -
50 -
51 -class GrpcServerProcess:
52 - """Manages a gRPC server running in a subprocess."""
53 -
54 - def __init__(self):
55 - self.process: subprocess.Popen | None = None
56 - self.port: int | None = None
57 -
58 - async def start(self):
59 - """Start the gRPC server process."""
60 - self.port = find_free_port()
61 -
62 - # Start the server as a subprocess
63 - self.process = subprocess.Popen(
64 - [
65 - sys.executable,
66 - "-m",
67 - "vllm.entrypoints.grpc_server",
68 - "--model",
69 - MODEL_NAME,
70 - "--host",
71 - "localhost",
72 - "--port",
73 - str(self.port),
74 - "--max-num-batched-tokens",
75 - "512",
76 - "--disable-log-stats-server",
77 - ],
78 - )
79 -
80 - # Wait for server to be ready
81 - if not await wait_for_server(self.port):
82 - self.stop()
83 - raise RuntimeError("gRPC server failed to start within timeout")
84 -
85 - def stop(self):
86 - """Stop the gRPC server process."""
87 - if self.process:
88 - self.process.terminate()
89 - try:
90 - self.process.wait(timeout=10)
91 - except subprocess.TimeoutExpired:
92 - self.process.kill()
93 - self.process.wait()
94 -
95 -
96 -@pytest_asyncio.fixture(scope="module")
97 -async def grpc_server():
98 - """Fixture providing a running gRPC server in a subprocess."""
99 - server = GrpcServerProcess()
100 - await server.start()
101 -
102 - yield server
103 -
104 - server.stop()
105 -
106 -
107 -@pytest_asyncio.fixture
108 -async def grpc_client(grpc_server):
109 - """Fixture providing a gRPC client connected to the server."""
110 - channel = grpc.aio.insecure_channel(f"localhost:{grpc_server.port}")
111 - stub = vllm_engine_pb2_grpc.VllmEngineStub(channel)
112 -
113 - yield stub
114 -
115 - await channel.close()
116 -
117 -
118 -@pytest.mark.asyncio
119 -async def test_health_check(grpc_client):
120 - """Test the HealthCheck RPC."""
121 - request = vllm_engine_pb2.HealthCheckRequest()
122 - response = await grpc_client.HealthCheck(request)
123 -
124 - assert response.healthy is True
125 - assert response.message == "Health"
126 -
127 -
128 -@pytest.mark.asyncio
129 -async def test_get_model_info(grpc_client):
130 - """Test the GetModelInfo RPC."""
131 - request = vllm_engine_pb2.GetModelInfoRequest()
132 - response = await grpc_client.GetModelInfo(request)
133 -
134 - assert response.model_path == MODEL_NAME
135 - assert response.is_generation is True
136 - assert response.max_context_length > 0
137 - assert response.vocab_size > 0
138 - assert response.supports_vision is False
139 -
140 -
141 -@pytest.mark.asyncio
142 -async def test_get_server_info(grpc_client):
143 - """Test the GetServerInfo RPC."""
144 - request = vllm_engine_pb2.GetServerInfoRequest()
145 - response = await grpc_client.GetServerInfo(request)
146 -
147 - assert response.active_requests >= 0
148 - assert response.is_paused is False
149 - assert response.uptime_seconds >= 0
150 - assert response.server_type == "vllm-grpc"
151 - assert response.last_receive_timestamp > 0
152 -
153 -
154 -@pytest.mark.asyncio
155 -async def test_generate_non_streaming(grpc_client):
156 - """Test the Generate RPC in non-streaming mode."""
157 - # Create a simple request
158 - request = vllm_engine_pb2.GenerateRequest(
159 - request_id="test-non-streaming-1",
160 - tokenized=vllm_engine_pb2.TokenizedInput(
161 - original_text="Hello, my name is",
162 - input_ids=[15496, 11, 616, 1438, 318], # GPT-2 tokens for the prompt
163 - ),
164 - sampling_params=vllm_engine_pb2.SamplingParams(
165 - temperature=0.0,
166 - max_tokens=10,
167 - n=1,
168 - ),
169 - stream=False,
170 - )
171 -
172 - # Collect all responses
173 - responses = []
174 - async for response in grpc_client.Generate(request):
175 - responses.append(response)
176 -
177 - # Should have exactly one response (complete)
178 - assert len(responses) == 1
179 -
180 - # Check the response
181 - final_response = responses[0]
182 - assert final_response.HasField("complete")
183 -
184 - complete = final_response.complete
185 - assert len(complete.output_ids) > 0
186 - assert complete.finish_reason in ["stop", "length"]
187 - assert complete.prompt_tokens > 0
188 - assert complete.completion_tokens > 0
189 -
190 -
191 -@pytest.mark.asyncio
192 -async def test_generate_streaming(grpc_client):
193 - """Test the Generate RPC in streaming mode."""
194 - request = vllm_engine_pb2.GenerateRequest(
195 - request_id="test-streaming-1",
196 - tokenized=vllm_engine_pb2.TokenizedInput(
197 - original_text="The capital of France is",
198 - input_ids=[464, 3139, 286, 4881, 318], # GPT-2 tokens
199 - ),
200 - sampling_params=vllm_engine_pb2.SamplingParams(
201 - temperature=0.0, max_tokens=10, n=1
202 - ),
203 - stream=True,
204 - )
205 -
206 - # Collect all responses
207 - chunks = []
208 - complete_response = None
209 -
210 - async for response in grpc_client.Generate(request):
211 - if response.HasField("chunk"):
212 - chunks.append(response.chunk)
213 - elif response.HasField("complete"):
214 - complete_response = response.complete
215 -
216 - # Should have received some chunks
217 - assert len(chunks) >= 0 # May have 0 chunks if generation is very fast
218 -
219 - # Should have a final complete response
220 - assert complete_response is not None
221 - assert complete_response.finish_reason in ["stop", "length"]
222 - assert complete_response.prompt_tokens > 0
223 -
224 - # Verify chunk structure
225 - for chunk in chunks:
226 - assert chunk.prompt_tokens > 0
227 - assert chunk.completion_tokens >= 0
228 -
229 -
230 -@pytest.mark.asyncio
231 -async def test_generate_with_different_sampling_params(grpc_client):
232 - """Test Generate with various sampling parameters."""
233 - # Test with temperature
234 - request = vllm_engine_pb2.GenerateRequest(
235 - request_id="test-sampling-temp",
236 - tokenized=vllm_engine_pb2.TokenizedInput(
237 - original_text="Hello",
238 - input_ids=[15496],
239 - ),
240 - sampling_params=vllm_engine_pb2.SamplingParams(
241 - temperature=0.8, top_p=0.95, max_tokens=5
242 - ),
243 - stream=False,
244 - )
245 -
246 - responses = [r async for r in grpc_client.Generate(request)]
247 - assert len(responses) == 1
248 - assert responses[0].HasField("complete")
249 -
250 - # Test with top_k
251 - request = vllm_engine_pb2.GenerateRequest(
252 - request_id="test-sampling-topk",
253 - tokenized=vllm_engine_pb2.TokenizedInput(
254 - original_text="Hello",
255 - input_ids=[15496],
256 - ),
257 - sampling_params=vllm_engine_pb2.SamplingParams(
258 - temperature=1.0, top_k=50, max_tokens=5
259 - ),
260 - stream=False,
261 - )
262 -
263 - responses = [r async for r in grpc_client.Generate(request)]
264 - assert len(responses) == 1
265 - assert responses[0].HasField("complete")
266 -
267 -
268 -@pytest.mark.asyncio
269 -async def test_generate_with_stop_strings(grpc_client):
270 - """Test Generate with stop strings."""
271 - request = vllm_engine_pb2.GenerateRequest(
272 - request_id="test-stop-strings",
273 - tokenized=vllm_engine_pb2.TokenizedInput(
274 - original_text="Hello",
275 - input_ids=[15496],
276 - ),
277 - sampling_params=vllm_engine_pb2.SamplingParams(
278 - temperature=0.0,
279 - max_tokens=20,
280 - stop=["\n", "END"],
281 - ),
282 - stream=False,
283 - )
284 -
285 - responses = [r async for r in grpc_client.Generate(request)]
286 - assert len(responses) == 1
287 - assert responses[0].HasField("complete")
288 -
289 - complete = responses[0].complete
290 - assert complete.finish_reason in ["stop", "length"]
291 -
292 -
293 -@pytest.mark.asyncio
294 -async def test_generate_multiple_requests(grpc_client):
295 - """Test handling multiple concurrent Generate requests."""
296 -
297 - async def make_request(request_id: str):
298 - request = vllm_engine_pb2.GenerateRequest(
299 - request_id=request_id,
300 - tokenized=vllm_engine_pb2.TokenizedInput(
301 - original_text="Hello",
302 - input_ids=[15496],
303 - ),
304 - sampling_params=vllm_engine_pb2.SamplingParams(
305 - temperature=0.0, max_tokens=5
306 - ),
307 - stream=False,
308 - )
309 -
310 - responses = [r async for r in grpc_client.Generate(request)]
311 - return responses[0]
312 -
313 - # Send multiple requests concurrently
314 - tasks = [make_request(f"test-concurrent-{i}") for i in range(3)]
315 - responses = await asyncio.gather(*tasks)
316 -
317 - # Verify all requests completed successfully
318 - assert len(responses) == 3
319 - for i, response in enumerate(responses):
320 - assert response.HasField("complete")
321 -
322 -
323 -@pytest.mark.asyncio
324 -async def test_generate_with_seed(grpc_client):
325 - """Test Generate with a fixed seed for reproducibility."""
326 -
327 - def make_request(request_id: str, seed: int):
328 - return vllm_engine_pb2.GenerateRequest(
329 - request_id=request_id,
330 - tokenized=vllm_engine_pb2.TokenizedInput(
331 - original_text="The future of AI is",
332 - input_ids=[464, 2003, 286, 9552, 318],
333 - ),
334 - sampling_params=vllm_engine_pb2.SamplingParams(
335 - temperature=1.0, max_tokens=10, seed=seed
336 - ),
337 - stream=False,
338 - )
339 -
340 - # Make two requests with the same seed
341 - request1 = make_request("test-seed-1", 42)
342 - request2 = make_request("test-seed-2", 42)
343 -
344 - response_list1 = [r async for r in grpc_client.Generate(request1)]
345 - response_list2 = [r async for r in grpc_client.Generate(request2)]
346 -
347 - # Both should complete successfully
348 - assert len(response_list1) == 1
349 - assert len(response_list2) == 1
350 - assert response_list1[0].HasField("complete")
351 - assert response_list2[0].HasField("complete")
352 -
353 - # With the same seed, outputs should be identical
354 - output_ids1 = list(response_list1[0].complete.output_ids)
355 - output_ids2 = list(response_list2[0].complete.output_ids)
356 - assert output_ids1 == output_ids2
357 -
358 -
359 -@pytest.mark.asyncio
360 -async def test_generate_error_handling(grpc_client):
361 - """Test error handling in Generate RPC."""
362 - # Request with invalid top_p value (-33)
363 - request = vllm_engine_pb2.GenerateRequest(
364 - request_id="test-error-invalid-topp",
365 - sampling_params=vllm_engine_pb2.SamplingParams(
366 - temperature=0.0, max_tokens=10, top_p=-33
367 - ),
368 - stream=False,
369 - )
370 -
371 - # Should raise an error response
372 - with pytest.raises(grpc.RpcError) as exc_info:
373 - _ = [r async for r in grpc_client.Generate(request)]
374 -
375 - assert exc_info.value.code() == grpc.StatusCode.INVALID_ARGUMENT
376 - assert "top_p must be in (0, 1], got -33.0" in exc_info.value.details()
377 -
378 -
379 -@pytest.mark.asyncio
380 -async def test_abort_request(grpc_client):
381 - """Test the out-of-band Abort RPC."""
382 - request_id = "test-abort-1"
383 -
384 - # Start a long-running streaming generate request
385 - generate_request = vllm_engine_pb2.GenerateRequest(
386 - request_id=request_id,
387 - tokenized=vllm_engine_pb2.TokenizedInput(
388 - original_text="Hello",
389 - input_ids=[15496],
390 - ),
391 - sampling_params=vllm_engine_pb2.SamplingParams(
392 - temperature=0.0,
393 - min_tokens=500,
394 - max_tokens=500, # Request many tokens to ensure it runs long enough
395 - ),
396 - stream=True,
397 - )
398 -
399 - # Track whether we were aborted
400 - was_aborted = False
401 - received_chunks = 0
402 -
403 - async def run_generate():
404 - nonlocal was_aborted, received_chunks
405 - async for response in grpc_client.Generate(generate_request):
406 - if response.HasField("chunk"):
407 - received_chunks += 1
408 -
409 - if response.HasField("complete"):
410 - complete = response.complete
411 - was_aborted = complete.finish_reason == "abort"
412 - else:
413 - was_aborted = False
414 -
415 - async def abort_after_delay():
416 - # Small delay to ensure generate has started
417 - await asyncio.sleep(0.1)
418 - abort_request = vllm_engine_pb2.AbortRequest(request_ids=[request_id])
419 - await grpc_client.Abort(abort_request)
420 -
421 - # Run generate and abort concurrently
422 - await asyncio.gather(run_generate(), abort_after_delay())
423 -
424 - # The request should have been aborted (received final chunk with
425 - # "abort" finish reason) and finished early due to the abort.
426 - assert was_aborted and received_chunks < 500, (
427 - "Request should have been aborted before generating all 500 tokens"
428 - )
vllm/entrypoints/cli/serve.py
+13
@@ -51,6 +51,12 @@ class ServeSubcommand(CLISubcommand):
51 if hasattr(args, "model_tag") and args.model_tag is not None:
52 args.model = args.model_tag
53
54 + if getattr(args, "grpc", False):
55 + from vllm.entrypoints.grpc_server import serve_grpc
56 +
57 + uvloop.run(serve_grpc(args))
58 + return
59 +
60 if args.headless:
61 if args.api_server_count is not None and args.api_server_count > 0:
62 raise ValueError(
@@ -127,6 +133,13 @@ class ServeSubcommand(CLISubcommand):
133 )
134
135 serve_parser = make_arg_parser(serve_parser)
136 + serve_parser.add_argument(
137 + "--grpc",
138 + action="store_true",
139 + default=False,
140 + help="Launch a gRPC server instead of the HTTP OpenAI-compatible "
141 + "server. Requires: pip install vllm[grpc].",
142 + )
143 serve_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name)
144 return serve_parser
145
vllm/entrypoints/grpc_server.py
+41 -410
@@ -5,7 +5,8 @@
5 """
6 vLLM gRPC Server
7
8 -Starts a gRPC server for vLLM using the VllmEngine protocol.
8 +Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer
9 +from the smg-grpc-servicer package.
10
11 Usage:
12 python -m vllm.entrypoints.grpc_server --model <model_path>
@@ -22,19 +23,23 @@ import asyncio
23 import signal
24 import sys
25 import time
25 -from collections.abc import AsyncGenerator
26
27 -import grpc
27 +try:
28 + import grpc
29 + from grpc_reflection.v1alpha import reflection
30 + from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
31 + from smg_grpc_servicer.vllm.servicer import VllmEngineServicer
32 +except ImportError:
33 + raise ImportError(
34 + "smg-grpc-servicer is required for gRPC mode. "
35 + "Install it with: pip install vllm[grpc]"
36 + ) from None
37 +
38 import uvloop
29 -from grpc_reflection.v1alpha import reflection
39
31 -from vllm import SamplingParams, TextPrompt, TokensPrompt
40 from vllm.engine.arg_utils import AsyncEngineArgs
41 from vllm.entrypoints.utils import log_version_and_model
34 -from vllm.grpc import vllm_engine_pb2, vllm_engine_pb2_grpc
42 from vllm.logger import init_logger
36 -from vllm.outputs import RequestOutput
37 -from vllm.sampling_params import RequestOutputKind, StructuredOutputsParams
43 from vllm.usage.usage_lib import UsageContext
44 from vllm.utils.argparse_utils import FlexibleArgumentParser
45 from vllm.v1.engine.async_llm import AsyncLLM
@@ -43,377 +48,9 @@ from vllm.version import __version__ as VLLM_VERSION
48 logger = init_logger(__name__)
49
50
46 -class VllmEngineServicer(vllm_engine_pb2_grpc.VllmEngineServicer):
47 - """
48 - gRPC servicer implementing the VllmEngine service.
49 -
50 - Handles 6 RPCs:
51 - - Generate: Streaming text generation
52 - - Embed: Embeddings (TODO)
53 - - HealthCheck: Health probe
54 - - Abort: Cancel requests out-of-band
55 - - GetModelInfo: Model metadata
56 - - GetServerInfo: Server state
57 - """
58 -
59 - def __init__(self, async_llm: AsyncLLM, start_time: float):
60 - """
61 - Initialize the servicer.
62 -
63 - Args:
64 - async_llm: The AsyncLLM instance
65 - start_time: The server start time, in seconds since epoch
66 - """
67 - self.async_llm = async_llm
68 - self.start_time = start_time
69 - logger.info("VllmEngineServicer initialized")
70 -
71 - async def Generate(
72 - self,
73 - request: vllm_engine_pb2.GenerateRequest,
74 - context: grpc.aio.ServicerContext,
75 - ) -> AsyncGenerator[vllm_engine_pb2.GenerateResponse, None]:
76 - """
77 - Handle streaming generation requests.
78 -
79 - Args:
80 - request: The GenerateRequest protobuf
81 - context: gRPC context
82 -
83 - Yields:
84 - GenerateResponse protobuf messages (streaming)
85 - """
86 - request_id = request.request_id
87 - logger.debug("Generate request %s received.", request_id)
88 -
89 - try:
90 - # Extract tokenized input
91 - if request.WhichOneof("input") == "tokenized":
92 - prompt: TokensPrompt = {
93 - "prompt_token_ids": list(request.tokenized.input_ids)
94 - }
95 - if request.tokenized.original_text:
96 - prompt["prompt"] = request.tokenized.original_text
97 - else:
98 - prompt: TextPrompt = {"prompt": request.text}
99 -
100 - # Build sampling params with detokenize=False
101 - sampling_params = self._sampling_params_from_proto(
102 - request.sampling_params, stream=request.stream
103 - )
104 - tokenization_kwargs = self._tokenization_kwargs_from_proto(
105 - request.sampling_params
106 - )
107 -
108 - async for output in self.async_llm.generate(
109 - prompt=prompt,
110 - sampling_params=sampling_params,
111 - request_id=request_id,
112 - tokenization_kwargs=tokenization_kwargs,
113 - ):
114 - # Convert vLLM output to protobuf
115 - # For streaming, always send chunks
116 - if request.stream:
117 - yield self._chunk_response(output)
118 -
119 - # Send complete response when finished
120 - if output.finished:
121 - yield self._complete_response(output)
122 -
123 - except ValueError as e:
124 - # Invalid request error (equiv to 400).
125 - await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e))
126 - except Exception as e:
127 - logger.exception("Error in Generate for request %s", request_id)
128 - await context.abort(grpc.StatusCode.INTERNAL, str(e))
129 -
130 - async def Embed(
131 - self,
132 - request: vllm_engine_pb2.EmbedRequest,
133 - context: grpc.aio.ServicerContext,
134 - ) -> vllm_engine_pb2.EmbedResponse:
135 - """
136 - Handle embedding requests.
137 -
138 - TODO: Implement in Phase 4
139 -
140 - Args:
141 - request: The EmbedRequest protobuf
142 - context: gRPC context
143 -
144 - Returns:
145 - EmbedResponse protobuf
146 - """
147 - logger.warning("Embed RPC not yet implemented")
148 - await context.abort(
149 - grpc.StatusCode.UNIMPLEMENTED, "Embed RPC not yet implemented"
150 - )
151 -
152 - async def HealthCheck(
153 - self,
154 - request: vllm_engine_pb2.HealthCheckRequest,
155 - context: grpc.aio.ServicerContext,
156 - ) -> vllm_engine_pb2.HealthCheckResponse:
157 - """
158 - Handle health check requests.
159 -
160 - Args:
161 - request: The HealthCheckRequest protobuf
162 - context: gRPC context
163 -
164 - Returns:
165 - HealthCheckResponse protobuf
166 - """
167 - is_healthy = not self.async_llm.errored
168 - message = "Health" if is_healthy else "Engine is not alive"
169 -
170 - logger.debug("HealthCheck request: healthy=%s, message=%s", is_healthy, message)
171 -
172 - return vllm_engine_pb2.HealthCheckResponse(healthy=is_healthy, message=message)
173 -
174 - async def Abort(
175 - self,
176 - request: vllm_engine_pb2.AbortRequest,
177 - context: grpc.aio.ServicerContext,
178 - ) -> vllm_engine_pb2.AbortResponse:
179 - """
180 - Out-of-band abort requests.
181 -
182 - Args:
183 - request: The AbortRequest protobuf
184 - context: gRPC context
185 -
186 - Returns:
187 - AbortResponse protobuf
188 - """
189 - request_ids = request.request_ids
190 - logger.debug("Abort requests: %s", request_ids)
191 -
192 - await self.async_llm.abort(request_ids)
193 - return vllm_engine_pb2.AbortResponse()
194 -
195 - async def GetModelInfo(
196 - self,
197 - request: vllm_engine_pb2.GetModelInfoRequest,
198 - context: grpc.aio.ServicerContext,
199 - ) -> vllm_engine_pb2.GetModelInfoResponse:
200 - """
201 - Handle model info requests.
202 -
203 - Args:
204 - request: The GetModelInfoRequest protobuf
205 - context: gRPC context
206 -
207 - Returns:
208 - GetModelInfoResponse protobuf
209 - """
210 - model_config = self.async_llm.model_config
211 -
212 - return vllm_engine_pb2.GetModelInfoResponse(
213 - model_path=model_config.model,
214 - is_generation=model_config.runner_type == "generate",
215 - max_context_length=model_config.max_model_len,
216 - vocab_size=model_config.get_vocab_size(),
217 - supports_vision=model_config.is_multimodal_model,
218 - )
219 -
220 - async def GetServerInfo(
221 - self,
222 - request: vllm_engine_pb2.GetServerInfoRequest,
223 - context: grpc.aio.ServicerContext,
224 - ) -> vllm_engine_pb2.GetServerInfoResponse:
225 - """
226 - Handle server info requests.
227 -
228 - Args:
229 - request: The GetServerInfoRequest protobuf
230 - context: gRPC context
231 -
232 - Returns:
233 - GetServerInfoResponse protobuf
234 - """
235 - num_requests = self.async_llm.output_processor.get_num_unfinished_requests()
236 -
237 - return vllm_engine_pb2.GetServerInfoResponse(
238 - active_requests=num_requests,
239 - is_paused=False, # TODO
240 - last_receive_timestamp=time.time(), # TODO looks wrong?
241 - uptime_seconds=time.time() - self.start_time,
242 - server_type="vllm-grpc",
243 - )
244 -
245 - # ========== Helper methods ==========
246 -
247 - @staticmethod
248 - def _sampling_params_from_proto(
249 - params: vllm_engine_pb2.SamplingParams, stream: bool = True
250 - ) -> SamplingParams:
251 - """
252 - Convert protobuf SamplingParams to vLLM SamplingParams.
253 -
254 - Args:
255 - params: Protobuf SamplingParams message
256 - stream: Whether streaming is enabled
257 -
258 - Returns:
259 - vLLM SamplingParams with detokenize=False and structured_outputs
260 - """
261 - # Build stop sequences
262 - stop = list(params.stop) if params.stop else None
263 - stop_token_ids = list(params.stop_token_ids) if params.stop_token_ids else None
264 -
265 - # Handle structured outputs constraints
266 - structured_outputs = None
267 - constraint_field = params.WhichOneof("constraint")
268 - if constraint_field:
269 - if constraint_field == "json_schema":
270 - structured_outputs = StructuredOutputsParams(json=params.json_schema)
271 - elif constraint_field == "regex":
272 - structured_outputs = StructuredOutputsParams(regex=params.regex)
273 - elif constraint_field == "grammar":
274 - structured_outputs = StructuredOutputsParams(grammar=params.grammar)
275 - elif constraint_field == "structural_tag":
276 - structured_outputs = StructuredOutputsParams(
277 - structural_tag=params.structural_tag
278 - )
279 - elif constraint_field == "json_object":
280 - structured_outputs = StructuredOutputsParams(
281 - json_object=params.json_object
282 - )
283 - elif constraint_field == "choice":
284 - structured_outputs = StructuredOutputsParams(
285 - choice=list(params.choice.choices)
286 - )
287 -
288 - # Create SamplingParams
289 - # output_kind=DELTA: Return only new tokens in each chunk (for streaming)
290 - return SamplingParams(
291 - temperature=params.temperature if params.HasField("temperature") else 1.0,
292 - top_p=params.top_p if params.top_p != 0.0 else 1.0,
293 - top_k=params.top_k,
294 - min_p=params.min_p,
295 - frequency_penalty=params.frequency_penalty,
296 - presence_penalty=params.presence_penalty,
297 - repetition_penalty=params.repetition_penalty
298 - if params.repetition_penalty != 0.0
299 - else 1.0,
300 - max_tokens=params.max_tokens if params.HasField("max_tokens") else None,
301 - min_tokens=params.min_tokens,
302 - stop=stop,
303 - stop_token_ids=stop_token_ids,
304 - skip_special_tokens=params.skip_special_tokens,
305 - spaces_between_special_tokens=params.spaces_between_special_tokens,
306 - ignore_eos=params.ignore_eos,
307 - n=params.n if params.n > 0 else 1,
308 - logprobs=params.logprobs if params.HasField("logprobs") else None,
309 - prompt_logprobs=params.prompt_logprobs
310 - if params.HasField("prompt_logprobs")
311 - else None,
312 - seed=params.seed if params.HasField("seed") else None,
313 - include_stop_str_in_output=params.include_stop_str_in_output,
314 - logit_bias=dict(params.logit_bias) if params.logit_bias else None,
315 - structured_outputs=structured_outputs,
316 - # detokenize must be True if stop strings are used
317 - detokenize=bool(stop),
318 - output_kind=RequestOutputKind.DELTA
319 - if stream
320 - else RequestOutputKind.FINAL_ONLY,
321 - )
322 -
323 - @staticmethod
324 - def _tokenization_kwargs_from_proto(
325 - params: vllm_engine_pb2.SamplingParams,
326 - ) -> dict[str, int] | None:
327 - if params.HasField("truncate_prompt_tokens"):
328 - return {"truncate_prompt_tokens": params.truncate_prompt_tokens}
329 - return None
330 -
331 - @staticmethod
332 - def _chunk_response(output: RequestOutput) -> vllm_engine_pb2.GenerateResponse:
333 - """
334 - Build a streaming chunk response from vLLM output.
335 - When output_kind=DELTA, vLLM returns only new tokens automatically.
336 -
337 - Args:
338 - output: vLLM RequestOutput (with delta tokens when output_kind=DELTA)
339 -
340 - Returns:
341 - GenerateResponse with chunk field set
342 - """
343 - # Get the completion output (first one if n > 1)
344 - completion = output.outputs[0] if output.outputs else None
345 -
346 - if completion is None:
347 - # Empty chunk
348 - return vllm_engine_pb2.GenerateResponse(
349 - chunk=vllm_engine_pb2.GenerateStreamChunk(
350 - token_ids=[],
351 - prompt_tokens=0,
352 - completion_tokens=0,
353 - cached_tokens=0,
354 - ),
355 - )
356 -
357 - # When output_kind=DELTA, completion.token_ids contains only new tokens
358 - # vLLM handles the delta logic internally
359 - # completion_tokens = delta count (client will accumulate)
360 - return vllm_engine_pb2.GenerateResponse(
361 - chunk=vllm_engine_pb2.GenerateStreamChunk(
362 - token_ids=completion.token_ids,
363 - prompt_tokens=len(output.prompt_token_ids)
364 - if output.prompt_token_ids
365 - else 0,
366 - completion_tokens=len(completion.token_ids), # Delta count
367 - cached_tokens=output.num_cached_tokens,
368 - ),
369 - )
370 -
371 - @staticmethod
372 - def _complete_response(output: RequestOutput) -> vllm_engine_pb2.GenerateResponse:
373 - """
374 - Build a final completion response from vLLM output.
375 -
376 - Args:
377 - output: vLLM RequestOutput (finished=True)
378 -
379 - Returns:
380 - GenerateResponse with complete field set
381 - """
382 - # Get the completion output (first one if n > 1)
383 - completion = output.outputs[0] if output.outputs else None
384 -
385 - if completion is None:
386 - # Empty completion
387 - return vllm_engine_pb2.GenerateResponse(
388 - complete=vllm_engine_pb2.GenerateComplete(
389 - output_ids=[],
390 - finish_reason="error",
391 - prompt_tokens=0,
392 - completion_tokens=0,
393 - cached_tokens=0,
394 - ),
395 - )
396 -
397 - # Build complete response
398 - # When streaming (DELTA mode): completion.token_ids will be empty/last delta
399 - # When non-streaming (FINAL_ONLY mode): completion.token_ids has all tokens
400 - # Client will accumulate token counts for streaming
401 - return vllm_engine_pb2.GenerateResponse(
402 - complete=vllm_engine_pb2.GenerateComplete(
403 - output_ids=completion.token_ids,
404 - finish_reason=completion.finish_reason or "stop",
405 - prompt_tokens=len(output.prompt_token_ids)
406 - if output.prompt_token_ids
407 - else 0,
408 - completion_tokens=len(completion.token_ids),
409 - cached_tokens=output.num_cached_tokens,
410 - ),
411 - )
412 -
413 -
51 async def serve_grpc(args: argparse.Namespace):
52 """
416 - Main serving function.
53 + Main gRPC serving function.
54
55 Args:
56 args: Parsed command line arguments
@@ -428,7 +65,7 @@ async def serve_grpc(args: argparse.Namespace):
65
66 # Build vLLM config
67 vllm_config = engine_args.create_engine_config(
431 - usage_context=UsageContext.OPENAI_API_SERVER
68 + usage_context=UsageContext.OPENAI_API_SERVER,
69 )
70
71 # Create AsyncLLM
@@ -436,7 +73,7 @@ async def serve_grpc(args: argparse.Namespace):
73 vllm_config=vllm_config,
74 usage_context=UsageContext.OPENAI_API_SERVER,
75 enable_log_requests=args.enable_log_requests,
439 - disable_log_stats=args.disable_log_stats_server,
76 + disable_log_stats=args.disable_log_stats,
77 )
78
79 # Create servicer
@@ -447,6 +84,11 @@ async def serve_grpc(args: argparse.Namespace):
84 options=[
85 ("grpc.max_send_message_length", -1),
86 ("grpc.max_receive_message_length", -1),
87 + # Tolerate client keepalive pings every 10s (default 300s is too
88 + # strict for non-streaming requests where no DATA frames flow
89 + # during generation)
90 + ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000),
91 + ("grpc.keepalive_permit_without_calls", True),
92 ],
93 )
94
@@ -461,46 +103,42 @@ async def serve_grpc(args: argparse.Namespace):
103 reflection.enable_server_reflection(service_names, server)
104
105 # Bind to address
464 - address = f"{args.host}:{args.port}"
106 + host = args.host or "0.0.0.0"
107 + address = f"{host}:{args.port}"
108 server.add_insecure_port(address)
109
467 - # Start server
468 - await server.start()
469 - logger.info("vLLM gRPC server started on %s", address)
470 - logger.info("Server is ready to accept requests")
110 + try:
111 + # Start server
112 + await server.start()
113 + logger.info("vLLM gRPC server started on %s", address)
114 + logger.info("Server is ready to accept requests")
115
472 - # Handle shutdown signals
473 - loop = asyncio.get_running_loop()
474 - stop_event = asyncio.Event()
116 + # Handle shutdown signals
117 + loop = asyncio.get_running_loop()
118 + stop_event = asyncio.Event()
119
476 - def signal_handler():
477 - logger.info("Received shutdown signal")
478 - stop_event.set()
120 + def signal_handler():
121 + logger.info("Received shutdown signal")
122 + stop_event.set()
123
480 - for sig in (signal.SIGTERM, signal.SIGINT):
481 - loop.add_signal_handler(sig, signal_handler)
124 + for sig in (signal.SIGTERM, signal.SIGINT):
125 + loop.add_signal_handler(sig, signal_handler)
126
483 - # Serve until shutdown signal
484 - try:
485 - await stop_event.wait()
486 - except KeyboardInterrupt:
487 - logger.info("Interrupted by user")
127 + try:
128 + await stop_event.wait()
129 + except KeyboardInterrupt:
130 + logger.info("Interrupted by user")
131 finally:
132 logger.info("Shutting down vLLM gRPC server...")
490 -
491 - # Stop gRPC server
133 await server.stop(grace=5.0)
134 logger.info("gRPC server stopped")
494 -
495 - # Shutdown AsyncLLM
135 async_llm.shutdown()
136 logger.info("AsyncLLM engine stopped")
498 -
137 logger.info("Shutdown complete")
138
139
140 def main():
503 - """Main entry point."""
141 + """Main entry point for python -m vllm.entrypoints.grpc_server."""
142 parser = FlexibleArgumentParser(
143 description="vLLM gRPC Server",
144 )
@@ -518,13 +156,6 @@ def main():
156 default=50051,
157 help="Port to bind gRPC server to",
158 )
521 - parser.add_argument(
522 - "--disable-log-stats-server",
523 - action="store_true",
524 - help="Disable stats logging on server side",
525 - )
526 -
527 - # Add vLLM engine args
159 parser = AsyncEngineArgs.add_cli_args(parser)
160
161 args = parser.parse_args()
vllm/grpc/__init__.py deleted
-17
@@ -1,17 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -vLLM gRPC protocol definitions.
5 -
6 -This module contains the protocol buffer definitions for vLLM's gRPC API.
7 -The protobuf files are compiled into Python code using grpcio-tools.
8 -"""
9 -
10 -# These imports will be available after protobuf compilation
11 -# from vllm.grpc import vllm_engine_pb2
12 -# from vllm.grpc import vllm_engine_pb2_grpc
13 -
14 -__all__ = [
15 - "vllm_engine_pb2",
16 - "vllm_engine_pb2_grpc",
17 -]
vllm/grpc/compile_protos.py deleted
-94
@@ -1,94 +0,0 @@
1 -#!/usr/bin/env python3
2 -# SPDX-License-Identifier: Apache-2.0
3 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 -"""
5 -Compile vLLM protobuf definitions into Python code.
6 -
7 -This script uses grpcio-tools to generate *_pb2.py, *_pb2_grpc.py, and
8 -*_pb2.pyi (type stubs) files from the vllm_engine.proto definition.
9 -
10 -NOTE: Proto compilation happens automatically during package build (via setup.py).
11 -This script is provided for developers who want to regenerate protos manually,
12 -e.g., after modifying vllm_engine.proto.
13 -
14 -Usage:
15 - python vllm/grpc/compile_protos.py
16 -
17 -Requirements:
18 - pip install grpcio-tools
19 -"""
20 -
21 -import sys
22 -from pathlib import Path
23 -
24 -
25 -def compile_protos():
26 - """Compile protobuf definitions."""
27 - # Get the vllm package root directory
28 - script_dir = Path(__file__).parent
29 - vllm_package_root = script_dir.parent.parent # vllm/vllm/grpc -> vllm/
30 -
31 - proto_file = script_dir / "vllm_engine.proto"
32 -
33 - if not proto_file.exists():
34 - print(f"Error: Proto file not found at {proto_file}")
35 - return 1
36 -
37 - print(f"Compiling protobuf: {proto_file}")
38 - print(f"Output directory: {script_dir}")
39 -
40 - # Compile the proto file
41 - # We use vllm/vllm as the proto_path so that the package is vllm.grpc.engine
42 - try:
43 - from grpc_tools import protoc
44 -
45 - result = protoc.main(
46 - [
47 - "grpc_tools.protoc",
48 - f"--proto_path={vllm_package_root}",
49 - f"--python_out={vllm_package_root}",
50 - f"--grpc_python_out={vllm_package_root}",
51 - f"--pyi_out={vllm_package_root}", # Generate type stubs
52 - str(script_dir / "vllm_engine.proto"),
53 - ]
54 - )
55 -
56 - if result == 0:
57 - # Add SPDX headers to generated files
58 - spdx_header = (
59 - "# SPDX-License-Identifier: Apache-2.0\n"
60 - "# SPDX-FileCopyrightText: Copyright contributors to the vLLM project\n"
61 - )
62 -
63 - for generated_file in [
64 - script_dir / "vllm_engine_pb2.py",
65 - script_dir / "vllm_engine_pb2_grpc.py",
66 - script_dir / "vllm_engine_pb2.pyi",
67 - ]:
68 - if generated_file.exists():
69 - content = generated_file.read_text()
70 - if not content.startswith("# SPDX-License-Identifier"):
71 - # Add mypy ignore-errors comment for all generated files
72 - header = spdx_header + "# mypy: ignore-errors\n"
73 - generated_file.write_text(header + content)
74 -
75 - print("✓ Protobuf compilation successful!")
76 - print(f" Generated: {script_dir / 'vllm_engine_pb2.py'}")
77 - print(f" Generated: {script_dir / 'vllm_engine_pb2_grpc.py'}")
78 - print(f" Generated: {script_dir / 'vllm_engine_pb2.pyi'} (type stubs)")
79 - return 0
80 - else:
81 - print(f"Error: protoc returned {result}")
82 - return result
83 -
84 - except ImportError:
85 - print("Error: grpcio-tools not installed")
86 - print("Install with: pip install grpcio-tools")
87 - return 1
88 - except Exception as e:
89 - print(f"Error during compilation: {e}")
90 - return 1
91 -
92 -
93 -if __name__ == "__main__":
94 - sys.exit(compile_protos())
vllm/grpc/vllm_engine.proto deleted
-195
@@ -1,195 +0,0 @@
1 -syntax = "proto3";
2 -
3 -package vllm.grpc.engine;
4 -
5 -// Service definition for vLLM engine communication
6 -// This protocol is designed for efficient binary communication between
7 -// the Rust router and vLLM Python engine (AsyncLLM).
8 -service VllmEngine {
9 - // Submit a generation request (supports streaming)
10 - rpc Generate(GenerateRequest) returns (stream GenerateResponse);
11 -
12 - // Submit an embedding request
13 - rpc Embed(EmbedRequest) returns (EmbedResponse);
14 -
15 - // Health check
16 - rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse);
17 -
18 - // Abort a running request
19 - rpc Abort(AbortRequest) returns (AbortResponse);
20 -
21 - // Get model information
22 - rpc GetModelInfo(GetModelInfoRequest) returns (GetModelInfoResponse);
23 -
24 - // Get server information
25 - rpc GetServerInfo(GetServerInfoRequest) returns (GetServerInfoResponse);
26 -}
27 -
28 -// =====================
29 -// Common Types
30 -// =====================
31 -
32 -// Sampling parameters for text generation
33 -message SamplingParams {
34 - optional float temperature = 1;
35 - float top_p = 2;
36 - uint32 top_k = 3;
37 - float min_p = 4;
38 - float frequency_penalty = 5;
39 - float presence_penalty = 6;
40 - float repetition_penalty = 7;
41 -
42 - optional uint32 max_tokens = 8;
43 - uint32 min_tokens = 9;
44 -
45 - repeated string stop = 10;
46 - repeated uint32 stop_token_ids = 11;
47 -
48 - bool skip_special_tokens = 12;
49 - bool spaces_between_special_tokens = 13;
50 - bool ignore_eos = 14;
51 -
52 - uint32 n = 15; // Number of parallel samples
53 -
54 - // Logprobs configuration
55 - optional int32 logprobs = 22; // Number of log probabilities per output token (-1 for all)
56 - optional int32 prompt_logprobs = 23; // Number of log probabilities per prompt token (-1 for all)
57 -
58 - // Additional vLLM fields
59 - optional int32 seed = 24; // Random seed for reproducibility
60 - bool include_stop_str_in_output = 25; // Whether to include stop strings in output
61 - map<int32, float> logit_bias = 26; // Token ID to bias mapping (-100 to 100)
62 - optional int32 truncate_prompt_tokens = 27; // Prompt truncation (-1 for model max)
63 -
64 - // Structured outputs (one of) - matches vLLM's StructuredOutputsParams
65 - oneof constraint {
66 - string json_schema = 16; // JSON schema for structured output
67 - string regex = 17; // Regex pattern
68 - string grammar = 18; // Grammar/EBNF for structured output
69 - string structural_tag = 19; // Structural tag (e.g., Harmony models)
70 - bool json_object = 20; // Force JSON object output
71 - ChoiceConstraint choice = 21; // List of allowed choices
72 - }
73 -}
74 -
75 -// Choice constraint for structured outputs
76 -message ChoiceConstraint {
77 - repeated string choices = 1;
78 -}
79 -
80 -// Pre-tokenized input from Rust router
81 -message TokenizedInput {
82 - string original_text = 1; // For reference/debugging
83 - repeated uint32 input_ids = 2; // Actual token IDs to process
84 -}
85 -
86 -// =====================
87 -// Generate Request
88 -// =====================
89 -
90 -message GenerateRequest {
91 - string request_id = 1;
92 -
93 - // Prompt input
94 - oneof input {
95 - TokenizedInput tokenized = 2;
96 - string text = 3;
97 - }
98 -
99 - // Generation parameters (includes logprobs config)
100 - SamplingParams sampling_params = 4;
101 -
102 - // Streaming
103 - bool stream = 5;
104 -}
105 -
106 -// =====================
107 -// Generate Response
108 -// =====================
109 -
110 -message GenerateResponse {
111 - oneof response {
112 - GenerateStreamChunk chunk = 1; // For streaming
113 - GenerateComplete complete = 2; // For final/non-streaming
114 - }
115 -}
116 -
117 -message GenerateStreamChunk {
118 - repeated uint32 token_ids = 1; // Incremental tokens
119 - uint32 prompt_tokens = 2;
120 - uint32 completion_tokens = 3;
121 - uint32 cached_tokens = 4;
122 -
123 - // Logprobs support (TODO: implement in Phase 4)
124 - // OutputLogProbs output_logprobs = 5;
125 - // InputLogProbs input_logprobs = 6; // Only in first chunk
126 -}
127 -
128 -message GenerateComplete {
129 - repeated uint32 output_ids = 1; // All output tokens
130 - string finish_reason = 2; // "stop", "length", "abort"
131 - uint32 prompt_tokens = 3;
132 - uint32 completion_tokens = 4;
133 - uint32 cached_tokens = 5;
134 -
135 - // Logprobs support (TODO: implement in Phase 4)
136 - // OutputLogProbs output_logprobs = 6;
137 - // InputLogProbs input_logprobs = 7;
138 -}
139 -
140 -// =====================
141 -// Embedding Request
142 -// =====================
143 -
144 -message EmbedRequest {
145 - string request_id = 1;
146 - TokenizedInput tokenized = 2;
147 -}
148 -
149 -message EmbedResponse {
150 - repeated float embedding = 1;
151 - uint32 prompt_tokens = 2;
152 - uint32 embedding_dim = 3;
153 -}
154 -
155 -// =====================
156 -// Management Operations
157 -// =====================
158 -
159 -message HealthCheckRequest {}
160 -
161 -message HealthCheckResponse {
162 - bool healthy = 1;
163 - string message = 2;
164 -}
165 -
166 -message AbortRequest {
167 - repeated string request_ids = 1;
168 -}
169 -
170 -message AbortResponse {
171 -}
172 -
173 -// =====================
174 -// Model and Server Info
175 -// =====================
176 -
177 -message GetModelInfoRequest {}
178 -
179 -message GetModelInfoResponse {
180 - string model_path = 1;
181 - bool is_generation = 2;
182 - uint32 max_context_length = 3;
183 - uint32 vocab_size = 4;
184 - bool supports_vision = 5;
185 -}
186 -
187 -message GetServerInfoRequest {}
188 -
189 -message GetServerInfoResponse {
190 - uint32 active_requests = 1;
191 - bool is_paused = 2;
192 - double last_receive_timestamp = 3;
193 - double uptime_seconds = 4;
194 - string server_type = 5; // "vllm-grpc"
195 -}