@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 791b2fc30a
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 10, 12:45 UTC

[grpc] Support gRPC server entrypoint (#30190)

Signed-off-by: Chang Su <chang.s.su@oracle.com> Signed-off-by: njhill <nickhill123@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: njhill <nickhill123@gmail.com> Co-authored-by: Simon Mo <simon.mo@hey.com>

Chang Su committed Jan 7, 2026 at 23:24 UTC 791b2fc30a25cc48e06a6bd7ce4fd62d765ac004
12 files changed +1357 -6
.gitignore
+5
@@ -227,3 +227,8 @@ ep_kernels_workspace/
227
228 # Allow tracked library source folders under submodules (e.g., benchmarks/lib)
229 !vllm/benchmarks/lib/
230 +
231 +# Generated gRPC protobuf files (compiled at build time from vllm_engine.proto)
232 +vllm/grpc/vllm_engine_pb2.py
233 +vllm/grpc/vllm_engine_pb2_grpc.py
234 +vllm/grpc/vllm_engine_pb2.pyi
mkdocs.yaml
+3 -1
@@ -80,6 +80,7 @@ plugins:
80 - "re:vllm\\._.*" # Internal modules
81 - "vllm.third_party"
82 - "vllm.vllm_flash_attn"
83 + - "re:vllm\\.grpc\\..*_pb2.*" # Auto-generated protobuf files
84 - !ENV [API_AUTONAV_EXCLUDE, "re:^$"] # Match nothing by default
85 - mkdocstrings:
86 handlers:
@@ -87,7 +88,8 @@ plugins:
88 options:
89 show_symbol_type_heading: true
90 show_symbol_type_toc: true
90 - filters: []
91 + filters:
92 + - "!.*_pb2_grpc" # Exclude auto-generated gRPC stubs
93 summary:
94 modules: true
95 show_if_no_docstring: true
pyproject.toml
+5
@@ -9,6 +9,7 @@ requires = [
9 "torch == 2.9.1",
10 "wheel",
11 "jinja2",
12 + "grpcio-tools>=1.76.0",
13 ]
14 build-backend = "setuptools.build_meta"
15
@@ -55,6 +56,10 @@ include = ["vllm*"]
56 "vllm/third_party/**" = ["ALL"]
57 "vllm/version.py" = ["F401"]
58 "vllm/_version.py" = ["ALL"]
59 +# Exclude generated protobuf files
60 +"vllm/grpc/*_pb2.py" = ["ALL"]
61 +"vllm/grpc/*_pb2_grpc.py" = ["ALL"]
62 +"vllm/grpc/*_pb2.pyi" = ["ALL"]
63
64 [tool.ruff.lint]
65 select = [
requirements/build.txt
+2
@@ -9,3 +9,5 @@ wheel
9 jinja2>=3.1.6
10 regex
11 build
12 +protobuf>=6.33.2
13 +grpcio-tools>=1.76.0
requirements/common.txt
+3 -1
@@ -9,7 +9,7 @@ blake3
9 py-cpuinfo
10 transformers >= 4.56.0, < 5
11 tokenizers >= 0.21.1 # Required for fast incremental detokenization.
12 -protobuf # Required by LlamaTokenizer.
12 +protobuf >= 6.30.0 # Required by LlamaTokenizer, gRPC.
13 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
14 aiohttp
15 openai >= 1.99.1 # For Responses API with reasoning content
@@ -52,3 +52,5 @@ openai-harmony >= 0.0.3 # Required for gpt-oss
52 anthropic == 0.71.0
53 model-hosting-container-standards >= 0.1.10, < 1.0.0
54 mcp
55 +grpcio>=1.76.0
56 +grpcio-reflection>=1.76.0
requirements/test.txt
+3 -2
@@ -297,7 +297,7 @@ graphql-relay==3.2.0
297 # via graphene
298 greenlet==3.2.3
299 # via sqlalchemy
300 -grpcio==1.71.0
300 +grpcio==1.76.0
301 # via ray
302 gunicorn==23.0.0
303 # via mlflow
@@ -758,7 +758,7 @@ propcache==0.2.0
758 # yarl
759 proto-plus==1.26.1
760 # via google-api-core
761 -protobuf==5.28.3
761 +protobuf==6.33.2
762 # via
763 # google-api-core
764 # googleapis-common-protos
@@ -1249,6 +1249,7 @@ typing-extensions==4.15.0
1249 # chz
1250 # fastapi
1251 # graphene
1252 + # grpcio
1253 # huggingface-hub
1254 # librosa
1255 # lightning
setup.py
+71 -2
@@ -18,6 +18,7 @@ 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_scm import get_version
23 from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
24
@@ -79,6 +80,73 @@ def is_freethreaded():
80 return bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
81
82
83 +def compile_grpc_protos():
84 + """Compile gRPC protobuf definitions during build.
85 +
86 + This generates *_pb2.py, *_pb2_grpc.py, and *_pb2.pyi files from
87 + the vllm_engine.proto definition.
88 + """
89 + try:
90 + from grpc_tools import protoc
91 + except ImportError:
92 + logger.warning(
93 + "grpcio-tools not installed, skipping gRPC proto compilation. "
94 + "gRPC server functionality will not be available."
95 + )
96 + return False
97 +
98 + proto_file = ROOT_DIR / "vllm" / "grpc" / "vllm_engine.proto"
99 + if not proto_file.exists():
100 + logger.warning("Proto file not found at %s, skipping compilation", proto_file)
101 + return False
102 +
103 + logger.info("Compiling gRPC protobuf: %s", proto_file)
104 +
105 + result = protoc.main(
106 + [
107 + "grpc_tools.protoc",
108 + f"--proto_path={ROOT_DIR}",
109 + f"--python_out={ROOT_DIR}",
110 + f"--grpc_python_out={ROOT_DIR}",
111 + f"--pyi_out={ROOT_DIR}",
112 + str(proto_file),
113 + ]
114 + )
115 +
116 + if result != 0:
117 + logger.error("protoc failed with exit code %s", result)
118 + return False
119 +
120 + # Add SPDX headers and mypy ignore to generated files
121 + spdx_header = (
122 + "# SPDX-License-Identifier: Apache-2.0\n"
123 + "# SPDX-FileCopyrightText: Copyright contributors to the vLLM project\n"
124 + "# mypy: ignore-errors\n"
125 + )
126 +
127 + grpc_dir = ROOT_DIR / "vllm" / "grpc"
128 + for generated_file in [
129 + grpc_dir / "vllm_engine_pb2.py",
130 + grpc_dir / "vllm_engine_pb2_grpc.py",
131 + grpc_dir / "vllm_engine_pb2.pyi",
132 + ]:
133 + if generated_file.exists():
134 + content = generated_file.read_text()
135 + if not content.startswith("# SPDX-License-Identifier"):
136 + generated_file.write_text(spdx_header + content)
137 +
138 + logger.info("gRPC protobuf compilation successful")
139 + return True
140 +
141 +
142 +class BuildPyAndGenerateGrpc(build_py):
143 + """Build Python modules and generate gRPC stubs from proto files."""
144 +
145 + def run(self):
146 + compile_grpc_protos()
147 + super().run()
148 +
149 +
150 class CMakeExtension(Extension):
151 def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
152 super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa)
@@ -882,12 +950,13 @@ if _no_device():
950 ext_modules = []
951
952 if not ext_modules:
885 - cmdclass = {}
953 + cmdclass = {"build_py": BuildPyAndGenerateGrpc}
954 else:
955 cmdclass = {
956 "build_ext": precompiled_build_ext
957 if envs.VLLM_USE_PRECOMPILED
890 - else cmake_build_ext
958 + else cmake_build_ext,
959 + "build_py": BuildPyAndGenerateGrpc,
960 }
961
962 setup(
tests/entrypoints/test_grpc_server.py new
+428
@@ -0,0 +1,428 @@
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 = 30.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/grpc_server.py new
+531
@@ -0,0 +1,531 @@
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 +"""
6 +vLLM gRPC Server
7 +
8 +Starts a gRPC server for vLLM using the VllmEngine protocol.
9 +
10 +Usage:
11 + python -m vllm.entrypoints.grpc_server --model <model_path>
12 +
13 +Example:
14 + python -m vllm.entrypoints.grpc_server \
15 + --model meta-llama/Llama-2-7b-hf \
16 + --host 0.0.0.0 \
17 + --port 50051
18 +"""
19 +
20 +import argparse
21 +import asyncio
22 +import signal
23 +import sys
24 +import time
25 +from collections.abc import AsyncGenerator
26 +
27 +import grpc
28 +import uvloop
29 +from grpc_reflection.v1alpha import reflection
30 +
31 +from vllm import SamplingParams, TextPrompt, TokensPrompt
32 +from vllm.engine.arg_utils import AsyncEngineArgs
33 +from vllm.grpc import vllm_engine_pb2, vllm_engine_pb2_grpc
34 +from vllm.logger import init_logger
35 +from vllm.outputs import RequestOutput
36 +from vllm.sampling_params import RequestOutputKind, StructuredOutputsParams
37 +from vllm.usage.usage_lib import UsageContext
38 +from vllm.utils.argparse_utils import FlexibleArgumentParser
39 +from vllm.v1.engine.async_llm import AsyncLLM
40 +from vllm.version import __version__ as VLLM_VERSION
41 +
42 +logger = init_logger(__name__)
43 +
44 +
45 +class VllmEngineServicer(vllm_engine_pb2_grpc.VllmEngineServicer):
46 + """
47 + gRPC servicer implementing the VllmEngine service.
48 +
49 + Handles 6 RPCs:
50 + - Generate: Streaming text generation
51 + - Embed: Embeddings (TODO)
52 + - HealthCheck: Health probe
53 + - Abort: Cancel requests out-of-band
54 + - GetModelInfo: Model metadata
55 + - GetServerInfo: Server state
56 + """
57 +
58 + def __init__(self, async_llm: AsyncLLM, start_time: float):
59 + """
60 + Initialize the servicer.
61 +
62 + Args:
63 + async_llm: The AsyncLLM instance
64 + start_time: The server start time, in seconds since epoch
65 + """
66 + self.async_llm = async_llm
67 + self.start_time = start_time
68 + logger.info("VllmEngineServicer initialized")
69 +
70 + async def Generate(
71 + self,
72 + request: vllm_engine_pb2.GenerateRequest,
73 + context: grpc.aio.ServicerContext,
74 + ) -> AsyncGenerator[vllm_engine_pb2.GenerateResponse, None]:
75 + """
76 + Handle streaming generation requests.
77 +
78 + Args:
79 + request: The GenerateRequest protobuf
80 + context: gRPC context
81 +
82 + Yields:
83 + GenerateResponse protobuf messages (streaming)
84 + """
85 + request_id = request.request_id
86 + logger.debug("Generate request %s received.", request_id)
87 +
88 + try:
89 + # Extract tokenized input
90 + if request.WhichOneof("input") == "tokenized":
91 + prompt: TokensPrompt = {
92 + "prompt_token_ids": list(request.tokenized.input_ids)
93 + }
94 + if request.tokenized.original_text:
95 + prompt["prompt"] = request.tokenized.original_text
96 + else:
97 + prompt: TextPrompt = {"prompt": request.text}
98 +
99 + # Build sampling params with detokenize=False
100 + sampling_params = self._sampling_params_from_proto(
101 + request.sampling_params, stream=request.stream
102 + )
103 +
104 + async for output in self.async_llm.generate(
105 + prompt=prompt,
106 + sampling_params=sampling_params,
107 + request_id=request_id,
108 + ):
109 + # Convert vLLM output to protobuf
110 + # For streaming, always send chunks
111 + if request.stream:
112 + yield self._chunk_response(output)
113 +
114 + # Send complete response when finished
115 + if output.finished:
116 + yield self._complete_response(output)
117 +
118 + except ValueError as e:
119 + # Invalid request error (equiv to 400).
120 + await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e))
121 + except Exception as e:
122 + logger.exception("Error in Generate for request %s", request_id)
123 + await context.abort(grpc.StatusCode.INTERNAL, str(e))
124 +
125 + async def Embed(
126 + self,
127 + request: vllm_engine_pb2.EmbedRequest,
128 + context: grpc.aio.ServicerContext,
129 + ) -> vllm_engine_pb2.EmbedResponse:
130 + """
131 + Handle embedding requests.
132 +
133 + TODO: Implement in Phase 4
134 +
135 + Args:
136 + request: The EmbedRequest protobuf
137 + context: gRPC context
138 +
139 + Returns:
140 + EmbedResponse protobuf
141 + """
142 + logger.warning("Embed RPC not yet implemented")
143 + await context.abort(
144 + grpc.StatusCode.UNIMPLEMENTED, "Embed RPC not yet implemented"
145 + )
146 +
147 + async def HealthCheck(
148 + self,
149 + request: vllm_engine_pb2.HealthCheckRequest,
150 + context: grpc.aio.ServicerContext,
151 + ) -> vllm_engine_pb2.HealthCheckResponse:
152 + """
153 + Handle health check requests.
154 +
155 + Args:
156 + request: The HealthCheckRequest protobuf
157 + context: gRPC context
158 +
159 + Returns:
160 + HealthCheckResponse protobuf
161 + """
162 + is_healthy = not self.async_llm.errored
163 + message = "Health" if is_healthy else "Engine is not alive"
164 +
165 + logger.debug("HealthCheck request: healthy=%s, message=%s", is_healthy, message)
166 +
167 + return vllm_engine_pb2.HealthCheckResponse(healthy=is_healthy, message=message)
168 +
169 + async def Abort(
170 + self,
171 + request: vllm_engine_pb2.AbortRequest,
172 + context: grpc.aio.ServicerContext,
173 + ) -> vllm_engine_pb2.AbortResponse:
174 + """
175 + Out-of-band abort requests.
176 +
177 + Args:
178 + request: The AbortRequest protobuf
179 + context: gRPC context
180 +
181 + Returns:
182 + AbortResponse protobuf
183 + """
184 + request_ids = request.request_ids
185 + logger.debug("Abort requests: %s", request_ids)
186 +
187 + await self.async_llm.abort(request_ids)
188 + return vllm_engine_pb2.AbortResponse()
189 +
190 + async def GetModelInfo(
191 + self,
192 + request: vllm_engine_pb2.GetModelInfoRequest,
193 + context: grpc.aio.ServicerContext,
194 + ) -> vllm_engine_pb2.GetModelInfoResponse:
195 + """
196 + Handle model info requests.
197 +
198 + Args:
199 + request: The GetModelInfoRequest protobuf
200 + context: gRPC context
201 +
202 + Returns:
203 + GetModelInfoResponse protobuf
204 + """
205 + model_config = self.async_llm.model_config
206 +
207 + return vllm_engine_pb2.GetModelInfoResponse(
208 + model_path=model_config.model,
209 + is_generation=model_config.runner_type == "generate",
210 + max_context_length=model_config.max_model_len,
211 + vocab_size=model_config.get_vocab_size(),
212 + supports_vision=model_config.is_multimodal_model,
213 + )
214 +
215 + async def GetServerInfo(
216 + self,
217 + request: vllm_engine_pb2.GetServerInfoRequest,
218 + context: grpc.aio.ServicerContext,
219 + ) -> vllm_engine_pb2.GetServerInfoResponse:
220 + """
221 + Handle server info requests.
222 +
223 + Args:
224 + request: The GetServerInfoRequest protobuf
225 + context: gRPC context
226 +
227 + Returns:
228 + GetServerInfoResponse protobuf
229 + """
230 + num_requests = self.async_llm.output_processor.get_num_unfinished_requests()
231 +
232 + return vllm_engine_pb2.GetServerInfoResponse(
233 + active_requests=num_requests,
234 + is_paused=False, # TODO
235 + last_receive_timestamp=time.time(), # TODO looks wrong?
236 + uptime_seconds=time.time() - self.start_time,
237 + server_type="vllm-grpc",
238 + )
239 +
240 + # ========== Helper methods ==========
241 +
242 + @staticmethod
243 + def _sampling_params_from_proto(
244 + params: vllm_engine_pb2.SamplingParams, stream: bool = True
245 + ) -> SamplingParams:
246 + """
247 + Convert protobuf SamplingParams to vLLM SamplingParams.
248 +
249 + Args:
250 + params: Protobuf SamplingParams message
251 + stream: Whether streaming is enabled
252 +
253 + Returns:
254 + vLLM SamplingParams with detokenize=False and structured_outputs
255 + """
256 + # Build stop sequences
257 + stop = list(params.stop) if params.stop else None
258 + stop_token_ids = list(params.stop_token_ids) if params.stop_token_ids else None
259 +
260 + # Handle structured outputs constraints
261 + structured_outputs = None
262 + constraint_field = params.WhichOneof("constraint")
263 + if constraint_field:
264 + if constraint_field == "json_schema":
265 + structured_outputs = StructuredOutputsParams(json=params.json_schema)
266 + elif constraint_field == "regex":
267 + structured_outputs = StructuredOutputsParams(regex=params.regex)
268 + elif constraint_field == "grammar":
269 + structured_outputs = StructuredOutputsParams(grammar=params.grammar)
270 + elif constraint_field == "structural_tag":
271 + structured_outputs = StructuredOutputsParams(
272 + structural_tag=params.structural_tag
273 + )
274 + elif constraint_field == "json_object":
275 + structured_outputs = StructuredOutputsParams(
276 + json_object=params.json_object
277 + )
278 + elif constraint_field == "choice":
279 + structured_outputs = StructuredOutputsParams(
280 + choice=list(params.choice.choices)
281 + )
282 +
283 + # Create SamplingParams
284 + # output_kind=DELTA: Return only new tokens in each chunk (for streaming)
285 + return SamplingParams(
286 + temperature=params.temperature if params.HasField("temperature") else 1.0,
287 + top_p=params.top_p if params.top_p != 0.0 else 1.0,
288 + top_k=params.top_k,
289 + min_p=params.min_p,
290 + frequency_penalty=params.frequency_penalty,
291 + presence_penalty=params.presence_penalty,
292 + repetition_penalty=params.repetition_penalty
293 + if params.repetition_penalty != 0.0
294 + else 1.0,
295 + max_tokens=params.max_tokens if params.HasField("max_tokens") else None,
296 + min_tokens=params.min_tokens,
297 + stop=stop,
298 + stop_token_ids=stop_token_ids,
299 + skip_special_tokens=params.skip_special_tokens,
300 + spaces_between_special_tokens=params.spaces_between_special_tokens,
301 + ignore_eos=params.ignore_eos,
302 + n=params.n if params.n > 0 else 1,
303 + logprobs=params.logprobs if params.HasField("logprobs") else None,
304 + prompt_logprobs=params.prompt_logprobs
305 + if params.HasField("prompt_logprobs")
306 + else None,
307 + seed=params.seed if params.HasField("seed") else None,
308 + include_stop_str_in_output=params.include_stop_str_in_output,
309 + logit_bias=dict(params.logit_bias) if params.logit_bias else None,
310 + truncate_prompt_tokens=params.truncate_prompt_tokens
311 + if params.HasField("truncate_prompt_tokens")
312 + else None,
313 + structured_outputs=structured_outputs,
314 + # detokenize must be True if stop strings are used
315 + detokenize=bool(stop),
316 + output_kind=RequestOutputKind.DELTA
317 + if stream
318 + else RequestOutputKind.FINAL_ONLY,
319 + )
320 +
321 + @staticmethod
322 + def _chunk_response(output: RequestOutput) -> vllm_engine_pb2.GenerateResponse:
323 + """
324 + Build a streaming chunk response from vLLM output.
325 + When output_kind=DELTA, vLLM returns only new tokens automatically.
326 +
327 + Args:
328 + output: vLLM RequestOutput (with delta tokens when output_kind=DELTA)
329 +
330 + Returns:
331 + GenerateResponse with chunk field set
332 + """
333 + # Get the completion output (first one if n > 1)
334 + completion = output.outputs[0] if output.outputs else None
335 +
336 + if completion is None:
337 + # Empty chunk
338 + return vllm_engine_pb2.GenerateResponse(
339 + chunk=vllm_engine_pb2.GenerateStreamChunk(
340 + token_ids=[],
341 + prompt_tokens=0,
342 + completion_tokens=0,
343 + cached_tokens=0,
344 + ),
345 + )
346 +
347 + # When output_kind=DELTA, completion.token_ids contains only new tokens
348 + # vLLM handles the delta logic internally
349 + # completion_tokens = delta count (client will accumulate)
350 + return vllm_engine_pb2.GenerateResponse(
351 + chunk=vllm_engine_pb2.GenerateStreamChunk(
352 + token_ids=completion.token_ids,
353 + prompt_tokens=len(output.prompt_token_ids)
354 + if output.prompt_token_ids
355 + else 0,
356 + completion_tokens=len(completion.token_ids), # Delta count
357 + cached_tokens=output.num_cached_tokens,
358 + ),
359 + )
360 +
361 + @staticmethod
362 + def _complete_response(output: RequestOutput) -> vllm_engine_pb2.GenerateResponse:
363 + """
364 + Build a final completion response from vLLM output.
365 +
366 + Args:
367 + output: vLLM RequestOutput (finished=True)
368 +
369 + Returns:
370 + GenerateResponse with complete field set
371 + """
372 + # Get the completion output (first one if n > 1)
373 + completion = output.outputs[0] if output.outputs else None
374 +
375 + if completion is None:
376 + # Empty completion
377 + return vllm_engine_pb2.GenerateResponse(
378 + complete=vllm_engine_pb2.GenerateComplete(
379 + output_ids=[],
380 + finish_reason="error",
381 + prompt_tokens=0,
382 + completion_tokens=0,
383 + cached_tokens=0,
384 + ),
385 + )
386 +
387 + # Build complete response
388 + # When streaming (DELTA mode): completion.token_ids will be empty/last delta
389 + # When non-streaming (FINAL_ONLY mode): completion.token_ids has all tokens
390 + # Client will accumulate token counts for streaming
391 + return vllm_engine_pb2.GenerateResponse(
392 + complete=vllm_engine_pb2.GenerateComplete(
393 + output_ids=completion.token_ids,
394 + finish_reason=completion.finish_reason or "stop",
395 + prompt_tokens=len(output.prompt_token_ids)
396 + if output.prompt_token_ids
397 + else 0,
398 + completion_tokens=len(completion.token_ids),
399 + cached_tokens=output.num_cached_tokens,
400 + ),
401 + )
402 +
403 +
404 +async def serve_grpc(args: argparse.Namespace):
405 + """
406 + Main serving function.
407 +
408 + Args:
409 + args: Parsed command line arguments
410 + """
411 + logger.info("vLLM gRPC server version %s", VLLM_VERSION)
412 + logger.info("args: %s", args)
413 +
414 + start_time = time.time()
415 +
416 + # Create engine args
417 + engine_args = AsyncEngineArgs.from_cli_args(args)
418 +
419 + # Build vLLM config
420 + vllm_config = engine_args.create_engine_config(
421 + usage_context=UsageContext.OPENAI_API_SERVER
422 + )
423 +
424 + # Create AsyncLLM
425 + async_llm = AsyncLLM.from_vllm_config(
426 + vllm_config=vllm_config,
427 + usage_context=UsageContext.OPENAI_API_SERVER,
428 + enable_log_requests=args.enable_log_requests,
429 + disable_log_stats=args.disable_log_stats_server,
430 + )
431 +
432 + # Create servicer
433 + servicer = VllmEngineServicer(async_llm, start_time)
434 +
435 + # Create gRPC server
436 + server = grpc.aio.server(
437 + options=[
438 + ("grpc.max_send_message_length", -1),
439 + ("grpc.max_receive_message_length", -1),
440 + ],
441 + )
442 +
443 + # Add servicer to server
444 + vllm_engine_pb2_grpc.add_VllmEngineServicer_to_server(servicer, server)
445 +
446 + # Enable reflection for grpcurl and other tools
447 + service_names = (
448 + vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name,
449 + reflection.SERVICE_NAME,
450 + )
451 + reflection.enable_server_reflection(service_names, server)
452 +
453 + # Bind to address
454 + address = f"{args.host}:{args.port}"
455 + server.add_insecure_port(address)
456 +
457 + # Start server
458 + await server.start()
459 + logger.info("vLLM gRPC server started on %s", address)
460 + logger.info("Server is ready to accept requests")
461 +
462 + # Handle shutdown signals
463 + loop = asyncio.get_running_loop()
464 + stop_event = asyncio.Event()
465 +
466 + def signal_handler():
467 + logger.info("Received shutdown signal")
468 + stop_event.set()
469 +
470 + for sig in (signal.SIGTERM, signal.SIGINT):
471 + loop.add_signal_handler(sig, signal_handler)
472 +
473 + # Serve until shutdown signal
474 + try:
475 + await stop_event.wait()
476 + except KeyboardInterrupt:
477 + logger.info("Interrupted by user")
478 + finally:
479 + logger.info("Shutting down vLLM gRPC server...")
480 +
481 + # Stop gRPC server
482 + await server.stop(grace=5.0)
483 + logger.info("gRPC server stopped")
484 +
485 + # Shutdown AsyncLLM
486 + async_llm.shutdown()
487 + logger.info("AsyncLLM engine stopped")
488 +
489 + logger.info("Shutdown complete")
490 +
491 +
492 +def main():
493 + """Main entry point."""
494 + parser = FlexibleArgumentParser(
495 + description="vLLM gRPC Server",
496 + )
497 +
498 + # Server args
499 + parser.add_argument(
500 + "--host",
501 + type=str,
502 + default="0.0.0.0",
503 + help="Host to bind gRPC server to",
504 + )
505 + parser.add_argument(
506 + "--port",
507 + type=int,
508 + default=50051,
509 + help="Port to bind gRPC server to",
510 + )
511 + parser.add_argument(
512 + "--disable-log-stats-server",
513 + action="store_true",
514 + help="Disable stats logging on server side",
515 + )
516 +
517 + # Add vLLM engine args
518 + parser = AsyncEngineArgs.add_cli_args(parser)
519 +
520 + args = parser.parse_args()
521 +
522 + # Run server
523 + try:
524 + uvloop.run(serve_grpc(args))
525 + except Exception as e:
526 + logger.exception("Server failed: %s", e)
527 + sys.exit(1)
528 +
529 +
530 +if __name__ == "__main__":
531 + main()
vllm/grpc/__init__.py new
+17
@@ -0,0 +1,17 @@
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 new
+94
@@ -0,0 +1,94 @@
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 new
+195
@@ -0,0 +1,195 @@
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 +}