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

[Rust Frontend][gRPC] Add explicit data-parallel rank routing (#51178)

Signed-off-by: Connor Carpenter <connorc@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Connor Carpenter committed Aug 10, 2026 at 15:05 UTC fa722b9f0188179a66728f2f882e9d5bc92c9b0c
19 files changed +215 -60
rust/src/cmd/src/cli.rs
+8
@@ -457,6 +457,7 @@ impl SharedRuntimeArgs {
457 coordinator_address: Option<String>,
458 engine_start_index: u32,
459 engine_count: usize,
460 + data_parallel_size: usize,
461 ) -> Config {
462 let ready_timeout = self.ready_timeout();
463 let shutdown_timeout = self.shutdown_timeout();
@@ -474,6 +475,7 @@ impl SharedRuntimeArgs {
475 engine_count,
476 ready_timeout,
477 },
478 + data_parallel_size,
479 coordinator_mode: match coordinator_address {
480 Some(address) => CoordinatorMode::External { address },
481 None => CoordinatorMode::None,
@@ -530,6 +532,7 @@ impl SharedRuntimeArgs {
532 local_input_address,
533 local_output_address,
534 },
535 + data_parallel_size: engine_count,
536 coordinator_mode: CoordinatorMode::MaybeInProc,
537 model: self.model,
538 served_model_name: self.served_model_name,
@@ -661,6 +664,9 @@ pub struct FrontendArgs {
664 /// Total number of data-parallel engines expected for this frontend.
665 #[arg(long, default_value_t = 1)]
666 pub engine_count: usize,
667 + /// Deployment-wide configured data-parallel size. Defaults to engine count.
668 + #[arg(long)]
669 + pub data_parallel_size: Option<usize>,
670
671 /// Shared frontend arguments as one JSON object.
672 #[arg(long = "args-json", value_parser = parse_runtime_args_json, value_name = "JSON")]
@@ -670,6 +676,7 @@ pub struct FrontendArgs {
676 impl FrontendArgs {
677 /// Convert the CLI arguments into the OpenAI server's runtime config.
678 pub fn into_config(self) -> Config {
679 + let data_parallel_size = self.data_parallel_size.unwrap_or(self.engine_count);
680 self.runtime.into_bootstrapped_config(
681 self.listen_fd,
682 self.input_address,
@@ -677,6 +684,7 @@ impl FrontendArgs {
684 self.coordinator_address,
685 self.engine_start_index,
686 self.engine_count,
687 + data_parallel_size,
688 )
689 }
690 }
rust/src/cmd/src/cli/tests.rs
+9
@@ -785,6 +785,7 @@ fn frontend_args_accept_json() {
785 ),
786 engine_start_index: 0,
787 engine_count: 1,
788 + data_parallel_size: None,
789 runtime: SharedRuntimeArgs {
790 model: "Qwen/Qwen3-0.6B",
791 engine_ready_timeout_secs: 600,
@@ -846,6 +847,8 @@ fn frontend_args_json_applies_defaults() {
847 "ipc:///tmp/input.sock",
848 "--output-address",
849 "ipc:///tmp/output.sock",
850 + "--engine-count",
851 + "4",
852 "--args-json",
853 r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
854 ])
@@ -861,6 +864,7 @@ fn frontend_args_json_applies_defaults() {
864 assert_eq!(args.runtime.renderer, RendererSelection::Auto);
865 assert_eq!(args.runtime.max_logprobs, None);
866 assert_eq!(args.runtime.shutdown_timeout, 0);
867 + assert_eq!(args.into_config().data_parallel_size, 4);
868 }
869
870 #[test]
@@ -1517,6 +1521,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
1521 "<ipc output>",
1522 ),
1523 },
1524 + data_parallel_size: 4,
1525 coordinator_mode: MaybeInProc,
1526 model: "Qwen/Qwen3-0.6B",
1527 served_model_name: [],
@@ -1602,6 +1607,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
1607 local_input_address: None,
1608 local_output_address: None,
1609 },
1610 + data_parallel_size: 4,
1611 coordinator_mode: MaybeInProc,
1612 model: "Qwen/Qwen3-0.6B",
1613 served_model_name: [],
@@ -1685,6 +1691,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
1691 "3",
1692 "--engine-count",
1693 "1",
1694 + "--data-parallel-size",
1695 + "4",
1696 "--args-json",
1697 r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
1698 ])
@@ -1704,6 +1712,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
1712 engine_count: 1,
1713 ready_timeout: 600s,
1714 },
1715 + data_parallel_size: 4,
1716 coordinator_mode: External {
1717 address: "tcp://127.0.0.1:7000",
1718 },
rust/src/engine-core-client/src/client.rs
-9
@@ -451,15 +451,6 @@ impl EngineCoreClient {
451 .world_size
452 }
453
454 - /// Return the data parallel size from the parallel config, if available.
455 - pub fn data_parallel_size(&self) -> u64 {
456 - self.engines
457 - .first()
458 - .expect("engine core client requires at least one engine")
459 - .ready_response
460 - .data_parallel_size
461 - }
462 -
454 /// Get the model name associated with this client used for metrics
455 /// labeling.
456 pub fn model_name(&self) -> &str {
rust/src/engine-core-client/src/client/state.rs
+22 -18
@@ -161,15 +161,16 @@ impl RequestRegistry {
161
162 fn choose_engine_for_request(&mut self, data_parallel_rank: Option<u32>) -> Result<EngineId> {
163 if let Some(rank) = data_parallel_rank {
164 - // Route to the engine at the specified rank index.
165 - let engine_id = EngineId::from_engine_index(rank);
166 - return self
167 - .routing_per_engine
168 - .contains_key(&engine_id)
169 - .then_some(engine_id)
164 + let engine_id = u16::try_from(rank).ok().map(EngineId::from_engine_index);
165 + return engine_id
166 + .filter(|engine_id| self.routing_per_engine.contains_key(engine_id))
167 .ok_or_else(|| Error::InvalidDataParallelRank {
168 rank,
172 - num_engines: self.routing_per_engine.len() as u32,
169 + connected_ranks: self
170 + .routing_per_engine
171 + .keys()
172 + .filter_map(EngineId::engine_index)
173 + .collect(),
174 });
175 }
176
@@ -343,6 +344,9 @@ impl RequestRegistry {
344 }
345
346 fn apply_scheduler_counts(&mut self, engine_index: u32, next: EngineLoadSnapshot) -> bool {
347 + let Ok(engine_index) = u16::try_from(engine_index) else {
348 + return false;
349 + };
350 let engine_id = EngineId::from_engine_index(engine_index);
351 let Some(state) = self.routing_per_engine.get_mut(&engine_id) else {
352 return false;
@@ -840,26 +844,26 @@ mod tests {
844 error,
845 crate::error::Error::InvalidDataParallelRank {
846 rank: 2,
843 - num_engines: 2,
844 - }
847 + connected_ranks,
848 + } if connected_ranks == vec![0, 1]
849 ));
850 }
851
852 #[test]
849 - fn register_with_rank_on_single_engine_only_accepts_zero() {
850 - let engine_0 = EngineId::from_engine_index(0);
851 - let mut registry = RequestRegistry::new(&[connected_engine(engine_0.clone())]);
853 + fn register_with_rank_uses_global_engine_identity() {
854 + let engine_3 = EngineId::from_engine_index(3);
855 + let mut registry = RequestRegistry::new(&[connected_engine(engine_3.clone())]);
856
853 - let (chosen, _) = registry.register("req-ok".to_string(), None, Some(0)).unwrap();
854 - assert_eq!(chosen, engine_0);
857 + let (chosen, _) = registry.register("req-ok".to_string(), None, Some(3)).unwrap();
858 + assert_eq!(chosen, engine_3);
859
856 - let error = registry.register("req-bad".to_string(), None, Some(1)).unwrap_err();
860 + let error = registry.register("req-bad".to_string(), None, Some(0)).unwrap_err();
861 assert!(matches!(
862 error,
863 crate::error::Error::InvalidDataParallelRank {
860 - rank: 1,
861 - num_engines: 1,
862 - }
864 + rank: 0,
865 + connected_ranks,
866 + } if connected_ranks == vec![3]
867 ));
868 }
869
rust/src/engine-core-client/src/error.rs
+7 -2
@@ -68,8 +68,13 @@ pub enum Error {
68 ControlClosed { message: String },
69 #[error("request `{request_id}` is already in flight")]
70 DuplicateRequestId { request_id: String },
71 - #[error("data parallel rank {rank} is out of range for {num_engines} engine(s)")]
72 - InvalidDataParallelRank { rank: u32, num_engines: u32 },
71 + #[error(
72 + "data parallel rank {rank} is not connected to this frontend; connected ranks: {connected_ranks:?}"
73 + )]
74 + InvalidDataParallelRank {
75 + rank: u32,
76 + connected_ranks: Vec<u32>,
77 + },
78 #[error("engine-core output dispatcher closed: {message}")]
79 DispatcherClosed { message: String },
80 #[error("engine-core client is closed: {message}")]
rust/src/engine-core-client/src/transport.rs
+22 -4
@@ -66,8 +66,8 @@ impl EngineId {
66
67 /// Construct an engine id from the Python-compatible engine index encoding
68 /// (two-byte little-endian).
69 - pub fn from_engine_index(value: u32) -> Self {
70 - Self(Bytes::copy_from_slice(&(value as u16).to_le_bytes()))
69 + pub fn from_engine_index(value: u16) -> Self {
70 + Self(Bytes::copy_from_slice(&value.to_le_bytes()))
71 }
72 }
73
@@ -336,6 +336,22 @@ pub async fn connect_bootstrapped(
336 if engine_count == 0 {
337 bail_unexpected_handshake_message!("expected engine_count >= 1");
338 }
339 + let engine_start_index =
340 + u16::try_from(engine_start_index).map_err(|_| Error::UnexpectedHandshakeMessage {
341 + message: "engine_start_index exceeds the two-byte engine identity limit".to_string(),
342 + })?;
343 + let engine_end_index =
344 + usize::from(engine_start_index).checked_add(engine_count).ok_or_else(|| {
345 + Error::UnexpectedHandshakeMessage {
346 + message: "engine_start_index + engine_count overflows".to_string(),
347 + }
348 + })?;
349 + if engine_end_index > usize::from(u16::MAX) + 1 {
350 + return Err(Error::UnexpectedHandshakeMessage {
351 + message: "engine_start_index + engine_count exceeds the two-byte engine identity limit"
352 + .to_string(),
353 + });
354 + }
355
356 let mut input_socket = RouterSocket::new();
357 let input_address = input_socket.bind(input_address).await?.to_string();
@@ -345,8 +361,10 @@ pub async fn connect_bootstrapped(
361
362 let engines = wait_for_input_registrations(
363 &mut input_socket,
348 - (0..engine_count)
349 - .map(|offset| EngineId::from_engine_index(engine_start_index + offset as u32)),
364 + (0..engine_count).map(|offset| {
365 + let offset = u16::try_from(offset).expect("validated engine offset fits u16");
366 + EngineId::from_engine_index(engine_start_index + offset)
367 + }),
368 ready_timeout,
369 )
370 .await?;
rust/src/mock-engine/src/lib.rs
+7 -8
@@ -48,10 +48,13 @@ pub struct Opt {
48 }
49
50 /// Run one mock engine until shutdown or transport failure.
51 -async fn run_engine(engine_index: u32, opt: Opt, shutdown: CancellationToken) -> Result<()> {
51 +async fn run_engine(engine_index: usize, opt: Opt, shutdown: CancellationToken) -> Result<()> {
52 + let engine_index =
53 + u16::try_from(engine_index).context("mock engine index exceeds two-byte identity limit")?;
54 + let engine_id = EngineId::from_engine_index(engine_index);
55 let MockEngineSockets { data_sockets, .. } = connect_to_frontend(
56 &opt.handshake_address,
54 - EngineId::from_engine_index(engine_index),
57 + engine_id,
58 MockEngineConfig::default(),
59 )
60 .await
@@ -71,7 +74,7 @@ async fn run_engine(engine_index: u32, opt: Opt, shutdown: CancellationToken) ->
74 ));
75 // Engine loop: input_rx -> engine logic -> output_tx
76 let mut engine_loop = tokio::spawn(engine::run_engine_loop(
74 - engine_index,
77 + u32::from(engine_index),
78 opt,
79 input_rx,
80 output_tx,
@@ -111,11 +114,7 @@ pub async fn run(opt: Opt, shutdown: CancellationToken) -> Result<()> {
114
115 let mut engines = JoinSet::new();
116 for engine_index in 0..opt.engine_count {
114 - engines.spawn(run_engine(
115 - engine_index as u32,
116 - opt.clone(),
117 - shutdown.clone(),
118 - ));
117 + engines.spawn(run_engine(engine_index, opt.clone(), shutdown.clone()));
118 }
119
120 tokio::select! {
rust/src/server/examples/external_engine_openai_qwen.rs
+1
@@ -58,6 +58,7 @@ async fn main() -> Result<()> {
58 local_input_address: None,
59 local_output_address: None,
60 },
61 + data_parallel_size: args.engine_count,
62 coordinator_mode: CoordinatorMode::MaybeInProc,
63 model: args.model,
64 served_model_name: vec![],
rust/src/server/src/config.rs
+42
@@ -161,6 +161,8 @@ impl TlsConfig {
161 pub struct Config {
162 /// Frontend-to-engine transport setup.
163 pub transport_mode: TransportMode,
164 + /// Deployment-wide data-parallel size retained by the frontend.
165 + pub data_parallel_size: usize,
166 /// Requested frontend-side coordinator behavior.
167 pub coordinator_mode: CoordinatorMode,
168 /// Backend model identifier used for engine-core loading.
@@ -237,6 +239,46 @@ impl Config {
239 max_logprobs
240 );
241 }
242 + if self.data_parallel_size == 0 {
243 + bail!("data parallel size must be at least 1");
244 + }
245 + if self.data_parallel_size > usize::from(u16::MAX) + 1 {
246 + bail!(
247 + "data parallel size ({}) exceeds the two-byte engine identity limit",
248 + self.data_parallel_size
249 + );
250 + }
251 + match &self.transport_mode {
252 + TransportMode::HandshakeOwner { engine_count, .. } => {
253 + if *engine_count != self.data_parallel_size {
254 + bail!(
255 + "managed frontend engine count ({engine_count}) must equal data parallel size ({})",
256 + self.data_parallel_size
257 + );
258 + }
259 + }
260 + TransportMode::Bootstrapped {
261 + engine_start_index,
262 + engine_count,
263 + ..
264 + } => {
265 + if *engine_count == 0 {
266 + bail!("engine count must be at least 1");
267 + }
268 + let engine_start_index = usize::try_from(*engine_start_index)
269 + .map_err(|_| anyhow::anyhow!("engine start index does not fit usize"))?;
270 + let engine_end_index =
271 + engine_start_index.checked_add(*engine_count).ok_or_else(|| {
272 + anyhow::anyhow!("engine start index + engine count overflows")
273 + })?;
274 + if engine_end_index > self.data_parallel_size {
275 + bail!(
276 + "connected engine range [{engine_start_index}, {engine_end_index}) exceeds data parallel size {}",
277 + self.data_parallel_size
278 + );
279 + }
280 + }
281 + }
282
283 Ok(())
284 }
rust/src/server/src/grpc/control.rs
+1 -1
@@ -31,7 +31,7 @@ impl ControlServiceImpl {
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,
34 + data_parallel_size: self.state.data_parallel_size().min(u32::MAX as usize) as u32,
35 data_parallel_rank: ready.data_parallel_rank,
36 decode_context_parallel_size: ready.decode_context_parallel_size,
37 }
rust/src/server/src/grpc/inference.rs
+25 -2
@@ -22,6 +22,8 @@ use crate::state::AppState;
22
23 pub(crate) type InferenceGrpcService = InferenceServer<InferenceServiceImpl>;
24
25 +const DATA_PARALLEL_RANK_METADATA_KEY: &str = "x-data-parallel-rank";
26 +
27 /// gRPC inference service backed by the shared application state.
28 pub struct InferenceServiceImpl {
29 state: Arc<AppState>,
@@ -41,6 +43,7 @@ impl InferenceServiceImpl {
43 async fn prepare_request(
44 &self,
45 mut proto_request: pb::GenerateRequest,
46 + data_parallel_rank: Option<u32>,
47 stream: bool,
48 rpc: &'static str,
49 ) -> Result<PreparedGrpcRequest, Status> {
@@ -62,6 +65,7 @@ impl InferenceServiceImpl {
65 rpc,
66 %model,
67 media_count,
68 + ?data_parallel_rank,
69 );
70 info!(parent: &request_span, "gRPC inference request received");
71
@@ -70,6 +74,7 @@ impl InferenceServiceImpl {
74 let mut text_request =
75 convert::to_text_request(proto_request, stream, self.state.served_model_names())?;
76 text_request.arrival_time = Some(arrival_time);
77 + text_request.data_parallel_rank = data_parallel_rank;
78
79 let media = convert::media_parts_from_request(media)?;
80 if !media.is_empty() {
@@ -115,6 +120,20 @@ impl InferenceServiceImpl {
120 }
121 }
122
123 +fn data_parallel_rank_from_metadata(
124 + request: &Request<pb::GenerateRequest>,
125 +) -> Result<Option<u32>, Status> {
126 + let Some(value) = request.metadata().get(DATA_PARALLEL_RANK_METADATA_KEY) else {
127 + return Ok(None);
128 + };
129 + let value = value.to_str().map_err(|_| {
130 + Status::invalid_argument("x-data-parallel-rank metadata must be an unsigned 32-bit integer")
131 + })?;
132 + value.trim().parse::<u32>().map(Some).map_err(|_| {
133 + Status::invalid_argument("x-data-parallel-rank metadata must be an unsigned 32-bit integer")
134 + })
135 +}
136 +
137 #[tonic::async_trait]
138 impl pb::inference_server::Inference for InferenceServiceImpl {
139 type GenerateStreamStream =
@@ -125,13 +144,14 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
144 &self,
145 request: Request<pb::GenerateRequest>,
146 ) -> Result<Response<pb::GenerateResponse>, Status> {
147 + let data_parallel_rank = data_parallel_rank_from_metadata(&request)?;
148 let proto_req = request.into_inner();
149 let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref());
150 let PreparedGrpcRequest {
151 text_request,
152 request_span,
153 started_at,
134 - } = self.prepare_request(proto_req, false, "Generate").await?;
154 + } = self.prepare_request(proto_req, data_parallel_rank, false, "Generate").await?;
155
156 let stream = self
157 .state
@@ -186,13 +206,16 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
206 &self,
207 request: Request<pb::GenerateRequest>,
208 ) -> Result<Response<Self::GenerateStreamStream>, Status> {
209 + let data_parallel_rank = data_parallel_rank_from_metadata(&request)?;
210 let proto_req = request.into_inner();
211 let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref());
212 let PreparedGrpcRequest {
213 text_request,
214 request_span,
215 started_at,
195 - } = self.prepare_request(proto_req, true, "GenerateStream").await?;
216 + } = self
217 + .prepare_request(proto_req, data_parallel_rank, true, "GenerateStream")
218 + .await?;
219
220 let stream = self
221 .state
rust/src/server/src/grpc/tests.rs
+37 -5
@@ -843,6 +843,37 @@ async fn unary_generate_missing_prompt_returns_invalid_argument() {
843 server_task.abort();
844 }
845
846 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
847 +#[serial]
848 +async fn unary_generate_unconnected_data_parallel_rank_returns_invalid_argument() {
849 + let (mut client, server_task, _engine_task) = grpc_test_server(
850 + EngineId::from_engine_index(3),
851 + default_stream_output_specs(),
852 + )
853 + .await;
854 +
855 + let mut request = tonic::Request::new(pb::GenerateRequest {
856 + request_id: "test-unconnected-dp-rank".to_string(),
857 + model: "test-model".to_string(),
858 + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
859 + ..Default::default()
860 + });
861 + request.metadata_mut().insert(
862 + "x-data-parallel-rank",
863 + "0".parse().expect("valid metadata value"),
864 + );
865 +
866 + let status = client
867 + .generate(request)
868 + .await
869 + .expect_err("rank 0 should not select globally ranked engine 3");
870 +
871 + assert_eq!(status.code(), tonic::Code::InvalidArgument);
872 + assert!(status.message().contains("connected ranks: [3]"));
873 +
874 + server_task.abort();
875 +}
876 +
877 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
878 #[serial]
879 async fn unary_generate_min_tokens_above_max_tokens_returns_invalid_argument() {
@@ -1523,7 +1554,9 @@ async fn control_aggregates_multi_engine_capacity() {
1554 ready_1.data_parallel_rank = 1;
1555
1556 let engine_tasks = [ready_0, ready_1].map(|ready| {
1526 - let engine_id = EngineId::from_engine_index(ready.data_parallel_rank);
1557 + let engine_id = EngineId::from_engine_index(
1558 + ready.data_parallel_rank.try_into().expect("test rank fits engine identity"),
1559 + );
1560 MockEngineTask::new(spawn_mock_engine_task_with_ready(
1561 handshake_address.clone(),
1562 engine_id,
@@ -1551,10 +1584,8 @@ async fn control_aggregates_multi_engine_capacity() {
1584 Llm::new(client),
1585 Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
1586 );
1554 - let service = ControlServiceImpl::new(Arc::new(AppState::new(
1555 - vec!["test-model".to_string()],
1556 - chat,
1557 - )));
1587 + let state = AppState::new(vec!["test-model".to_string()], chat).with_data_parallel_size(4);
1588 + let service = ControlServiceImpl::new(Arc::new(state));
1589
1590 let server = pb::control_server::Control::get_server_info(
1591 &service,
@@ -1565,6 +1596,7 @@ async fn control_aggregates_multi_engine_capacity() {
1596 .into_inner();
1597 assert_eq!(server.max_model_len, 4_096);
1598 assert_eq!(server.total_kv_blocks, 30);
1599 + assert_eq!(server.parallelism.unwrap().data_parallel_size, 4);
1600
1601 drop(engine_tasks);
1602 }
rust/src/server/src/lib.rs
+1
@@ -141,6 +141,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
141 .with_model_path(config.model.clone())
142 .with_api_server_options(config.api_server_options)
143 .with_server_info(ServerInfoSnapshot::from_config(config))
144 + .with_data_parallel_size(config.data_parallel_size)
145 .with_api_keys(config.api_keys.clone())
146 .with_cors(config.cors.clone())
147 .with_profiler(config.profiler.clone()),
rust/src/server/src/routes/tests.rs
+9 -9
@@ -725,11 +725,11 @@ async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router {
725 /// stays alive for the duration of the test.
726 async fn test_dev_mode_app_with_ready(
727 ready_response: vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse,
728 + data_parallel_size: usize,
729 ) -> (axum::Router, MockEngineTask) {
730 let ipc = IpcNamespace::new().expect("create ipc namespace");
731 let handshake_address = ipc.handshake_endpoint();
732 let engine_id = b"engine-world-size".to_vec();
732 -
733 let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready(
734 handshake_address.clone(),
735 engine_id.clone(),
@@ -750,10 +750,10 @@ async fn test_dev_mode_app_with_ready(
750
751 let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
752 let app = build_router_with_dev_mode(
753 - Arc::new(AppState::new(
754 - vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
755 - chat,
756 - )),
753 + Arc::new(
754 + AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
755 + .with_data_parallel_size(data_parallel_size),
756 + ),
757 true,
758 );
759 (app, engine_task)
@@ -6313,7 +6313,7 @@ async fn tokenize_allows_prompts_at_or_above_max_model_len() {
6313 max_model_len: 4,
6314 ..default_ready_response()
6315 };
6316 - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;
6316 + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 1).await;
6317
6318 let (completion_status, completion_json) = post_json(
6319 &mut app,
@@ -6569,10 +6569,10 @@ async fn world_size_endpoint_is_dev_mode_only() {
6569 async fn world_size_includes_data_parallelism_by_default() {
6570 let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse {
6571 world_size: 2,
6572 - data_parallel_size: 4,
6572 + data_parallel_size: 1,
6573 ..default_ready_response()
6574 };
6575 - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;
6575 + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 4).await;
6576
6577 let response = app
6578 .call(
@@ -6598,7 +6598,7 @@ async fn world_size_excludes_data_parallelism_when_include_dp_false() {
6598 data_parallel_size: 4,
6599 ..default_ready_response()
6600 };
6601 - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;
6601 + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 4).await;
6602
6603 let response = app
6604 .call(
rust/src/server/src/routes/world_size.rs
+1 -1
@@ -47,7 +47,7 @@ pub async fn get_world_size(
47 let ws = client.world_size();
48
49 let world_size = if params.include_dp {
50 - let dp = client.data_parallel_size();
50 + let dp = state.data_parallel_size() as u64;
51 ws * dp
52 } else {
53 ws
rust/src/server/src/state.rs
+15
@@ -40,6 +40,8 @@ pub struct AppState {
40 pub cors: CorsConfig,
41 /// Runtime server information returned by `/server_info`, when available.
42 server_info: Option<ServerInfoSnapshot>,
43 + /// Deployment-wide data-parallel size retained by the frontend.
44 + data_parallel_size: usize,
45 /// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes.
46 api_key_hashes: Vec<ApiKeyHash>,
47 /// Number of in-flight inference requests currently owned by this frontend.
@@ -69,12 +71,14 @@ impl AppState {
71 !served_model_names.is_empty(),
72 "served_model_names must not be empty"
73 );
74 + let data_parallel_size = chat.engine_core_client().engine_count();
75 Self {
76 served_model_names,
77 chat,
78 api_server_options: ApiServerOptions::default(),
79 cors: CorsConfig::default(),
80 server_info: None,
81 + data_parallel_size,
82 api_key_hashes: Vec::new(),
83 server_load: AtomicU64::new(0),
84 lora_manager: LoraManager::new(),
@@ -114,6 +118,17 @@ impl AppState {
118 self
119 }
120
121 + /// Set the deployment-wide data-parallel size reported by frontend APIs.
122 + pub(crate) fn with_data_parallel_size(mut self, size: usize) -> Self {
123 + self.data_parallel_size = size;
124 + self
125 + }
126 +
127 + /// Return the deployment-wide data-parallel size.
128 + pub(crate) fn data_parallel_size(&self) -> usize {
129 + self.data_parallel_size
130 + }
131 +
132 /// Build a `/server_info` response payload.
133 pub(crate) fn server_info_response(
134 &self,
rust/src/text/src/error.rs
+4 -1
@@ -62,7 +62,10 @@ impl Error {
62 | Self::InvalidRepetitionDetection { .. }
63 // An empty tokenized prompt detected later, at request prepare
64 // time, surfaces through the transparent Llm wrapper.
65 - | Self::Llm(LlmError::EmptyPromptTokenIds { .. }) => true,
65 + | Self::Llm(LlmError::EmptyPromptTokenIds { .. })
66 + | Self::Llm(LlmError::EngineCoreClient(
67 + EngineCoreError::InvalidDataParallelRank { .. },
68 + )) => true,
69 _ => false,
70 }
71 }
vllm/entrypoints/cli/serve.py
+1
@@ -352,6 +352,7 @@ def run_multi_api_server(args: argparse.Namespace):
352 output_address=addresses.outputs[0],
353 engine_start_index=expected_engine_start_index,
354 engine_count=expected_engine_count,
355 + data_parallel_size=parallel_config.data_parallel_size,
356 stats_update_address=stats_update_address,
357 )
358 else:
vllm/v1/utils.py
+3
@@ -339,6 +339,7 @@ class RustFrontendProcessManager:
339 output_address: str,
340 engine_start_index: int,
341 engine_count: int,
342 + data_parallel_size: int,
343 stats_update_address: str | None = None,
344 ):
345 import os
@@ -360,6 +361,8 @@ class RustFrontendProcessManager:
361 str(engine_start_index),
362 "--engine-count",
363 str(engine_count),
364 + "--data-parallel-size",
365 + str(data_parallel_size),
366 ]
367 if stats_update_address is not None:
368 cmd.extend(["--coordinator-address", stats_update_address])