[Feature] gRPC Engine-Client Backend PR1: gRPC Engine Service & Standalone Worker Daemon - #250
[Feature] gRPC Engine-Client Backend PR1: gRPC Engine Service & Standalone Worker Daemon#250herotai214 wants to merge 4 commits into
Conversation
…Python worker daemon (PR 1) Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: herotai214 <herotai214@gmail.com>
…s gRPC worker daemon Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: herotai214 <herotai214@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 95b7bef0a5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| Stream generated tokens for incoming GenerateRequest over gRPC HTTP/2. | ||
| Supports pre-tokenized prompt_token_ids and cancellation propagation. | ||
| """ | ||
| from vllm import TokensPrompt |
There was a problem hiding this comment.
In the checked .buildkite/pipeline.yml Python Tests environment, the Rust image installs only .[dev], which does not include vLLM, but every mock generation imports vllm.TokensPrompt here before invoking MockAsyncEngine; _build_sampling_params and the mock generator also import vLLM classes. Consequently the newly added mock generation tests fail with ModuleNotFoundError and --mock-engine cannot provide the advertised GPU-free, vLLM-free mode. Use lightweight mock-native objects or otherwise avoid all vLLM imports when the mock engine is selected.
Useful? React with 👍 / 👎.
| "grpcio>=1.60.0", | ||
| "protobuf>=5.0.0", |
There was a problem hiding this comment.
These constraints permit environments that the checked-in generated modules explicitly reject: engine_client_pb2_grpc.py:10-29 requires grpcio 1.83.1 or newer, while engine_client_pb2.py:13-15 validates against protobuf 7.35.1. An installation constrained to allowed grpcio 1.60–1.82 or an older protobuf runtime succeeds dependency resolution but then fails while importing vllm_router.proto, making the worker entry point unusable. Regenerate with the declared minimum toolchain or require the generated versions.
Useful? React with 👍 / 👎.
| engine_args_kwargs = { | ||
| "model": args.model, | ||
| "tensor_parallel_size": args.tensor_parallel_size, | ||
| "gpu_memory_utilization": args.gpu_memory_utilization, | ||
| "trust_remote_code": args.trust_remote_code, |
There was a problem hiding this comment.
When the daemon is launched with --dp-size N for N > 1, that value is only copied into the discovery response; it is absent from the AsyncEngineArgs configuration, so the engine retains its default data-parallel size while GetModelInfo advertises N replicas. Multi-DP deployments therefore start the wrong topology and expose misleading routing metadata.
Useful? React with 👍 / 👎.
| generator = self.engine.generate( | ||
| prompt=prompt, | ||
| sampling_params=sampling_params, | ||
| request_id=request_id, | ||
| ) |
There was a problem hiding this comment.
For requests with a nonzero dp_rank, PREFILL_ONLY/DECODE_ONLY execution mode, or multimodal_data, this call forwards only the text/tokens, sampling parameters, and request ID, silently treating the request as an ordinary text generation. That can send work to the wrong DP rank, perform a full decode for a prefill-only request, or return an answer that ignores image/audio input; these fields should be translated into engine inputs or rejected instead of being accepted and discarded.
Useful? React with 👍 / 👎.
| delta_token_id = 0 | ||
| if len(current_token_ids) > prev_token_count: | ||
| delta_token_id = current_token_ids[-1] | ||
| prev_token_count = len(current_token_ids) |
There was a problem hiding this comment.
If a RequestOutput adds more than one token since the previous yield, as can occur with batched, multi-step, or speculative output, this code emits only current_token_ids[-1] while text_delta contains the text for all newly accepted tokens. The streaming token sequence and unary output_token_ids then lose tokens despite the output text being complete; process the entire slice beginning at prev_token_count rather than retaining only its last element.
Useful? React with 👍 / 👎.
| self._running_requests: int = 0 | ||
| self._waiting_requests: int = 0 |
There was a problem hiding this comment.
_waiting_requests is initialized to zero and never updated anywhere in the daemon, so every streamed WorkerMetrics claims the queue is empty even while requests are waiting inside vLLM. A router using this newly exposed field for load balancing will continue directing traffic toward an overloaded worker; derive it from actual engine statistics or omit the unsupported metric.
Useful? React with 👍 / 👎.
| text_match = http_text == grpc_text | ||
| finish_match = http_finish == grpc_finish |
There was a problem hiding this comment.
The differential harness advertises token-by-token parity, but its verdict checks only decoded text and finish reason; grpc_tokens is never compared with any HTTP-derived token sequence. A regression that drops or changes gRPC token IDs—while preserving the decoded text—therefore still reports complete parity, so the test should obtain or reconstruct the HTTP token IDs and assert equality.
Useful? React with 👍 / 👎.
…ize, and validations - Decouple mock path from vLLM imports so CPU unit tests run in zero-GPU / zero-vLLM CI environments - Emit all token IDs across multi-token steps (multi-step / speculative decoding) to prevent dropped tokens - Forward requested dp_size to AsyncEngineArgs if supported, or log operational guidance - Validate dp_rank, execution_mode, and multimodal_data in GenerateStream with proper gRPC status codes - Dynamically track waiting requests and query engine statistics for worker metrics - Verify token sequence equivalence against tokenizer in differential parity test harness - Relax generated protobuf/grpcio runtime version check to match pyproject.toml constraints Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: herotai214 <herotai214@gmail.com>
There was a problem hiding this comment.
🟡 Changes recommended
There are correctness-breaking issues (generated stub/runtime version constraints vs declared dependencies, unary request_id mismatch, and an API contract ambiguity around token_id==0) that can cause runtime failures or incorrect semantics.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR introduces a new Python-side gRPC “worker daemon” backend for vllm-router, along with the protobuf service contract, generated Python stubs, and verification harnesses (mock-unit, real-GPU, and HTTP-vs-gRPC differential parity).
Changes:
- Add
EngineServiceprotobuf contract (proto/engine_client.proto) plus generated Python gRPC/protobuf stubs underpy_src/vllm_router/proto/. - Add a standalone
grpc.aioworker daemon (py_src/vllm_router/worker_daemon.py) that wrapsvllm.AsyncLLMEngine(or a mock engine). - Add test harnesses and wiring: unit tests, real-GPU integration script, differential parity script, and new Python deps/script entrypoint.
File summaries
| File | Description |
|---|---|
| scripts/test_worker_daemon_real_gpu.py | Real-GPU end-to-end smoke test for the worker daemon. |
| scripts/differential_test_grpc_vs_http.py | Differential parity harness comparing stock HTTP server vs gRPC daemon. |
| pyproject.toml | Adds gRPC/protobuf deps and a vllm-worker-daemon entrypoint. |
| py_test/test_worker_daemon_grpc.py | Pytest suite for gRPC RPCs using the mock engine + in-process server. |
| py_src/vllm_router/worker_daemon.py | Implements the gRPC service, mock engine, CLI, and server lifecycle. |
| py_src/vllm_router/proto/engine_client_pb2.py | Generated protobuf message classes. |
| py_src/vllm_router/proto/engine_client_pb2_grpc.py | Generated gRPC stub/servicer scaffolding with runtime version check. |
| py_src/vllm_router/proto/init.py | Proto package init and exports. |
| py_src/vllm_router/init.py | Makes Router export conditional when Rust extension isn’t available. |
| proto/engine_client.proto | Defines the gRPC EngineService API and message schemas. |
| examples/simulate_consistent_hash.rs | Minor ring iteration tweak in the example. |
Review details
Files not reviewed (1)
- py_src/vllm_router/proto/engine_client_pb2.py: Generated file
Suppressed comments (2)
py_src/vllm_router/worker_daemon.py:210
Generate()generates its ownrequest_idwhen the incoming request omits one, butGenerateStream()separately generates a (different) ID. This can cause the unary responserequest_idto disagree with the ID actually used for engine generation and streamed chunks.
"""Unary generation returning full output in a single response."""
accumulated_text = []
accumulated_token_ids = []
finish_reason = ""
request_id = request.request_id or f"req-{uuid.uuid4().hex[:12]}"
async for chunk in self.GenerateStream(request, context):
if chunk.text_delta:
accumulated_text.append(chunk.text_delta)
if chunk.token_id > 0:
accumulated_token_ids.append(chunk.token_id)
if chunk.is_finished:
finish_reason = chunk.finish_reason
return engine_client_pb2.GenerateResponse(
request_id=request_id,
py_src/vllm_router/worker_daemon.py:350
--enable-prefix-cachingis declared withaction="store_true"anddefault=True, which makes it impossible to disable (the value is always True regardless of whether the flag is passed). If the intent is to make prefix caching configurable, this needs a--disable-prefix-caching(store_true) or a paired enable/disable flag setup.
parser.add_argument(
"--enable-prefix-caching",
action="store_true",
default=True,
help="Enable automatic prefix caching",
)
- Files reviewed: 10/11 changed files
- Comments generated: 5
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| "requests>=2.25.0", | ||
| "grpcio>=1.60.0", | ||
| "protobuf>=5.0.0", | ||
| ] |
| import os | ||
| import sys | ||
|
|
||
| _proto_dir = os.path.dirname(os.path.abspath(__file__)) | ||
| if _proto_dir not in sys.path: | ||
| sys.path.insert(0, _proto_dir) | ||
|
|
||
| from . import engine_client_pb2 as engine_client_pb2 # noqa: E402 | ||
| from . import engine_client_pb2_grpc as engine_client_pb2_grpc # noqa: E402 | ||
|
|
||
| __all__ = ["engine_client_pb2", "engine_client_pb2_grpc"] |
| # Compute deltas | ||
| text_delta = current_text[len(prev_text) :] | ||
| prev_text = current_text | ||
|
|
||
| delta_token_id = 0 | ||
| if len(current_token_ids) > prev_token_count: | ||
| delta_token_id = current_token_ids[-1] | ||
| prev_token_count = len(current_token_ids) | ||
|
|
| async for chunk in stream: | ||
| if chunk.text_delta: | ||
| text_chunks.append(chunk.text_delta) | ||
| if chunk.token_id > 0: | ||
| token_ids.append(chunk.token_id) | ||
| if chunk.is_finished: |
| dev = [ | ||
| "pytest>=7.0.0", | ||
| "pytest-asyncio>=0.21.0", | ||
| "pytest-cov>=4.0.0", | ||
| "grpcio-tools>=1.60.0", | ||
| ] |
Signed-off-by: herotai214 <herotai214@gmail.com>
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
[Feature] gRPC Engine-Client Backend PR1: gRPC Engine Service & Standalone Worker Daemon
PR1 of #249
Multi-PR plan and lines of change
Five PRs, each reviewable on its own. This PR is PR 1 only.
VllmEngineServicerproto/vllm_engine.proto, in-tree Servicer wrappingAsyncLLMEngine, stubs generated atpip install, mock CPU tests + optional GPU / differential harnesses. No Rust router client.BackendClient, e2e streamingGrpcEngineBackend, router tokenize + minijinja, OpenAI SSE out. First timePOST /v1/chat/completionshits this Servicer.dp_rankrouting, cache-aware dispatch./v1/models, HTTP/reset_prefix_cache, andStartProfile/StopProfileonce--profiler-configexists.PR 1 measured vs
main(ab7e453): +1894 / −6 across 12 files. Breakdown:py_src/vllm_router/vllm_servicer.pyproto/vllm_engine.protosetup.pypy_src/vllm_router/proto/__init__.pypyproject.tomlgrpcio,protobuf,vllm-servicer.gitignore/MANIFEST.in/ package__init__.py*_pb2.pypy_test/test_vllm_servicer.pypy_test/e2e/test_differential_grpc_vs_http.pypy_test/e2e/test_vllm_servicer_real_gpu.pyGenerated
vllm_engine_pb2.py/vllm_engine_pb2_grpc.pyare not in that count (not checked in).Summary
This PR adds an in-tree gRPC Engine Servicer for vLLM Router: a Python process that wraps
vllm.AsyncLLMEngineand speaks Protobuf over HTTP/2.It is only the engine side. The Rust router still talks HTTP/JSON to
vllm serve. Nothing in this PR sends a request from the router to the Servicer, or streams tokens from the Servicer back through the router.What landed
proto/vllm_engine.proto—package vllm.engine.v1, serviceVllmEnginevllm_router.vllm_servicer.VllmEngineServicer— wrapsAsyncLLMEngine(or a self-contained mock)vllm-servicer/python -m vllm_router.vllm_servicervllm_engine_pb2.py,vllm_engine_pb2_grpc.py) generated atpip install/ wheel build, not checked inStructure (what this PR is, and what it is not)
Target end state (PR 1 + PR 2):
This PR implements only the right box. Router → Servicer and Servicer → Router are not done.
grpcurl→ Servicer →AsyncLLMEnginePublic HTTP routes (
/v1/chat/completions, …) are unchanged. They still hit stock HTTP backends.Design
Today the router is an HTTP load balancer in front of
vllm serve. Every token pays FastAPI + JSON SSE + HTTP/1.1 on the worker, then JSON parse again in Rust. PR 2 will replace that internal hop with the gRPC link above. PR 1 only proves the Servicer is a correct engine front.Rules we locked:
vllm serve --grpcorsmg-grpc-servicer.prompt_token_ids→vllm.TokensPrompt. The Servicer does not run the chat template.prompt_textis a fallback for tests and debugging.engine.abort(request_id). There is no public OpenAI/v1/abort, so we do not add anAbortRPC in this PR.vllm_servicer.py, classVllmEngineServicer, CLIvllm-servicer. Not “daemon”, not a genericEngineService.What is in this PR
proto/vllm_engine.protopy_src/vllm_router/vllm_servicer.pyVllmEngineServicer+MockAsyncEngine+ CLIsetup.pygenerate_grpc_stubs()protocat install; relative imports; grpcio/protobuf floorsvllm-servicerentry pointpy_test/test_vllm_servicer.pypy_test/e2e/test_vllm_servicer_real_gpu.pypy_test/e2e/test_differential_grpc_vs_http.pyMissing vLLM and no
--mock-engine→ log +SystemExit. The mock is explicit (is_mock), not a silent fake vLLM.Prefix caching defaults on (
--no-enable-prefix-cachingto turn it off).Generated stubs: the two
pb2filesThe
.protois the source of truth. Python bindings are not checked in.setup.pycompiles them when you install the package.grpcio-tools>=1.60.0is in[build-system] requires, so a PEP 517 install gets a compiler. Afterprotoc,setup.pypatches the output:_grpcimport package-relative (from . import vllm_engine_pb2)GRPC_GENERATED_VERSION = "1.60.0"to the runtime floor inpyproject.tomlprotobufdoes not hard-crash importwarningsimportBoth files are in
.gitignore. A fresh clone does not contain them.What each file is for
vllm_engine_pb2.pyGenerateRequest,SamplingParams,GenerateStreamResponse, …vllm_engine_pb2_grpc.pyVllmEnginevllm_engine_pb2_grpc.pydefines three classes. Only two matter:VllmEngineStubchannel.unary_stream(...)/unary_unary(...).VllmEngineServicerUNIMPLEMENTED+NotImplementedError. Compiler output, not a TODO list.vllm_router.vllm_servicer.VllmEngineServicersubclasses it and overrides the five RPCs. The dummy methods never run.VllmEngineThis class is part of an EXPERIMENTAL API.Each method is a@staticmethodthat callsgrpc.experimental.unary_stream/unary_unarywith a target string, instead of aChannel.add_VllmEngineServicer_to_server(...)is the registration helper. We call that inserve().Import path after install:
If the two files are missing,
vllm_router/proto/__init__.pyraises a clearImportErrortelling you to re-install.When to install again
Re-run
pip install -e .from the router repo root when:proto/vllm_engine.proto— oldpb2files will not match the new RPC/field numbers.setup.py, orgrpcio-toolswas absent).git clean, wipedpy_src/vllm_router/proto/vllm_engine_pb2*.py).You do not need to reinstall after editing
vllm_servicer.pyor tests. Those import the already-generated modules.Supported command:
pip install -e .pip install -e .is the path reviewers and CI should use. A one-offfrom setup import generate_grpc_stubsalso writes the files (as a side effect of importingsetup.py) but is not the supported workflow.Intentionally not in this PR
BackendClient,grpc://workersPOST /v1/chat/completions→ gRPCGetTokenizerdp_rankrouting / prefix-cache aware dispatchdp_rankis rejectedExecutionMode(PREFILL_ONLY/DECODE_ONLY)UNIMPLEMENTEDUNIMPLEMENTEDStartProfile/StopProfile--profiler-config, which this CLI does not expose. Shipping the RPCs now would besuccess=Falseon a default launch/start_profile/stop_profilebroadcast + a real trace test)/v1/modelslive probe, cluster admin fan-outAbortRPCengine.abortPR 2 milestone: client → router HTTP → Rust gRPC → this Servicer → tokens back as OpenAI SSE.
The proto: what is in it, and why
Wire path:
/vllm.engine.v1.VllmEngine/GenerateStream.package vllm.engine.v1is a protobuf namespace, not a Python import of vLLM.Five RPCs
GenerateStream(core). Token-first generate. Each chunk can carryoptional uint32 token_id,text_delta,is_finished,finish_reason, andWorkerMetrics(running / waiting / KV %).optionalmatters: token id0is a real token, not “unset”.Generate(unary). Same request, one complete response (output_token_ids,output_text,finish_reason). Useful for tests and non-streaming callers. Samerequest_idspace as the stream.GetModelInfo. Model name,max_model_len,dp_size,block_size. PR 2 uses this for discovery.stop_tokensis reserved; we do not invent a stop list here.HealthCheck.SERVING/NOT_SERVING. Also registergrpc.health.v1whengrpcio-health-checkingis installed.ResetPrefixCache. Forwards toengine.reset_prefix_cache(). Prefix caching is on by default, and a live GPU run already returnedsuccess=True. This is a real admin call, not a stub.Request / response choices
prompt_token_idsprompt_textTokenizedInput.original_text(IDs plus a string for logging/renderer).SamplingParamsvllm.SamplingParams(temperature, top_p, top_k, max_tokens, stops, penalties, seed, ignore_eos).optional seedWorkerMetricson every chunkGetLoadsRPC on the hot path. Best-effort; counters if the engine has no stats.optional token_iddp_rank< dp_size. PR 1 does not implement DP routing.ExecutionMode/MultimodalItemNORMALand any MM payload.What the proto is not
It is not a clone of SMG’s proto. It is not
vllm serve’s HTTP API. It is the smallest contract that lets a router send tokens in and get tokens + finish + light load signals out.Diff vs SMG (
smg-grpc-servicer/vllm serve --grpc)SMG already runs this kind of architecture in production. We are not wrapping their package. We share the idea (Rust gateway, Python engine, Protobuf) and a few names (
VllmEngine,VllmEngineServicer,vllm_engine.proto). The contract and the process are ours.vllm-project/router)smg-grpc-servicer)vllm serve --grpcis a convenience flagvllm.engine.v1vllm.grpc.enginevllm-servicer/python -m vllm_router.vllm_servicervllm serve --grpcor their servicerGenerateStream+ unaryGenerateGenerate+streambool; response isoneof { chunk, complete }prompt_token_ids+ optionalprompt_textoneof input { TokenizedInput tokenized; string text }ResetPrefixCacheonlyEmbed,Abort,GetServerInfo,GetLoads,GetTokenizer,SubscribeKvEventsGetLoadsengine.abortAbortRPC (list of ids)GetTokenizerstreams a zip of tokenizer files--profiler-configexistsWe stayed smaller on purpose: one repo, one process, five RPCs we can test, no PyPI pin on their servicer, no requirement to change vLLM.
Tests: what, why, expected behavior
CI runs
pytest py_test/ -v --ignore=py_test/e2e. GPU tests never run on CPU.A.
py_test/test_vllm_servicer.py(CPU, mock)In-process
grpc.aioserver +MockAsyncEngine. No GPU, no real weights.test_health_checkstatus == SERVINGtest_get_model_infomax_model_len/dp_size/block_sizematch fixturetest_generate_stream_prompt_token_idsis_finishedtest_generate_stream_text_fallbackprompt_textstill workstest_generate_unarytest_reset_prefix_cachesuccess == Truetest_stream_cancellationtest_unsupported_execution_mode_rejectedPREFILL_ONLY/DECODE_ONLY→UNIMPLEMENTEDtest_unsupported_multimodal_rejectedmultimodal_data→UNIMPLEMENTEDtest_invalid_dp_rank_rejecteddp_rank >= dp_size→INVALID_ARGUMENTtest_multi_token_step_emissiontest_stream_token_zero_preservedoptional token_id0is present, not treated as missingtest_prefix_caching_default_enabledtest_prefix_caching_can_be_disabled--no-enable-prefix-cachingworkstest_serve_without_vllm_exits_unless_mockSystemExitwithout vLLM unless--mock-engine; skipped if vLLM is installedLast run (env with vLLM installed): 14 passed, 1 skipped.
B.
py_test/e2e/test_vllm_servicer_real_gpu.py(1 GPU, not CI)@pytest.mark.e2e+@pytest.mark.skipif(not HAS_VLLM).Spawns
python -m vllm_router.vllm_serviceron a real model (defaultQwen/Qwen3.5-4B). Client uses HuggingFaceAutoTokenizer, sendsprompt_token_ids, not text.Expected:
HealthCheckSERVING,GetModelInfomatches the model,GenerateStreamproduces a completion that contains"Paris"for the France prompt,ResetPrefixCachereturnssuccess=True.This does not run chat-template / Rust tokenize. That is PR 2. It proves the Servicer’s token-id ingress into a live engine.
C.
py_test/e2e/test_differential_grpc_vs_http.py(2 GPUs, not CI)Same model, same sampling (
temperature=0, fixed seed):python -m vllm.entrypoints.openai.api_server(POST /v1/completions, text)GenerateStream, pre-tokenized ids)Three prompts (factual, arithmetic, code). Compare generated text, finish_reason, and token sequence.
Expected: byte-for-byte identical text, identical tokens, identical finish reason. That is the claim that the gRPC path is the same engine, not a different sampler.
Already observed on live H800s for those three cases.
How to run
From the router repo root, after
pip install -e .(stubs are generated bysetup.py).CPU / CI (no GPU):
pytest py_test/test_vllm_servicer.py -v # same ignore CI uses: pytest py_test/ -v --ignore=py_test/e2eLint:
Regenerate stubs only (if you edited the proto without reinstalling):
VLLM_ROUTER_BUILD_NO_RUST=1 python -c "from setup import generate_grpc_stubs; generate_grpc_stubs()"(
import setupalso runssetup(); the stubs are written as a side effect.pip install -e .is the supported path.)Mock Servicer (no GPU):
vllm-servicer --model mock-model --mock-engine --host 127.0.0.1 --port 50051 # or: python -m vllm_router.vllm_servicer --model mock-model --mock-engine --host 127.0.0.1 --port 50051Real GPU Servicer:
GPU e2e:
pytest py_test/e2e/test_vllm_servicer_real_gpu.py -v # or: python py_test/e2e/test_vllm_servicer_real_gpu.pyDifferential (2 GPUs):
pytest py_test/e2e/test_differential_grpc_vs_http.py -v # or: python py_test/e2e/test_differential_grpc_vs_http.pyMODEL_PATHoverrides the default checkpoint for the e2e tests.Test plan
pytest py_test/test_vllm_servicer.py -v— 14 passed, 1 skipped when vLLM is installedResetPrefixCache(prior H800 run; re-run after this proto trim if you want a fresh log)Essential Elements of an Effective PR Description Checklist