@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 41f179b57a

[Rust Frontend] Simplify data-parallel size ownership (#52575)

Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Bugen Zhao committed Aug 18, 2026 at 19:06 UTC 41f179b57aa8ab6f634f508128ce1f1efadd0eb1
19 files changed +188 -134
.buildkite/test-amd.yaml
+3
@@ -2417,6 +2417,7 @@ steps:
2417 - vllm/v1/engine/
2418 - vllm/v1/worker/
2419 - tests/utils.py
2420 + - tests/v1/distributed/test_dense_dp_world_size.py
2421 - tests/v1/distributed/test_external_lb_dp.py
2422 - tests/v1/distributed/test_hybrid_lb_dp.py
2423 - tests/v1/distributed/test_internal_lb_dp.py
@@ -2424,6 +2425,7 @@ steps:
2425 commands:
2426 - export VLLM_USE_RUST_FRONTEND=1
2427 - export VLLM_WORKER_MULTIPROC_METHOD=spawn
2428 + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py
2429 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
2430 - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
2431 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
@@ -2882,6 +2884,7 @@ steps:
2884 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
2885 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
2886 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
2887 + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py
2888 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
2889 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
2890 - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
.buildkite/test_areas/distributed.yaml
+1
@@ -154,6 +154,7 @@ steps:
154 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
155 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
156 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
157 + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py
158 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
159 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
160 - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
.buildkite/test_areas/rust_frontend.yaml
+2
@@ -108,6 +108,7 @@ steps:
108 - vllm/v1/engine/
109 - vllm/v1/worker/
110 - tests/utils.py
111 + - tests/v1/distributed/test_dense_dp_world_size.py
112 - tests/v1/distributed/test_external_lb_dp.py
113 - tests/v1/distributed/test_hybrid_lb_dp.py
114 - tests/v1/distributed/test_internal_lb_dp.py
@@ -115,6 +116,7 @@ steps:
116 - export VLLM_USE_RUST_FRONTEND=1
117 - export VLLM_WORKER_MULTIPROC_METHOD=spawn
118 - export NCCL_CUMEM_HOST_ENABLE=0
119 + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py
120 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
121 - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
122 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
rust/src/cmd/src/cli.rs
+1 -2
@@ -474,9 +474,9 @@ impl SharedRuntimeArgs {
474 output_address,
475 engine_start_index,
476 engine_count,
477 + data_parallel_size,
478 ready_timeout,
479 },
479 - data_parallel_size,
480 coordinator_mode: match coordinator_address {
481 Some(address) => CoordinatorMode::External { address },
482 None => CoordinatorMode::None,
@@ -533,7 +533,6 @@ impl SharedRuntimeArgs {
533 local_input_address,
534 local_output_address,
535 },
536 - data_parallel_size: engine_count,
536 coordinator_mode: CoordinatorMode::MaybeInProc,
537 model: self.model,
538 served_model_name: self.served_model_name,
rust/src/cmd/src/cli/tests.rs
+2 -4
@@ -864,7 +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);
867 + assert_eq!(args.into_config().transport_mode.data_parallel_size(), 4);
868 }
869
870 #[test]
@@ -1521,7 +1521,6 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
1521 "<ipc output>",
1522 ),
1523 },
1524 - data_parallel_size: 4,
1524 coordinator_mode: MaybeInProc,
1525 model: "Qwen/Qwen3-0.6B",
1526 served_model_name: [],
@@ -1607,7 +1606,6 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
1606 local_input_address: None,
1607 local_output_address: None,
1608 },
1610 - data_parallel_size: 4,
1609 coordinator_mode: MaybeInProc,
1610 model: "Qwen/Qwen3-0.6B",
1611 served_model_name: [],
@@ -1710,9 +1708,9 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
1708 output_address: "ipc:///tmp/output.sock",
1709 engine_start_index: 3,
1710 engine_count: 1,
1711 + data_parallel_size: 4,
1712 ready_timeout: 600s,
1713 },
1715 - data_parallel_size: 4,
1714 coordinator_mode: External {
1715 address: "tcp://127.0.0.1:7000",
1716 },
rust/src/engine-core-client/src/client.rs
+73 -1
@@ -15,7 +15,7 @@ use tracing::{debug, info, trace};
15
16 use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop};
17 use crate::coordinator::CoordinatorHandle;
18 -use crate::error::{Error, Result};
18 +use crate::error::{Error, Result, bail_invalid_client_config};
19 use crate::protocol::dtype::ModelDtype;
20 use crate::protocol::handshake::EngineCoreReadyResponse;
21 use crate::protocol::lora::LoraRequest;
@@ -68,11 +68,70 @@ pub enum TransportMode {
68 engine_start_index: u32,
69 /// Total number of engines expected to register on this transport.
70 engine_count: usize,
71 + /// Deployment-wide data-parallel size. This may be larger than
72 + /// `engine_count` when a supervisor partitions ranks across frontends.
73 + data_parallel_size: usize,
74 /// Maximum time to wait for all expected engines to register.
75 ready_timeout: Duration,
76 },
77 }
78
79 +impl TransportMode {
80 + /// Return the deployment-wide data-parallel size for this transport.
81 + pub fn data_parallel_size(&self) -> usize {
82 + match self {
83 + Self::HandshakeOwner { engine_count, .. } => *engine_count,
84 + Self::Bootstrapped {
85 + data_parallel_size, ..
86 + } => *data_parallel_size,
87 + }
88 + }
89 +
90 + /// Validate the transport topology before opening any sockets.
91 + pub fn validate(&self) -> Result<()> {
92 + let data_parallel_size = self.data_parallel_size();
93 + if data_parallel_size == 0 {
94 + bail_invalid_client_config!("data parallel size must be at least 1");
95 + }
96 + if data_parallel_size > usize::from(u16::MAX) + 1 {
97 + bail_invalid_client_config!(
98 + "data parallel size ({data_parallel_size}) exceeds the two-byte engine identity limit"
99 + );
100 + }
101 +
102 + match self {
103 + Self::HandshakeOwner { .. } => {}
104 + Self::Bootstrapped {
105 + engine_start_index,
106 + engine_count,
107 + ..
108 + } => {
109 + if *engine_count == 0 {
110 + bail_invalid_client_config!("engine count must be at least 1");
111 + }
112 + let engine_start_index = usize::try_from(*engine_start_index).map_err(|_| {
113 + Error::InvalidClientConfig {
114 + message: "engine start index does not fit usize".to_string(),
115 + }
116 + })?;
117 + let engine_end_index =
118 + engine_start_index.checked_add(*engine_count).ok_or_else(|| {
119 + Error::InvalidClientConfig {
120 + message: "engine start index + engine count overflows".to_string(),
121 + }
122 + })?;
123 + if engine_end_index > data_parallel_size {
124 + bail_invalid_client_config!(
125 + "connected engine range [{engine_start_index}, {engine_end_index}) exceeds data parallel size ({data_parallel_size})"
126 + );
127 + }
128 + }
129 + }
130 +
131 + Ok(())
132 + }
133 +}
134 +
135 /// Which coordinator implementation should be active when one is present for a
136 /// frontend client.
137 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -117,6 +176,11 @@ impl EngineCoreClientConfig {
176 }
177 }
178
179 + /// Validate the client topology before opening any transport sockets.
180 + pub fn validate(&self) -> Result<()> {
181 + self.transport_mode.validate()
182 + }
183 +
184 /// Set the model name used by frontend-side metrics and diagnostics.
185 pub fn with_model_name(mut self, model_name: impl Into<String>) -> Self {
186 self.model_name = model_name.into();
@@ -226,6 +290,7 @@ impl EngineCoreClient {
290 /// handshake. In bootstrapped mode it binds the provided frontend
291 /// sockets and waits for the expected engine registration frames.
292 pub async fn connect(config: EngineCoreClientConfig) -> Result<Self> {
293 + config.validate()?;
294 let connected = match &config.transport_mode {
295 TransportMode::HandshakeOwner {
296 handshake_address,
@@ -261,6 +326,7 @@ impl EngineCoreClient {
326 engine_start_index,
327 engine_count,
328 ready_timeout,
329 + ..
330 } => {
331 if let Some(CoordinatorMode::InProc) = config.coordinator_mode {
332 panic!("cannot use in-process coordinator with bootstrapped transport mode")
@@ -378,6 +444,12 @@ impl EngineCoreClient {
444 self.engines.len()
445 }
446
447 + /// Return the deployment-wide data-parallel size configured for this
448 + /// client.
449 + pub fn data_parallel_size(&self) -> usize {
450 + self.config.transport_mode.data_parallel_size()
451 + }
452 +
453 /// Return the engine-side indices connected to this client.
454 pub fn engine_indices(&self) -> Vec<u32> {
455 self.engines
rust/src/engine-core-client/src/error.rs
+2
@@ -49,6 +49,8 @@ pub enum Error {
49 UnexpectedHandshakeIdentity { expected: Vec<u8>, actual: Vec<u8> },
50 #[error("unexpected startup handshake message: {message}")]
51 UnexpectedHandshakeMessage { message: String },
52 + #[error("invalid engine-core client configuration: {message}")]
53 + InvalidClientConfig { message: String },
54 #[error("unexpected non-control output on coordinator path: {message}")]
55 UnexpectedCoordinatorOutput { message: String },
56 #[error("unexpected output on main dispatcher path: {message}")]
rust/src/engine-core-client/src/mock_engine.rs
+1 -1
@@ -56,7 +56,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
56 dtype: ModelDtype::Float32,
57 vllm_version: "test-vllm-version".to_string(),
58 world_size: 1,
59 - data_parallel_size: 1,
59 + effective_data_parallel_size: 1,
60 tensor_parallel_size: 1,
61 pipeline_parallel_size: 1,
62 decode_context_parallel_size: 1,
rust/src/engine-core-client/src/protocol/handshake.rs
+5 -2
@@ -63,8 +63,11 @@ pub struct EngineCoreReadyResponse {
63 pub vllm_version: String,
64 /// World size (TP * PP) from the parallel config.
65 pub world_size: u64,
66 - /// Data parallelism size from the parallel config.
67 - pub data_parallel_size: u64,
66 + /// Data-parallel size from this EngineCore's effective parallel config.
67 + /// Dense independent-DP ranks are reconfigured to report `1`; the client
68 + /// transport owns the deployment-wide data-parallel size.
69 + #[serde(rename = "data_parallel_size")]
70 + pub effective_data_parallel_size: u64,
71 // Required discovery metadata; EngineCore and client versions must match.
72 /// Tensor-parallel size of this engine.
73 pub tensor_parallel_size: u32,
rust/src/engine-core-client/src/tests/client.rs
+33
@@ -303,6 +303,7 @@ fn bootstrapped_test_config(
303 output_address,
304 engine_start_index: 0,
305 engine_count,
306 + data_parallel_size: engine_count,
307 ready_timeout,
308 },
309 coordinator_mode,
@@ -330,15 +331,47 @@ fn bootstrapped_test_config_with_start_index(
331 );
332 let TransportMode::Bootstrapped {
333 engine_start_index: start,
334 + data_parallel_size,
335 ..
336 } = &mut config.transport_mode
337 else {
338 unreachable!("bootstrapped_test_config returns bootstrapped transport")
339 };
340 *start = engine_start_index;
341 + *data_parallel_size = usize::try_from(engine_start_index)
342 + .expect("test start index fits usize")
343 + .checked_add(engine_count)
344 + .expect("test engine range fits usize");
345 config
346 }
347
348 +#[test]
349 +fn client_config_validates_bootstrapped_dp_range() {
350 + let mut config = bootstrapped_test_config_with_start_index(
351 + "ipc://unused-input".to_string(),
352 + "ipc://unused-output".to_string(),
353 + 1,
354 + 2,
355 + Duration::from_secs(1),
356 + 0,
357 + None,
358 + );
359 + config.validate().expect("frontend may own a subset of global DP ranks");
360 +
361 + let TransportMode::Bootstrapped {
362 + data_parallel_size, ..
363 + } = &mut config.transport_mode
364 + else {
365 + unreachable!("expected bootstrapped transport")
366 + };
367 + *data_parallel_size = 2;
368 + let error = config.validate().expect_err("engine range above DP size must fail");
369 + expect_test::expect![[
370 + "invalid engine-core client configuration: connected engine range [1, 3) exceeds data parallel size (2)"
371 + ]]
372 + .assert_eq(&error.to_string());
373 +}
374 +
375 async fn recv_xpub_message(xpub: &mut XPubSocket) -> Vec<bytes::Bytes> {
376 xpub.recv().await.unwrap().into_vec()
377 }
rust/src/server/examples/external_engine_openai_qwen.rs
-1
@@ -58,7 +58,6 @@ async fn main() -> Result<()> {
58 local_input_address: None,
59 local_output_address: None,
60 },
61 - data_parallel_size: args.engine_count,
61 coordinator_mode: CoordinatorMode::MaybeInProc,
62 model: args.model,
63 served_model_name: vec![],
rust/src/server/src/config.rs
+1 -42
@@ -161,8 +161,6 @@ 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,
164 /// Requested frontend-side coordinator behavior.
165 pub coordinator_mode: CoordinatorMode,
166 /// Backend model identifier used for engine-core loading.
@@ -239,46 +237,7 @@ impl Config {
237 max_logprobs
238 );
239 }
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 - }
240 + self.transport_mode.validate()?;
241
242 Ok(())
243 }
rust/src/server/src/grpc/control.rs
+1 -1
@@ -43,7 +43,7 @@ impl ControlServiceImpl {
43 pb::ParallelismInfo {
44 tensor_parallel_size: ready.tensor_parallel_size,
45 pipeline_parallel_size: ready.pipeline_parallel_size,
46 - data_parallel_size: self.state.data_parallel_size().min(u32::MAX as usize) as u32,
46 + data_parallel_size: self.client().data_parallel_size().min(u32::MAX as usize) as u32,
47 data_parallel_rank: ready.data_parallel_rank,
48 decode_context_parallel_size: ready.decode_context_parallel_size,
49 }
rust/src/server/src/grpc/tests.rs
+4 -4
@@ -1625,7 +1625,7 @@ async fn control_aggregates_multi_engine_capacity() {
1625 let mut ready_0 = default_ready_response();
1626 ready_0.max_model_len = 8_192;
1627 ready_0.num_gpu_blocks = 10;
1628 - ready_0.data_parallel_size = 2;
1628 + ready_0.effective_data_parallel_size = 2;
1629 ready_0.weight_transfer_backend = Some("nccl".to_string());
1630 ready_0.enable_sleep_mode = true;
1631 ready_0.supports_draft_weight_updates = true;
@@ -1633,7 +1633,7 @@ async fn control_aggregates_multi_engine_capacity() {
1633 let mut ready_1 = default_ready_response();
1634 ready_1.max_model_len = 4_096;
1635 ready_1.num_gpu_blocks = 20;
1636 - ready_1.data_parallel_size = 2;
1636 + ready_1.effective_data_parallel_size = 2;
1637 ready_1.data_parallel_rank = 1;
1638
1639 let engine_tasks = [ready_0, ready_1].map(|ready| {
@@ -1667,7 +1667,7 @@ async fn control_aggregates_multi_engine_capacity() {
1667 Llm::new(client),
1668 Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
1669 );
1670 - let state = AppState::new(vec!["test-model".to_string()], chat).with_data_parallel_size(4);
1670 + let state = AppState::new(vec!["test-model".to_string()], chat);
1671 let service = ControlServiceImpl::new(Arc::new(state));
1672
1673 let server = pb::control_server::Control::get_server_info(
@@ -1684,7 +1684,7 @@ async fn control_aggregates_multi_engine_capacity() {
1684 assert!(rl.weight_transfer_backend.is_empty());
1685 assert!(!rl.sleep_mode_enabled);
1686 assert!(!rl.draft_weight_updates_enabled);
1687 - assert_eq!(server.parallelism.unwrap().data_parallel_size, 4);
1687 + assert_eq!(server.parallelism.unwrap().data_parallel_size, 2);
1688
1689 drop(engine_tasks);
1690 }
rust/src/server/src/lib.rs
-1
@@ -141,7 +141,6 @@ 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)
144 .with_api_keys(config.api_keys.clone())
145 .with_cors(config.cors.clone())
146 .with_profiler(config.profiler.clone()),
rust/src/server/src/routes/tests.rs
+6 -59
@@ -704,14 +704,13 @@ async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router {
704 /// stays alive for the duration of the test.
705 async fn test_dev_mode_app_with_ready(
706 ready_response: vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse,
707 - data_parallel_size: usize,
707 ) -> (axum::Router, MockEngineTask) {
708 let ipc = IpcNamespace::new().expect("create ipc namespace");
709 let handshake_address = ipc.handshake_endpoint();
710 let engine_id = b"engine-world-size".to_vec();
711 let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready(
712 handshake_address.clone(),
714 - engine_id.clone(),
713 + engine_id,
714 ready_response,
715 |_dealer, _push| boxed_test_future(async {}),
716 ));
@@ -729,10 +728,10 @@ async fn test_dev_mode_app_with_ready(
728
729 let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));
730 let app = build_router_with_dev_mode(
732 - Arc::new(
733 - AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
734 - .with_data_parallel_size(data_parallel_size),
735 - ),
731 + Arc::new(AppState::new(
732 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
733 + chat,
734 + )),
735 true,
736 );
737 (app, engine_task)
@@ -6292,7 +6291,7 @@ async fn tokenize_allows_prompts_at_or_above_max_model_len() {
6291 max_model_len: 4,
6292 ..default_ready_response()
6293 };
6295 - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 1).await;
6294 + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;
6295
6296 let (completion_status, completion_json) = post_json(
6297 &mut app,
@@ -6543,58 +6542,6 @@ async fn world_size_endpoint_is_dev_mode_only() {
6542 assert_eq!(response.status(), StatusCode::NOT_FOUND);
6543 }
6544
6546 -#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6547 -#[serial]
6548 -async fn world_size_includes_data_parallelism_by_default() {
6549 - let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse {
6550 - world_size: 2,
6551 - data_parallel_size: 1,
6552 - ..default_ready_response()
6553 - };
6554 - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 4).await;
6555 -
6556 - let response = app
6557 - .call(
6558 - Request::builder()
6559 - .uri("/get_world_size")
6560 - .body(Body::empty())
6561 - .expect("build request"),
6562 - )
6563 - .await
6564 - .expect("call app");
6565 -
6566 - assert_eq!(response.status(), StatusCode::OK);
6567 - let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
6568 - let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
6569 - assert_eq!(json, json!({"world_size": 8}));
6570 -}
6571 -
6572 -#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6573 -#[serial]
6574 -async fn world_size_excludes_data_parallelism_when_include_dp_false() {
6575 - let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse {
6576 - world_size: 2,
6577 - data_parallel_size: 4,
6578 - ..default_ready_response()
6579 - };
6580 - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 4).await;
6581 -
6582 - let response = app
6583 - .call(
6584 - Request::builder()
6585 - .uri("/get_world_size?include_dp=false")
6586 - .body(Body::empty())
6587 - .expect("build request"),
6588 - )
6589 - .await
6590 - .expect("call app");
6591 -
6592 - assert_eq!(response.status(), StatusCode::OK);
6593 - let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
6594 - let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
6595 - assert_eq!(json, json!({"world_size": 2}));
6596 -}
6597 -
6545 // ========================= Profiler route tests =========================
6546
6547 async fn test_profiling_app_with_engine_script<F>(script: F) -> (axum::Router, MockEngineTask)
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 = state.data_parallel_size() as u64;
50 + let dp = client.data_parallel_size() as u64;
51 ws * dp
52 } else {
53 ws
rust/src/server/src/state.rs
-15
@@ -40,8 +40,6 @@ 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,
43 /// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes.
44 api_key_hashes: Vec<ApiKeyHash>,
45 /// Number of in-flight inference requests currently owned by this frontend.
@@ -71,14 +69,12 @@ impl AppState {
69 !served_model_names.is_empty(),
70 "served_model_names must not be empty"
71 );
74 - let data_parallel_size = chat.engine_core_client().engine_count();
72 Self {
73 served_model_names,
74 chat,
75 api_server_options: ApiServerOptions::default(),
76 cors: CorsConfig::default(),
77 server_info: None,
81 - data_parallel_size,
78 api_key_hashes: Vec::new(),
79 server_load: AtomicU64::new(0),
80 lora_manager: LoraManager::new(),
@@ -118,17 +114,6 @@ impl AppState {
114 self
115 }
116
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 -
117 /// Build a `/server_info` response payload.
118 pub(crate) fn server_info_response(
119 &self,
tests/v1/distributed/test_dense_dp_world_size.py new
+52
@@ -0,0 +1,52 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import os
4 +
5 +import requests
6 +
7 +from tests.utils import RemoteOpenAIServer
8 +from vllm.platforms import current_platform
9 +
10 +MODEL_NAME = "Qwen/Qwen3-0.6B"
11 +DP_SIZE = int(os.getenv("DP_SIZE", "2"))
12 +TP_SIZE = int(os.getenv("TP_SIZE", "1"))
13 +
14 +
15 +def test_dense_dp_world_size():
16 + server_args = [
17 + "--dtype",
18 + "bfloat16",
19 + "--max-model-len",
20 + "2048",
21 + "--max-num-seqs",
22 + "128",
23 + "--enforce-eager",
24 + "--data-parallel-size",
25 + str(DP_SIZE),
26 + "--data-parallel-size-local",
27 + str(DP_SIZE),
28 + "--tensor-parallel-size",
29 + str(TP_SIZE),
30 + ]
31 + env_dict = {
32 + "VLLM_SERVER_DEV_MODE": "1",
33 + current_platform.device_control_env_var: ",".join(
34 + str(current_platform.device_id_to_physical_device_id(i))
35 + for i in range(DP_SIZE * TP_SIZE)
36 + ),
37 + }
38 +
39 + with RemoteOpenAIServer(
40 + MODEL_NAME,
41 + server_args,
42 + env_dict=env_dict,
43 + ) as server:
44 + response = requests.get(server.url_for("get_world_size"))
45 + response.raise_for_status()
46 + assert response.json() == {"world_size": TP_SIZE * DP_SIZE}
47 +
48 + response = requests.get(
49 + server.url_for("get_world_size"), params={"include_dp": "false"}
50 + )
51 + response.raise_for_status()
52 + assert response.json() == {"world_size": TP_SIZE}