@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e3c2fc3b3c

[Rust Frontend][gRPC] Add server and model discovery (#49491)

Signed-off-by: Connor Carpenter <connorc@nvidia.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Connor Carpenter committed Jul 27, 2026 at 09:53 UTC e3c2fc3b3ca3f41466126cd2aa0b228eb8465844
18 files changed +615 -269
rust/proto/control.proto new
+53
@@ -0,0 +1,53 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +syntax = "proto3";
5 +package vllm;
6 +
7 +service Control {
8 + rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {}
9 + rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {}
10 + rpc Abort (AbortRequest) returns (AbortResponse) {}
11 +}
12 +
13 +message GetServerInfoRequest {}
14 +
15 +message ServerInfo {
16 + string engine_version = 1;
17 + string api_version = 2;
18 + string instance_id = 3;
19 + ParallelismInfo parallelism = 4;
20 + uint32 max_model_len = 5;
21 + uint32 kv_block_size = 6;
22 + uint64 total_kv_blocks = 7;
23 + uint64 max_running_requests = 8;
24 + uint64 max_batched_tokens = 9;
25 +}
26 +
27 +message ParallelismInfo {
28 + uint32 tensor_parallel_size = 1;
29 + uint32 pipeline_parallel_size = 2;
30 + uint32 data_parallel_size = 3;
31 + uint32 data_parallel_rank = 4;
32 + uint32 decode_context_parallel_size = 5;
33 +}
34 +
35 +message GetModelInfoRequest {}
36 +
37 +message ModelInfo {
38 + string model_id = 1;
39 + string served_model_name = 2;
40 + repeated string served_model_aliases = 3;
41 +
42 + bool supports_text_input = 20;
43 + bool supports_token_ids_input = 21;
44 + bool supports_multimodal = 23;
45 + string reasoning_parser = 24;
46 + string tool_call_parser = 25;
47 +}
48 +
49 +message AbortRequest {
50 + repeated string request_ids = 1;
51 +}
52 +
53 +message AbortResponse {}
rust/proto/inference.proto renamed
+1 -15
@@ -7,17 +7,13 @@ package vllm;
7 import "google/protobuf/struct.proto";
8
9
10 -service Generate {
10 +service Inference {
11 // Generates text given a prompt
12 rpc Generate (GenerateRequest) returns (GenerateResponse) {}
13 // Generates text given a prompt, streaming the outputs
14 rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
15 }
16
17 -service Control {
18 - rpc Abort (AbortRequest) returns (AbortResponse) {}
19 -}
20 -
17 // ======================================================================================
18 // Generate Request
19 // ======================================================================================
@@ -204,13 +200,3 @@ message CandidateTokenInfo {
200 message TokenIds {
201 repeated uint32 ids = 1;
202 }
207 -
208 -// ======================================================================================
209 -// Control
210 -// ======================================================================================
211 -
212 -message AbortRequest {
213 - repeated string request_ids = 1;
214 -}
215 -
216 -message AbortResponse {}
rust/src/chat/src/lib.rs
+27
@@ -255,6 +255,33 @@ impl ChatLlm {
255 self.text.engine_core_client()
256 }
257
258 + /// Whether the loaded backend has a registered multimodal processor.
259 + pub fn supports_multimodal(&self) -> bool {
260 + self.processor.backend.multimodal_model_info().is_some()
261 + }
262 +
263 + /// Effective tool-call parser name for this model, if parsing is enabled.
264 + pub fn tool_call_parser_name(&self) -> Option<&str> {
265 + match &self.tool_call_parser {
266 + ParserSelection::Auto => {
267 + ToolParserFactory::global().resolve_name_for_model(self.model_id())
268 + }
269 + ParserSelection::None => None,
270 + ParserSelection::Explicit(name) => Some(name),
271 + }
272 + }
273 +
274 + /// Effective reasoning parser name for this model, if parsing is enabled.
275 + pub fn reasoning_parser_name(&self) -> Option<&str> {
276 + match &self.reasoning_parser {
277 + ParserSelection::Auto => {
278 + ReasoningParserFactory::global().resolve_name_for_model(self.model_id())
279 + }
280 + ParserSelection::None => None,
281 + ParserSelection::Explicit(name) => Some(name),
282 + }
283 + }
284 +
285 /// Render, tokenize, and submit one chat request.
286 pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
287 let (text_request, output_processor) = self
rust/src/engine-core-client/src/client.rs
+11
@@ -394,6 +394,17 @@ impl EngineCoreClient {
394 self.engines.iter().map(|engine| &engine.ready_response).collect()
395 }
396
397 + /// Return the first engine's ready response.
398 + ///
399 + /// Per-engine fields such as `data_parallel_rank` should be read through
400 + /// [`ready_responses`](Self::ready_responses).
401 + pub fn ready_response(&self) -> &EngineCoreReadyResponse {
402 + &self
403 + .engines
404 + .first()
405 + .expect("engine core client requires at least one engine")
406 + .ready_response
407 + }
408 /// Return the engine-reported effective model dtype.
409 pub fn model_dtype(&self) -> ModelDtype {
410 self.engines
rust/src/engine-core-client/src/mock_engine.rs
+7
@@ -57,6 +57,13 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
57 vllm_version: "test-vllm-version".to_string(),
58 world_size: 1,
59 data_parallel_size: 1,
60 + tensor_parallel_size: 1,
61 + pipeline_parallel_size: 1,
62 + decode_context_parallel_size: 1,
63 + data_parallel_rank: 0,
64 + max_num_seqs: 256,
65 + max_num_batched_tokens: 8192,
66 + instance_id: "test-instance".to_string(),
67 kv_cache_size_tokens: None,
68 kv_cache_max_concurrency: None,
69 }
rust/src/engine-core-client/src/protocol/handshake.rs
+15
@@ -52,6 +52,21 @@ pub struct EngineCoreReadyResponse {
52 pub world_size: u64,
53 /// Data parallelism size from the parallel config.
54 pub data_parallel_size: u64,
55 + // Required discovery metadata; EngineCore and client versions must match.
56 + /// Tensor-parallel size of this engine.
57 + pub tensor_parallel_size: u32,
58 + /// Pipeline-parallel size of this engine.
59 + pub pipeline_parallel_size: u32,
60 + /// Decode-context-parallel size of this engine.
61 + pub decode_context_parallel_size: u32,
62 + /// This engine's data-parallel rank.
63 + pub data_parallel_rank: u32,
64 + /// Scheduler cap on concurrently running sequences.
65 + pub max_num_seqs: u64,
66 + /// Scheduler cap on batched tokens per step.
67 + pub max_num_batched_tokens: u64,
68 + /// Unique identifier for this server instance.
69 + pub instance_id: String,
70 /// Total KV cache capacity in tokens, if reported.
71 pub kv_cache_size_tokens: Option<u64>,
72 /// Maximum achievable request concurrency given the KV cache, if reported.
rust/src/engine-core-client/src/tests/python_compat.py
+14
@@ -363,6 +363,13 @@ class EngineCoreReadyResponse:
363 vllm_version: str
364 world_size: int
365 data_parallel_size: int
366 + tensor_parallel_size: int
367 + pipeline_parallel_size: int
368 + decode_context_parallel_size: int
369 + data_parallel_rank: int
370 + max_num_seqs: int
371 + max_num_batched_tokens: int
372 + instance_id: str
373 kv_cache_size_tokens: int | None = None
374 kv_cache_max_concurrency: float | None = None
375
@@ -376,6 +383,13 @@ ready_response = EngineCoreReadyResponse(
383 vllm_version="0.0.0",
384 data_parallel_size=1,
385 world_size=1,
386 + tensor_parallel_size=1,
387 + pipeline_parallel_size=1,
388 + decode_context_parallel_size=1,
389 + data_parallel_rank=0,
390 + max_num_seqs=256,
391 + max_num_batched_tokens=8192,
392 + instance_id="test-instance",
393 )
394
395 print(msgspec.msgpack.encode(request).hex())
rust/src/server/build.rs
+7 -1
@@ -9,7 +9,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
9 .build_server(true)
10 .build_client(true)
11 .protoc_arg("--experimental_allow_proto3_optional") // be compatible with old compilers
12 - .compile_protos(&[format!("{proto_dir}/vllm_grpc.proto")], &[proto_dir])?;
12 + .compile_protos(
13 + &[
14 + format!("{proto_dir}/control.proto"),
15 + format!("{proto_dir}/inference.proto"),
16 + ],
17 + &[proto_dir],
18 + )?;
19
20 Ok(())
21 }
rust/src/server/src/grpc/control.rs new
+106
@@ -0,0 +1,106 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +
6 +use thiserror_ext::AsReport as _;
7 +use tonic::{Request, Response, Status};
8 +use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse;
9 +
10 +use super::{ControlServer, pb};
11 +use crate::state::AppState;
12 +
13 +pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
14 +
15 +/// gRPC control service backed by the shared application state.
16 +pub struct ControlServiceImpl {
17 + state: Arc<AppState>,
18 +}
19 +
20 +impl ControlServiceImpl {
21 + pub fn new(state: Arc<AppState>) -> Self {
22 + Self { state }
23 + }
24 +
25 + fn ready(&self) -> &EngineCoreReadyResponse {
26 + self.state.engine_core_client().ready_response()
27 + }
28 +
29 + fn parallelism_info(&self) -> pb::ParallelismInfo {
30 + let ready = self.ready();
31 + pb::ParallelismInfo {
32 + tensor_parallel_size: ready.tensor_parallel_size,
33 + pipeline_parallel_size: ready.pipeline_parallel_size,
34 + data_parallel_size: ready.data_parallel_size.min(u64::from(u32::MAX)) as u32,
35 + data_parallel_rank: ready.data_parallel_rank,
36 + decode_context_parallel_size: ready.decode_context_parallel_size,
37 + }
38 + }
39 +}
40 +
41 +const GRPC_API_VERSION: &str = "vllm";
42 +
43 +#[tonic::async_trait]
44 +impl pb::control_server::Control for ControlServiceImpl {
45 + async fn get_server_info(
46 + &self,
47 + _request: Request<pb::GetServerInfoRequest>,
48 + ) -> Result<Response<pb::ServerInfo>, Status> {
49 + let ready = self.ready();
50 + Ok(Response::new(pb::ServerInfo {
51 + engine_version: ready.vllm_version.clone(),
52 + api_version: GRPC_API_VERSION.to_string(),
53 + instance_id: ready.instance_id.clone(),
54 + parallelism: Some(self.parallelism_info()),
55 + max_model_len: self.state.engine_core_client().max_model_len(),
56 + kv_block_size: ready.block_size.min(u64::from(u32::MAX)) as u32,
57 + total_kv_blocks: self.state.engine_core_client().total_num_gpu_blocks(),
58 + max_running_requests: ready.max_num_seqs,
59 + max_batched_tokens: ready.max_num_batched_tokens,
60 + }))
61 + }
62 +
63 + async fn get_model_info(
64 + &self,
65 + _request: Request<pb::GetModelInfoRequest>,
66 + ) -> Result<Response<pb::ModelInfo>, Status> {
67 + let served = self.state.served_model_names();
68 + Ok(Response::new(pb::ModelInfo {
69 + model_id: self.state.chat.text().model_id().to_string(),
70 + served_model_name: self.state.primary_model_name().to_string(),
71 + served_model_aliases: served.iter().skip(1).cloned().collect(),
72 + // GenerateRequest accepts both prompt representations.
73 + supports_text_input: true,
74 + supports_token_ids_input: true,
75 + supports_multimodal: self.state.chat.supports_multimodal(),
76 + reasoning_parser: self
77 + .state
78 + .chat
79 + .reasoning_parser_name()
80 + .unwrap_or_default()
81 + .to_string(),
82 + tool_call_parser: self
83 + .state
84 + .chat
85 + .tool_call_parser_name()
86 + .unwrap_or_default()
87 + .to_string(),
88 + }))
89 + }
90 +
91 + async fn abort(
92 + &self,
93 + request: Request<pb::AbortRequest>,
94 + ) -> Result<Response<pb::AbortResponse>, Status> {
95 + let request_ids = request.into_inner().request_ids;
96 + if request_ids.is_empty() {
97 + return Ok(Response::new(pb::AbortResponse {}));
98 + }
99 + self.state
100 + .chat
101 + .abort(&request_ids)
102 + .await
103 + .map_err(|error| Status::internal(error.to_report_string()))?;
104 + Ok(Response::new(pb::AbortResponse {}))
105 + }
106 +}
rust/src/server/src/grpc/health.rs
+4 -4
@@ -8,14 +8,14 @@ use tonic_health::ServingStatus;
8 use tonic_health::server::HealthReporter;
9 use tracing::{info, warn};
10
11 -use super::{ControlGrpcService, GenerateGrpcService};
11 +use super::{ControlGrpcService, InferenceGrpcService};
12
13 pub(crate) async fn monitor_health(
14 mut health_reporter: HealthReporter,
15 mut engine_health: watch::Receiver<bool>,
16 shutdown: CancellationToken,
17 ) {
18 - let generate_service = GenerateGrpcService::NAME;
18 + let inference_service = InferenceGrpcService::NAME;
19 let control_service = ControlGrpcService::NAME;
20 let status = ServingStatus::NotServing;
21 let health_event_first = tokio::select! {
@@ -45,7 +45,7 @@ pub(crate) async fn monitor_health(
45 }
46 };
47
48 - health_reporter.set_not_serving::<GenerateGrpcService>().await;
48 + health_reporter.set_not_serving::<InferenceGrpcService>().await;
49 health_reporter.set_not_serving::<ControlGrpcService>().await;
50 // Both gRPC services use the same engine client, so overall server health
51 // mirrors their shared engine health.
@@ -59,7 +59,7 @@ pub(crate) async fn monitor_health(
59 );
60 }
61
62 - health_reporter.clear_service_status(generate_service).await;
62 + health_reporter.clear_service_status(inference_service).await;
63 health_reporter.clear_service_status(control_service).await;
64 health_reporter.clear_service_status("").await;
65 }
rust/src/server/src/grpc/inference.rs new
+156
@@ -0,0 +1,156 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::pin::Pin;
5 +use std::sync::Arc;
6 +
7 +use futures::{Stream, StreamExt as _};
8 +use thiserror_ext::AsReport as _;
9 +use tokio::sync::mpsc;
10 +use tokio_stream::wrappers::ReceiverStream;
11 +use tonic::{Request, Response, Status};
12 +use tracing::info;
13 +use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _};
14 +
15 +use super::convert::{self, ResponseOpts};
16 +use super::{InferenceServer, pb};
17 +use crate::state::AppState;
18 +
19 +pub(crate) type InferenceGrpcService = InferenceServer<InferenceServiceImpl>;
20 +
21 +/// gRPC inference service backed by the shared application state.
22 +pub struct InferenceServiceImpl {
23 + state: Arc<AppState>,
24 +}
25 +
26 +impl InferenceServiceImpl {
27 + pub fn new(state: Arc<AppState>) -> Self {
28 + Self { state }
29 + }
30 +}
31 +
32 +#[tonic::async_trait]
33 +impl pb::inference_server::Inference for InferenceServiceImpl {
34 + type GenerateStreamStream =
35 + Pin<Box<dyn Stream<Item = Result<pb::GenerateResponse, Status>> + Send>>;
36 +
37 + /// Unary generate: collect all output and return a single response.
38 + async fn generate(
39 + &self,
40 + request: Request<pb::GenerateRequest>,
41 + ) -> Result<Response<pb::GenerateResponse>, Status> {
42 + let proto_req = request.into_inner();
43 + let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref());
44 + let text_request =
45 + convert::to_text_request(proto_req, false, self.state.served_model_names())?;
46 +
47 + let request_id = text_request.request_id.clone();
48 + info!(%request_id, "grpc generate (unary)");
49 +
50 + let stream = self.state.chat.text().generate(text_request).await;
51 + let stream = stream.map_err(text_error_to_status)?;
52 +
53 + let collected = stream.collect_output().await.map_err(text_error_to_status)?;
54 +
55 + // Build the single aggregated response.
56 + let prompt_info = convert::to_prompt_info(
57 + &collected.prompt_token_ids,
58 + collected.prompt_logprobs.as_ref(),
59 + &response_opts,
60 + );
61 +
62 + let finish_info = vllm_text::Finished {
63 + usage: collected.usage,
64 + finish_reason: collected.finish_reason,
65 + kv_transfer_params: collected.kv_transfer_params,
66 + ec_transfer_params: collected.ec_transfer_params,
67 + };
68 +
69 + let outputs = convert::to_sequence_output(
70 + &collected.text,
71 + &collected.token_ids,
72 + collected.logprobs.as_ref(),
73 + Some(&finish_info),
74 + &response_opts,
75 + );
76 +
77 + Ok(Response::new(pb::GenerateResponse {
78 + prompt_info: Some(prompt_info),
79 + outputs: Some(outputs),
80 + }))
81 + }
82 +
83 + /// Streaming generate: yield incremental responses as tokens are produced.
84 + async fn generate_stream(
85 + &self,
86 + request: Request<pb::GenerateRequest>,
87 + ) -> Result<Response<Self::GenerateStreamStream>, Status> {
88 + let proto_req = request.into_inner();
89 + let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref());
90 + let text_request =
91 + convert::to_text_request(proto_req, true, self.state.served_model_names())?;
92 +
93 + let request_id = text_request.request_id.clone();
94 + info!(%request_id, "grpc generate (stream)");
95 +
96 + let stream = self.state.chat.text().generate(text_request).await;
97 + let stream = stream.map_err(text_error_to_status)?;
98 +
99 + let (tx, rx) = mpsc::channel(32);
100 +
101 + tokio::spawn(async move {
102 + futures::pin_mut!(stream);
103 + while let Some(event) = stream.next().await {
104 + let response = match event {
105 + Err(e) => Err(text_error_to_status(e)),
106 + Ok(DecodedTextEvent::Start {
107 + prompt_token_ids,
108 + prompt_logprobs,
109 + }) => {
110 + let prompt_info = convert::to_prompt_info(
111 + &prompt_token_ids,
112 + prompt_logprobs.as_ref(),
113 + &response_opts,
114 + );
115 + Ok(pb::GenerateResponse {
116 + prompt_info: Some(prompt_info),
117 + outputs: None,
118 + })
119 + }
120 + Ok(DecodedTextEvent::TextDelta {
121 + delta,
122 + token_ids,
123 + logprobs,
124 + finished,
125 + }) => Ok(pb::GenerateResponse {
126 + prompt_info: None,
127 + outputs: Some(convert::to_sequence_output(
128 + &delta,
129 + &token_ids,
130 + logprobs.as_ref(),
131 + finished.as_ref(),
132 + &response_opts,
133 + )),
134 + }),
135 + };
136 +
137 + if tx.send(response).await.is_err() {
138 + // Client disconnected.
139 + break;
140 + }
141 + }
142 + });
143 +
144 + let response_stream = ReceiverStream::new(rx);
145 + Ok(Response::new(Box::pin(response_stream)))
146 + }
147 +}
148 +
149 +fn text_error_to_status(error: vllm_text::Error) -> Status {
150 + let message = error.to_report_string();
151 + if error.is_request_validation_error() {
152 + Status::invalid_argument(message)
153 + } else {
154 + Status::internal(message)
155 + }
156 +}
rust/src/server/src/grpc/mod.rs
+8 -186
@@ -1,203 +1,25 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
4 +//! gRPC services backed by the shared application state.
5
6 +mod control;
7 mod convert;
8 mod health;
8 -
9 -use std::pin::Pin;
10 -use std::sync::Arc;
11 -
12 -use futures::{Stream, StreamExt as _};
13 -use thiserror_ext::AsReport as _;
14 -use tokio::sync::mpsc;
15 -use tokio_stream::wrappers::ReceiverStream;
16 -use tonic::{Request, Response, Status};
17 -use tracing::info;
18 -use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _};
19 -
20 -use self::convert::ResponseOpts;
21 -use crate::state::AppState;
9 +mod inference;
10
11 /// Generated protobuf/gRPC types for the `vllm` package.
12 pub mod pb {
13 tonic::include_proto!("vllm");
14 }
15
16 +pub(crate) use control::ControlGrpcService;
17 +pub use control::ControlServiceImpl;
18 pub(crate) use health::monitor_health;
19 +pub(crate) use inference::InferenceGrpcService;
20 +pub use inference::InferenceServiceImpl;
21 pub use pb::control_server::ControlServer;
30 -pub use pb::generate_server::GenerateServer;
31 -
32 -pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
33 -pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
22 +pub use pb::inference_server::InferenceServer;
23
24 #[cfg(test)]
25 mod tests;
37 -
38 -/// gRPC Generate service implementation backed by the shared application state.
39 -pub struct GenerateServiceImpl {
40 - state: Arc<AppState>,
41 -}
42 -
43 -impl GenerateServiceImpl {
44 - pub fn new(state: Arc<AppState>) -> Self {
45 - Self { state }
46 - }
47 -}
48 -
49 -/// gRPC control service backed by the shared application state.
50 -pub struct ControlServiceImpl {
51 - state: Arc<AppState>,
52 -}
53 -
54 -impl ControlServiceImpl {
55 - pub fn new(state: Arc<AppState>) -> Self {
56 - Self { state }
57 - }
58 -}
59 -
60 -#[tonic::async_trait]
61 -impl pb::control_server::Control for ControlServiceImpl {
62 - async fn abort(
63 - &self,
64 - request: Request<pb::AbortRequest>,
65 - ) -> Result<Response<pb::AbortResponse>, Status> {
66 - let request_ids = request.into_inner().request_ids;
67 - if request_ids.is_empty() {
68 - return Ok(Response::new(pb::AbortResponse {}));
69 - }
70 - self.state
71 - .chat
72 - .abort(&request_ids)
73 - .await
74 - .map_err(|error| Status::internal(error.to_report_string()))?;
75 - Ok(Response::new(pb::AbortResponse {}))
76 - }
77 -}
78 -
79 -#[tonic::async_trait]
80 -impl pb::generate_server::Generate for GenerateServiceImpl {
81 - type GenerateStreamStream =
82 - Pin<Box<dyn Stream<Item = Result<pb::GenerateResponse, Status>> + Send>>;
83 -
84 - /// Unary generate: collect all output and return a single response.
85 - async fn generate(
86 - &self,
87 - request: Request<pb::GenerateRequest>,
88 - ) -> Result<Response<pb::GenerateResponse>, Status> {
89 - let proto_req = request.into_inner();
90 - let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref());
91 - let text_request =
92 - convert::to_text_request(proto_req, false, self.state.served_model_names())?;
93 -
94 - let request_id = text_request.request_id.clone();
95 - info!(%request_id, "grpc generate (unary)");
96 -
97 - let stream = self.state.chat.text().generate(text_request).await;
98 - let stream = stream.map_err(text_error_to_status)?;
99 -
100 - let collected = stream.collect_output().await.map_err(text_error_to_status)?;
101 -
102 - // Build the single aggregated response.
103 - let prompt_info = convert::to_prompt_info(
104 - &collected.prompt_token_ids,
105 - collected.prompt_logprobs.as_ref(),
106 - &response_opts,
107 - );
108 -
109 - let finish_info = vllm_text::Finished {
110 - usage: collected.usage,
111 - finish_reason: collected.finish_reason,
112 - kv_transfer_params: collected.kv_transfer_params,
113 - ec_transfer_params: collected.ec_transfer_params,
114 - };
115 -
116 - let outputs = convert::to_sequence_output(
117 - &collected.text,
118 - &collected.token_ids,
119 - collected.logprobs.as_ref(),
120 - Some(&finish_info),
121 - &response_opts,
122 - );
123 -
124 - Ok(Response::new(pb::GenerateResponse {
125 - prompt_info: Some(prompt_info),
126 - outputs: Some(outputs),
127 - }))
128 - }
129 -
130 - /// Streaming generate: yield incremental responses as tokens are produced.
131 - async fn generate_stream(
132 - &self,
133 - request: Request<pb::GenerateRequest>,
134 - ) -> Result<Response<Self::GenerateStreamStream>, Status> {
135 - let proto_req = request.into_inner();
136 - let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref());
137 - let text_request =
138 - convert::to_text_request(proto_req, true, self.state.served_model_names())?;
139 -
140 - let request_id = text_request.request_id.clone();
141 - info!(%request_id, "grpc generate (stream)");
142 -
143 - let stream = self.state.chat.text().generate(text_request).await;
144 - let stream = stream.map_err(text_error_to_status)?;
145 -
146 - let (tx, rx) = mpsc::channel(32);
147 -
148 - tokio::spawn(async move {
149 - futures::pin_mut!(stream);
150 - while let Some(event) = stream.next().await {
151 - let response = match event {
152 - Err(e) => Err(text_error_to_status(e)),
153 - Ok(DecodedTextEvent::Start {
154 - prompt_token_ids,
155 - prompt_logprobs,
156 - }) => {
157 - let prompt_info = convert::to_prompt_info(
158 - &prompt_token_ids,
159 - prompt_logprobs.as_ref(),
160 - &response_opts,
161 - );
162 - Ok(pb::GenerateResponse {
163 - prompt_info: Some(prompt_info),
164 - outputs: None,
165 - })
166 - }
167 - Ok(DecodedTextEvent::TextDelta {
168 - delta,
169 - token_ids,
170 - logprobs,
171 - finished,
172 - }) => Ok(pb::GenerateResponse {
173 - prompt_info: None,
174 - outputs: Some(convert::to_sequence_output(
175 - &delta,
176 - &token_ids,
177 - logprobs.as_ref(),
178 - finished.as_ref(),
179 - &response_opts,
180 - )),
181 - }),
182 - };
183 -
184 - if tx.send(response).await.is_err() {
185 - // Client disconnected.
186 - break;
187 - }
188 - }
189 - });
190 -
191 - let response_stream = ReceiverStream::new(rx);
192 - Ok(Response::new(Box::pin(response_stream)))
193 - }
194 -}
195 -
196 -fn text_error_to_status(error: vllm_text::Error) -> Status {
197 - let message = error.to_report_string();
198 - if error.is_request_validation_error() {
199 - Status::invalid_argument(message)
200 - } else {
201 - Status::internal(message)
202 - }
203 -}
rust/src/server/src/grpc/tests.rs
+157 -38
@@ -25,12 +25,18 @@ use vllm_chat::{
25 ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
26 DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
27 };
28 +use vllm_engine_core_client::mock_engine::{
29 + DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS,
30 + default_ready_response,
31 +};
32 use vllm_engine_core_client::protocol::output::{
33 EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs,
34 };
35 use vllm_engine_core_client::protocol::request::EngineCoreRequest;
32 -use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
33 -use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
36 +use vllm_engine_core_client::test_utils::{
37 + IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready,
38 +};
39 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId, TransportMode};
40 use vllm_llm::Llm;
41 use vllm_text::tokenizer::DynTokenizer;
42 use vllm_text::{Prompt, TextBackend};
@@ -39,8 +45,8 @@ use zeromq::prelude::{SocketRecv, SocketSend};
45 use zeromq::{DealerSocket, PushSocket, ZmqMessage};
46
47 use super::pb::control_client::ControlClient;
42 -use super::pb::generate_client::GenerateClient;
43 -use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb};
48 +use super::pb::inference_client::InferenceClient;
49 +use super::{ControlServer, ControlServiceImpl, InferenceServer, InferenceServiceImpl, pb};
50 use crate::listener::{Listener, MaybeTlsListener};
51 use crate::state::AppState;
52 use crate::tls;
@@ -202,7 +208,7 @@ async fn setup_grpc_service(
208 engine_id: impl Into<EngineId>,
209 output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
210 ) -> (
205 - GenerateServer<GenerateServiceImpl>,
211 + InferenceServer<InferenceServiceImpl>,
212 ControlServer<ControlServiceImpl>,
213 tokio::sync::watch::Receiver<bool>,
214 MockEngineTask,
@@ -246,7 +252,7 @@ async fn setup_grpc_service(
252 );
253 let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
254 (
249 - GenerateServer::new(GenerateServiceImpl::new(state.clone())),
255 + InferenceServer::new(InferenceServiceImpl::new(state.clone())),
256 ControlServer::new(ControlServiceImpl::new(state)),
257 engine_health,
258 engine_task,
@@ -259,30 +265,30 @@ async fn grpc_test_server(
265 engine_id: impl Into<EngineId>,
266 output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
267 ) -> (
262 - GenerateClient<tonic::transport::Channel>,
268 + InferenceClient<tonic::transport::Channel>,
269 tokio::task::JoinHandle<()>,
270 MockEngineTask,
271 ) {
266 - let (generate_service, control_service, engine_health, engine_task) =
272 + let (inference_service, control_service, engine_health, engine_task) =
273 setup_grpc_service(engine_id, output_specs).await;
274 let (channel, server_task) = start_grpc_test_server(
269 - generate_service,
275 + inference_service,
276 control_service,
277 engine_health,
278 tokio_util::sync::CancellationToken::new(),
279 )
280 .await;
275 - (GenerateClient::new(channel), server_task, engine_task)
281 + (InferenceClient::new(channel), server_task, engine_task)
282 }
283
284 async fn start_grpc_test_server(
279 - generate_service: GenerateServer<GenerateServiceImpl>,
285 + inference_service: InferenceServer<InferenceServiceImpl>,
286 control_service: ControlServer<ControlServiceImpl>,
287 engine_health: tokio::sync::watch::Receiver<bool>,
288 shutdown: tokio_util::sync::CancellationToken,
289 ) -> (Channel, tokio::task::JoinHandle<()>) {
290 let (health_reporter, health_service) = health_reporter();
285 - health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
291 + health_reporter.set_serving::<InferenceServer<InferenceServiceImpl>>().await;
292 health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await;
293
294 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
@@ -293,7 +299,7 @@ async fn start_grpc_test_server(
299 let server = TonicServer::builder()
300 .add_service(health_service)
301 .add_service(control_service)
296 - .add_service(generate_service)
302 + .add_service(inference_service)
303 .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
304 let health_monitor =
305 super::monitor_health(health_reporter, engine_health, shutdown.clone());
@@ -323,7 +329,7 @@ async fn grpc_tls_test_server(
329 certs: &TestCerts,
330 cert_reqs: i32,
331 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
326 - let (generate_service, control_service, _engine_health, engine_task) =
332 + let (inference_service, control_service, _engine_health, engine_task) =
333 setup_grpc_service(engine_id, output_specs).await;
334 let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
335 .expect("build grpc tls config");
@@ -335,7 +341,7 @@ async fn grpc_tls_test_server(
341 let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
342 TonicServer::builder()
343 .add_service(control_service)
338 - .add_service(generate_service)
344 + .add_service(inference_service)
345 .serve_with_incoming(incoming)
346 .await
347 .expect("grpc tls server");
@@ -351,7 +357,7 @@ async fn grpc_tls_client(
357 certs: &TestCerts,
358 addr: &str,
359 identity: Option<&str>,
354 -) -> Result<GenerateClient<Channel>, tonic::transport::Error> {
360 +) -> Result<InferenceClient<Channel>, tonic::transport::Error> {
361 let ca = certs.path("ca.pem");
362 let identity = identity.map(|name| {
363 (
@@ -388,7 +394,7 @@ async fn grpc_tls_client(
394 .expect("grpc endpoint")
395 .connect_with_connector(connector)
396 .await?;
391 - Ok(GenerateClient::new(channel))
397 + Ok(InferenceClient::new(channel))
398 }
399
400 /// Complete a raw TLS handshake against the gRPC port (offering ALPN `h2`) for
@@ -415,7 +421,7 @@ async fn grpc_server_with_keepalive(
421 engine_id: impl Into<EngineId>,
422 keepalive: Option<Duration>,
423 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
418 - let (generate_service, control_service, _engine_health, engine_task) =
424 + let (inference_service, control_service, _engine_health, engine_task) =
425 setup_grpc_service(engine_id, default_stream_output_specs()).await;
426
427 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
@@ -432,7 +438,7 @@ async fn grpc_server_with_keepalive(
438 let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
439 builder
440 .add_service(control_service)
435 - .add_service(generate_service)
441 + .add_service(inference_service)
442 .serve_with_incoming(incoming)
443 .await
444 .expect("grpc server");
@@ -1083,20 +1089,20 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
1089 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1090 #[serial]
1091 async fn control_abort_resolves_external_id_and_empty_is_noop() {
1086 - let (generate_service, control_service, engine_health, engine_task) =
1092 + let (inference_service, control_service, engine_health, engine_task) =
1093 setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await;
1094 let (channel, server_task) = start_grpc_test_server(
1089 - generate_service,
1095 + inference_service,
1096 control_service,
1097 engine_health,
1098 tokio_util::sync::CancellationToken::new(),
1099 )
1100 .await;
1095 - let mut generate_client = GenerateClient::new(channel.clone());
1101 + let mut inference_client = InferenceClient::new(channel.clone());
1102 let mut control_client = ControlClient::new(channel);
1103 let request_id = "test-abort-active";
1104
1099 - let mut stream = generate_client
1105 + let mut stream = inference_client
1106 .generate_stream(pb::GenerateRequest {
1107 request_id: request_id.to_string(),
1108 model: "test-model".to_string(),
@@ -1171,14 +1177,127 @@ async fn control_abort_resolves_external_id_and_empty_is_noop() {
1177 server_task.abort();
1178 }
1179
1180 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1181 +#[serial]
1182 +async fn control_reports_server_and_model_info() {
1183 + let (generate_service, control_service, engine_health, _engine_task) =
1184 + setup_grpc_service(b"engine-grpc-info", default_stream_output_specs()).await;
1185 + let (channel, server_task) = start_grpc_test_server(
1186 + generate_service,
1187 + control_service,
1188 + engine_health,
1189 + tokio_util::sync::CancellationToken::new(),
1190 + )
1191 + .await;
1192 + let mut client = ControlClient::new(channel);
1193 +
1194 + let server = client
1195 + .get_server_info(pb::GetServerInfoRequest {})
1196 + .await
1197 + .expect("get server info")
1198 + .into_inner();
1199 + assert_eq!(server.engine_version, "test-vllm-version");
1200 + assert_eq!(server.api_version, "vllm");
1201 + assert_eq!(server.instance_id, "test-instance");
1202 + assert_eq!(server.max_model_len, DEFAULT_MOCK_MAX_MODEL_LEN as u32);
1203 + assert_eq!(server.kv_block_size, DEFAULT_MOCK_BLOCK_SIZE as u32);
1204 + assert_eq!(server.total_kv_blocks, DEFAULT_MOCK_NUM_GPU_BLOCKS);
1205 + assert_eq!(server.max_running_requests, 256);
1206 + assert_eq!(server.max_batched_tokens, 8_192);
1207 + let parallelism = server.parallelism.expect("parallelism metadata");
1208 + assert_eq!(parallelism.tensor_parallel_size, 1);
1209 + assert_eq!(parallelism.pipeline_parallel_size, 1);
1210 + assert_eq!(parallelism.data_parallel_size, 1);
1211 + assert_eq!(parallelism.data_parallel_rank, 0);
1212 + assert_eq!(parallelism.decode_context_parallel_size, 1);
1213 +
1214 + let model = client
1215 + .get_model_info(pb::GetModelInfoRequest {})
1216 + .await
1217 + .expect("get model info")
1218 + .into_inner();
1219 + assert_eq!(model.model_id, "test-model");
1220 + assert_eq!(model.served_model_name, "test-model");
1221 + assert!(model.served_model_aliases.is_empty());
1222 + assert!(model.supports_text_input);
1223 + assert!(model.supports_token_ids_input);
1224 + assert!(!model.supports_multimodal);
1225 + assert!(model.reasoning_parser.is_empty());
1226 + assert!(model.tool_call_parser.is_empty());
1227 +
1228 + server_task.abort();
1229 +}
1230 +
1231 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1232 +async fn control_aggregates_multi_engine_capacity() {
1233 + let ipc = IpcNamespace::new().expect("create ipc namespace");
1234 + let handshake_address = ipc.handshake_endpoint();
1235 +
1236 + let mut ready_0 = default_ready_response();
1237 + ready_0.max_model_len = 8_192;
1238 + ready_0.num_gpu_blocks = 10;
1239 + ready_0.data_parallel_size = 2;
1240 +
1241 + let mut ready_1 = default_ready_response();
1242 + ready_1.max_model_len = 4_096;
1243 + ready_1.num_gpu_blocks = 20;
1244 + ready_1.data_parallel_size = 2;
1245 + ready_1.data_parallel_rank = 1;
1246 +
1247 + let engine_tasks = [ready_0, ready_1].map(|ready| {
1248 + let engine_id = EngineId::from_engine_index(ready.data_parallel_rank);
1249 + MockEngineTask::new(spawn_mock_engine_task_with_ready(
1250 + handshake_address.clone(),
1251 + engine_id,
1252 + ready,
1253 + |_, _| boxed_test_future(async {}),
1254 + ))
1255 + });
1256 +
1257 + let client = EngineCoreClient::connect(EngineCoreClientConfig {
1258 + transport_mode: TransportMode::HandshakeOwner {
1259 + handshake_address,
1260 + advertised_host: "127.0.0.1".to_string(),
1261 + engine_count: 2,
1262 + ready_timeout: Duration::from_secs(2),
1263 + local_input_address: Some(ipc.input_endpoint()),
1264 + local_output_address: Some(ipc.output_endpoint()),
1265 + },
1266 + coordinator_mode: None,
1267 + model_name: "test-model".to_string(),
1268 + client_index: 0,
1269 + })
1270 + .await
1271 + .expect("connect multi-engine client");
1272 + let chat = ChatLlm::from_shared_backend(
1273 + Llm::new(client),
1274 + Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
1275 + );
1276 + let service = ControlServiceImpl::new(Arc::new(AppState::new(
1277 + vec!["test-model".to_string()],
1278 + chat,
1279 + )));
1280 +
1281 + let server = pb::control_server::Control::get_server_info(
1282 + &service,
1283 + tonic::Request::new(pb::GetServerInfoRequest {}),
1284 + )
1285 + .await
1286 + .expect("get server info")
1287 + .into_inner();
1288 + assert_eq!(server.max_model_len, 4_096);
1289 + assert_eq!(server.total_kv_blocks, 30);
1290 +
1291 + drop(engine_tasks);
1292 +}
1293 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1294 #[serial]
1295 async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
1177 - let (generate_service, control_service, _connected_engine_health, _engine_task) =
1296 + let (inference_service, control_service, _connected_engine_health, _engine_task) =
1297 setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
1298 let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
1299 let (channel, server_task) = start_grpc_test_server(
1181 - generate_service,
1300 + inference_service,
1301 control_service,
1302 engine_health,
1303 tokio_util::sync::CancellationToken::new(),
@@ -1187,7 +1306,7 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
1306 let mut health_client = HealthClient::new(channel);
1307
1308 let mut health_streams = Vec::new();
1190 - for service in ["vllm.Generate", "vllm.Control", ""] {
1309 + for service in ["vllm.Inference", "vllm.Control", ""] {
1310 let service_label = if service.is_empty() {
1311 "overall"
1312 } else {
@@ -1242,14 +1361,14 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
1361 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1362 #[serial]
1363 async fn grpc_health_watch_closes_on_graceful_shutdown() {
1245 - let (generate_service, control_service, engine_health, _engine_task) = setup_grpc_service(
1364 + let (inference_service, control_service, engine_health, _engine_task) = setup_grpc_service(
1365 b"engine-grpc-health-shutdown",
1366 default_stream_output_specs(),
1367 )
1368 .await;
1369 let shutdown = tokio_util::sync::CancellationToken::new();
1370 let (channel, server_task) = start_grpc_test_server(
1252 - generate_service,
1371 + inference_service,
1372 control_service,
1373 engine_health,
1374 shutdown.clone(),
@@ -1258,43 +1377,43 @@ async fn grpc_health_watch_closes_on_graceful_shutdown() {
1377 let mut health_client = HealthClient::new(channel);
1378 let mut stream = health_client
1379 .watch(HealthCheckRequest {
1261 - service: "vllm.Generate".to_string(),
1380 + service: "vllm.Inference".to_string(),
1381 })
1382 .await
1264 - .expect("start health watch for vllm.Generate")
1383 + .expect("start health watch for vllm.Inference")
1384 .into_inner();
1385
1386 let initial = stream
1387 .message()
1388 .await
1270 - .expect("read initial health status for vllm.Generate")
1389 + .expect("read initial health status for vllm.Inference")
1390 .expect("health watch ended before its initial status");
1391 assert_eq!(
1392 initial.status,
1393 HealthServingStatus::Serving as i32,
1275 - "unexpected initial health status for vllm.Generate"
1394 + "unexpected initial health status for vllm.Inference"
1395 );
1396
1397 shutdown.cancel();
1398
1399 let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
1400 .await
1282 - .expect("timed out waiting for shutdown health update for vllm.Generate")
1283 - .expect("failed to read shutdown health update for vllm.Generate")
1401 + .expect("timed out waiting for shutdown health update for vllm.Inference")
1402 + .expect("failed to read shutdown health update for vllm.Inference")
1403 .expect("health watch ended before its shutdown update");
1404 assert_eq!(
1405 update.status,
1406 HealthServingStatus::NotServing as i32,
1288 - "unexpected shutdown health status for vllm.Generate"
1407 + "unexpected shutdown health status for vllm.Inference"
1408 );
1409
1410 let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message())
1411 .await
1293 - .expect("timed out waiting for vllm.Generate health watch to close")
1294 - .expect("failed while closing vllm.Generate health watch");
1412 + .expect("timed out waiting for vllm.Inference health watch to close")
1413 + .expect("failed while closing vllm.Inference health watch");
1414 assert!(
1415 stream_end.is_none(),
1297 - "vllm.Generate health watch remained open"
1416 + "vllm.Inference health watch remained open"
1417 );
1418
1419 tokio::time::timeout(Duration::from_secs(2), server_task)
rust/src/server/src/lib.rs
+5 -5
@@ -187,7 +187,7 @@ where
187 let model = state.primary_model_name().to_owned();
188 let app = extend_router(build_router(state.clone()));
189
190 - // Optionally bind the gRPC Generate server on a separate port. Bind
190 + // Optionally bind the gRPC Inference server on a separate port. Bind
191 // synchronously here so bind errors (port in use, permission denied, ...)
192 // surface before serving rather than being deferred until shutdown.
193 let grpc_setup = if let Some(grpc_port) = config.grpc_port {
@@ -206,19 +206,19 @@ where
206 .context("invalid gRPC TLS configuration")?;
207 let (health_reporter, health_service) = health_reporter();
208 let engine_health = state.engine_core_client().subscribe_health();
209 - health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
209 + health_reporter.set_serving::<grpc::InferenceGrpcService>().await;
210 health_reporter.set_serving::<grpc::ControlGrpcService>().await;
211 let control_service =
212 grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()));
213 - let generate_service =
214 - grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
213 + let inference_service =
214 + grpc::InferenceGrpcService::new(grpc::InferenceServiceImpl::new(state.clone()));
215 let svc = TonicServer::builder()
216 .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
217 .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
218 .layer(middleware::request_runtime_layer(state.clone()))
219 .add_service(health_service)
220 .add_service(control_service)
221 - .add_service(generate_service);
221 + .add_service(inference_service);
222 info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
223 Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
224 } else {
rust/src/server/src/middleware/offload.rs
+4 -4
@@ -30,8 +30,8 @@ const OFFLOADED_PATHS: &[&str] = &[
30 "/detokenize",
31 "/inference/v1/generate",
32 // gRPC routes:
33 - "/vllm.Generate/Generate",
34 - "/vllm.Generate/GenerateStream",
33 + "/vllm.Inference/Generate",
34 + "/vllm.Inference/GenerateStream",
35 ];
36
37 /// Return a Tower layer that runs selected data-plane requests on the request runtime,
@@ -124,8 +124,8 @@ mod tests {
124 assert!(should_offload("/tokenize"));
125 assert!(should_offload("/detokenize"));
126 assert!(should_offload("/inference/v1/generate"));
127 - assert!(should_offload("/vllm.Generate/Generate"));
128 - assert!(should_offload("/vllm.Generate/GenerateStream"));
127 + assert!(should_offload("/vllm.Inference/Generate"));
128 + assert!(should_offload("/vllm.Inference/GenerateStream"));
129 }
130
131 #[test]
tests/v1/engine/test_engine_core_client.py
+7
@@ -328,6 +328,13 @@ def test_apply_ready_response_syncs_block_size():
328 vllm_version="test",
329 world_size=1,
330 data_parallel_size=1,
331 + tensor_parallel_size=1,
332 + pipeline_parallel_size=1,
333 + decode_context_parallel_size=1,
334 + data_parallel_rank=0,
335 + max_num_seqs=256,
336 + max_num_batched_tokens=8192,
337 + instance_id="test-instance",
338 )
339 )
340 client._apply_ready_response(payload)
vllm/v1/engine/__init__.py
+7
@@ -82,6 +82,13 @@ class EngineCoreReadyResponse:
82 vllm_version: str
83 world_size: int
84 data_parallel_size: int
85 + tensor_parallel_size: int
86 + pipeline_parallel_size: int
87 + decode_context_parallel_size: int
88 + data_parallel_rank: int
89 + max_num_seqs: int
90 + max_num_batched_tokens: int
91 + instance_id: str
92 # KV cache capacity (None for encoder-only/attention-free models).
93 kv_cache_size_tokens: int | None = None
94 kv_cache_max_concurrency: float | None = None
vllm/v1/engine/core.py
+26 -16
@@ -1613,6 +1613,31 @@ class EngineCoreProc(EngineCore):
1613 "to send. Please report this issue."
1614 )
1615
1616 + def _make_ready_response(self) -> EngineCoreReadyResponse:
1617 + parallel_config = self.vllm_config.parallel_config
1618 + scheduler_config = self.vllm_config.scheduler_config
1619 + return EngineCoreReadyResponse(
1620 + max_model_len=self.vllm_config.model_config.max_model_len,
1621 + num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0,
1622 + block_size=self.vllm_config.cache_config.block_size,
1623 + dp_stats_address=self.frontend_stats_publish_address,
1624 + dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."),
1625 + vllm_version=VLLM_VERSION,
1626 + world_size=self.vllm_config.parallel_config.world_size,
1627 + data_parallel_size=parallel_config.data_parallel_size,
1628 + kv_cache_size_tokens=self.vllm_config.cache_config.kv_cache_size_tokens,
1629 + kv_cache_max_concurrency=(
1630 + self.vllm_config.cache_config.kv_cache_max_concurrency
1631 + ),
1632 + tensor_parallel_size=parallel_config.tensor_parallel_size,
1633 + pipeline_parallel_size=parallel_config.pipeline_parallel_size,
1634 + decode_context_parallel_size=parallel_config.decode_context_parallel_size,
1635 + data_parallel_rank=self.engine_index,
1636 + max_num_seqs=scheduler_config.max_num_seqs,
1637 + max_num_batched_tokens=scheduler_config.max_num_batched_tokens,
1638 + instance_id=self.vllm_config.instance_id,
1639 + )
1640 +
1641 def process_input_sockets(
1642 self,
1643 input_addresses: list[str],
@@ -1654,22 +1679,7 @@ class EngineCoreProc(EngineCore):
1679
1680 # Register sockets with poller.
1681 poller = zmq.Poller()
1657 - ready_response = EngineCoreReadyResponse(
1658 - max_model_len=self.vllm_config.model_config.max_model_len,
1659 - num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0,
1660 - block_size=self.vllm_config.cache_config.block_size,
1661 - dp_stats_address=self.frontend_stats_publish_address,
1662 - dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."),
1663 - vllm_version=VLLM_VERSION,
1664 - world_size=self.vllm_config.parallel_config.world_size,
1665 - data_parallel_size=self.vllm_config.parallel_config.data_parallel_size,
1666 - kv_cache_size_tokens=(
1667 - self.vllm_config.cache_config.kv_cache_size_tokens
1668 - ),
1669 - kv_cache_max_concurrency=(
1670 - self.vllm_config.cache_config.kv_cache_max_concurrency
1671 - ),
1672 - )
1682 + ready_response = self._make_ready_response()
1683 ready_payload = msgspec.msgpack.encode(ready_response)
1684 for input_socket in input_sockets:
1685 # Send initial message to each input socket - this is required