@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4cb5e746b6

[Rust Frontend]: Add `/get_world_size` route with static parallel size (#44801)

Ashar committed Jun 18, 2026 at 18:40 UTC 4cb5e746b63707ed470f952cfb77778a3dd34400

11 files changed +258 -2

rust/src/engine-core-client/src/client.rs

+18

@@ -414,6 +414,24 @@ impl EngineCoreClient {

414 .expect("engine core client requires at least one engine")

415 }

416

417 + /// Return the world size (TP * PP) from the parallel config, if available.

418 + pub fn world_size(&self) -> u64 {

419 + self.engines

420 + .first()

421 + .expect("engine core client requires at least one engine")

422 + .ready_response

423 + .world_size

424 + }

425 +

426 + /// Return the data parallel size from the parallel config, if available.

427 + pub fn data_parallel_size(&self) -> u64 {

428 + self.engines

429 + .first()

430 + .expect("engine core client requires at least one engine")

431 + .ready_response

432 + .data_parallel_size

433 + }

434 +

435 /// Get the model name associated with this client used for metrics

436 /// labeling.

437 pub fn model_name(&self) -> &str {

rust/src/engine-core-client/src/mock_engine.rs

@@ -52,6 +52,8 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {

52 dp_stats_address: None,

53 dtype: ModelDtype::Float32,

54 vllm_version: "test-vllm-version".to_string(),

55 + world_size: 1,

56 + data_parallel_size: 1,

57 kv_cache_size_tokens: None,

58 kv_cache_max_concurrency: None,

59 }

rust/src/engine-core-client/src/protocol/handshake.rs

@@ -44,6 +44,10 @@ pub struct EngineCoreReadyResponse {

44 pub dtype: ModelDtype,

45 /// Python vLLM version reported by the engine process.

46 pub vllm_version: String,

47 + /// World size (TP * PP) from the parallel config.

48 + pub world_size: u64,

49 + /// Data parallelism size from the parallel config.

50 + pub data_parallel_size: u64,

51 /// Total KV cache capacity in tokens, if reported.

52 pub kv_cache_size_tokens: Option<u64>,

53 /// Maximum achievable request concurrency given the KV cache, if reported.

rust/src/engine-core-client/src/test_utils.rs

+56 -1

@@ -12,7 +12,7 @@ use crate::mock_engine::{

12 MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend,

13 default_ready_response,

14 };

15 -use crate::protocol::handshake::HandshakeInitMessage;

15 +use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage};

17 /// Per-test IPC endpoint namespace backed by a unique temporary directory.

18 ///

@@ -62,6 +62,15 @@ fn test_mock_engine_config() -> MockEngineConfig {

62 }

63 }

65 +fn test_mock_engine_config_with_ready(ready_response: EngineCoreReadyResponse) -> MockEngineConfig {

66 + MockEngineConfig {

67 + local: true,

68 + headless: true,

69 + ready_response,

70 + ..Default::default()

71 + }

72 +}

73 +

74 /// Complete the engine-core handshake and connect mock input/output sockets

75 /// plus optional coordinator sockets.

76 pub async fn setup_mock_engine_sockets(

@@ -147,3 +156,49 @@ where

156 });

157 (shutdown_tx, engine_task)

158 }

159 +

160 +/// Like [`setup_mock_engine`] but uses a custom ready response for the

161 +/// handshake, allowing tests to control `world_size`, `data_parallel_size`,

162 +/// etc.

163 +async fn setup_mock_engine_with_ready(

164 + engine_handshake: String,

165 + engine_id: impl Into<EngineId>,

166 + ready_response: EngineCoreReadyResponse,

167 +) -> (DealerSocket, PushSocket) {

168 + let config = test_mock_engine_config_with_ready(ready_response);

169 + let MockEngineSockets { data_sockets, .. } =

170 + connect_to_frontend(engine_handshake, engine_id, config)

171 + .await

172 + .expect("connect mock engine with custom ready response");

173 + let MockEngineDataSockets { dealer, push } =

174 + data_sockets.into_iter().next().expect("mock engine data socket");

175 + (dealer, push)

176 +}

177 +

178 +/// Like [`spawn_mock_engine_task`] but uses a custom ready response for the

179 +/// handshake, allowing tests to set `world_size` and `data_parallel_size` to

180 +/// non-default values.

181 +pub fn spawn_mock_engine_task_with_ready<F>(

182 + engine_handshake: String,

183 + engine_id: impl Into<EngineId>,

184 + ready_response: EngineCoreReadyResponse,

185 + run: F,

186 +) -> (oneshot::Sender<()>, tokio::task::JoinHandle<()>)

187 +where

188 + F: for<'a> FnOnce(

189 + &'a mut DealerSocket,

190 + &'a mut PushSocket,

191 + ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>

192 + + Send

193 + + 'static,

194 +{

195 + let (shutdown_tx, shutdown_rx) = oneshot::channel();

196 + let engine_id = engine_id.into();

197 + let engine_task = tokio::spawn(async move {

198 + let (mut dealer, mut push) =

199 + setup_mock_engine_with_ready(engine_handshake, engine_id, ready_response).await;

200 + run(&mut dealer, &mut push).await;

201 + let _ = shutdown_rx.await;

202 + });

203 + (shutdown_tx, engine_task)

204 +}

rust/src/engine-core-client/src/tests/python_compat.py

@@ -358,6 +358,8 @@ class EngineCoreReadyResponse:

358 dp_stats_address: str | None

359 dtype: str

360 vllm_version: str

361 + world_size: int

362 + data_parallel_size: int

363 kv_cache_size_tokens: int | None = None

364 kv_cache_max_concurrency: float | None = None

365

@@ -369,6 +371,8 @@ ready_response = EngineCoreReadyResponse(

371 dp_stats_address=None,

372 dtype="float32",

373 vllm_version="0.0.0",

374 + data_parallel_size=1,

375 + world_size=1,

376 )

377

378 print(msgspec.msgpack.encode(request).hex())

rust/src/server/src/routes.rs

@@ -12,6 +12,7 @@ mod server_info;

12 mod sleep;

13 mod tokenize;

14 mod version;

15 +mod world_size;

17 use std::sync::Arc;

@@ -100,6 +101,7 @@ fn build_router_with_options(

101 .route("/resume", post(pause::resume))

102 .route("/is_paused", get(pause::is_paused))

103 .route("/server_info", get(server_info::server_info))

104 + .route("/get_world_size", get(world_size::get_world_size))

105 }

106

107 let enable_request_id_headers = state.api_server_options.enable_request_id_headers;

rust/src/server/src/routes/tests.rs

+112 -1

@@ -23,6 +23,7 @@ use vllm_chat::{

23 ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer,

24 NewChatOutputProcessorOptions,

25 };

26 +use vllm_engine_core_client::mock_engine::default_ready_response;

27 use vllm_engine_core_client::protocol::logprobs::{

28 Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,

29 };

@@ -31,7 +32,9 @@ use vllm_engine_core_client::protocol::{

32 EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,

33 decode_value,

34 };

34 -use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};

35 +use vllm_engine_core_client::test_utils::{

36 + IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready,

37 +};

38 use vllm_engine_core_client::{

39 ENGINE_CORE_DEAD_SENTINEL, EngineCoreClient, EngineCoreClientConfig, EngineId,

40 };

@@ -788,6 +791,45 @@ async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router {

791 )

792 }

793

794 +/// Build a dev-mode router backed by a mock engine using a custom ready

795 +/// response, returning the router and the engine task handle so the engine

796 +/// stays alive for the duration of the test.

797 +async fn test_dev_mode_app_with_ready(

798 + ready_response: vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse,

799 +) -> (axum::Router, MockEngineTask) {

800 + let ipc = IpcNamespace::new().expect("create ipc namespace");

801 + let handshake_address = ipc.handshake_endpoint();

802 + let engine_id = b"engine-world-size".to_vec();

803 +

804 + let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready(

805 + handshake_address.clone(),

806 + engine_id.clone(),

807 + ready_response,

808 + |_dealer, _push| boxed_test_future(async {}),

809 + ));

810 +

811 + let client = EngineCoreClient::connect(

812 + EngineCoreClientConfig::new_single(handshake_address)

813 + .with_model_name("test-model")

814 + .with_local_input_output_addresses(

815 + Some(ipc.input_endpoint()),

816 + Some(ipc.output_endpoint()),

817 + ),

818 + )

819 + .await

820 + .expect("connect client");

821 +

822 + let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new()));

823 + let app = build_router_with_dev_mode(

824 + Arc::new(AppState::new(

825 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],

826 + chat,

827 + )),

828 + true,

829 + );

830 + (app, engine_task)

831 +}

832 +

833 async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {

834 let (chat, engine_task) = test_models_with_engine_outputs_and_backend(

835 b"engine-openai-request-id",

@@ -5876,3 +5918,72 @@ async fn tokenize_chat_continue_final_vs_new_assistant_differs() {

5918 let new_len = new_assistant["tokens"].as_array().unwrap().len();

5919 assert!(new_len > continue_len);

5920 }

5921 +

5922 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

5923 +#[serial]

5924 +async fn world_size_endpoint_is_dev_mode_only() {

5925 + let mut app = test_app().await;

5926 + let response = app

5927 + .call(

5928 + Request::builder()

5929 + .uri("/get_world_size")

5930 + .body(Body::empty())

5931 + .expect("build request"),

5932 + )

5933 + .await

5934 + .expect("call app");

5935 +

5936 + assert_eq!(response.status(), StatusCode::NOT_FOUND);

5937 +}

5938 +

5939 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

5940 +#[serial]

5941 +async fn world_size_includes_data_parallelism_by_default() {

5942 + let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse {

5943 + world_size: 2,

5944 + data_parallel_size: 4,

5945 + ..default_ready_response()

5946 + };

5947 + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;

5948 +

5949 + let response = app

5950 + .call(

5951 + Request::builder()

5952 + .uri("/get_world_size")

5953 + .body(Body::empty())

5954 + .expect("build request"),

5955 + )

5956 + .await

5957 + .expect("call app");

5958 +

5959 + assert_eq!(response.status(), StatusCode::OK);

5960 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

5961 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");

5962 + assert_eq!(json, json!({"world_size": 8}));

5963 +}

5964 +

5965 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

5966 +#[serial]

5967 +async fn world_size_excludes_data_parallelism_when_include_dp_false() {

5968 + let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse {

5969 + world_size: 2,

5970 + data_parallel_size: 4,

5971 + ..default_ready_response()

5972 + };

5973 + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await;

5974 +

5975 + let response = app

5976 + .call(

5977 + Request::builder()

5978 + .uri("/get_world_size?include_dp=false")

5979 + .body(Body::empty())

5980 + .expect("build request"),

5981 + )

5982 + .await

5983 + .expect("call app");

5984 +

5985 + assert_eq!(response.status(), StatusCode::OK);

5986 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

5987 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");

5988 + assert_eq!(json, json!({"world_size": 2}));

5989 +}

rust/src/server/src/routes/world_size.rs new

+54

@@ -0,0 +1,54 @@

1 +use std::sync::Arc;

2 +

3 +use axum::Json;

4 +use axum::extract::{Query, State};

5 +use serde::{Deserialize, Serialize};

6 +

7 +use crate::error::ApiError;

8 +use crate::state::AppState;

9 +

10 +#[derive(Debug, Deserialize)]

11 +pub(crate) struct WorldSizeParams {

12 + /// If true (default), returns the world size including data parallelism

13 + /// (TP * PP * DP). If false, returns the world size without data

14 + /// parallelism (TP * PP).

15 + #[serde(default = "default_true")]

16 + include_dp: bool,

17 +}

18 +

19 +const fn default_true() -> bool {

20 + true

21 +}

22 +

23 +#[derive(Serialize)]

24 +pub(crate) struct WorldSizeResponse {

25 + world_size: u64,

26 +}

27 +

28 +/// Get the world size from the parallel config.

29 +///

30 +/// Currently reads static values captured during the engine startup handshake.

31 +///

32 +/// TODO: If the world size can change at runtime (e.g. elastic EP scaling,

33 +/// DP rank recovery), this should be switched to either:

34 +/// - A `call_utility("get_world_size", (include_dp,))` RPC to the Python

35 +/// engine for live values (simple, adds one ZMQ round-trip per request), or

36 +/// - A push-based approach where the engine sends config updates via the

37 +/// output stream into shared state (zero per-request overhead, more complex).

38 +pub async fn get_world_size(

39 + State(state): State<Arc<AppState>>,

40 + Query(params): Query<WorldSizeParams>,

41 +) -> Result<Json<WorldSizeResponse>, ApiError> {

42 + let client = state.engine_core_client();

43 +

44 + let ws = client.world_size();

45 +

46 + let world_size = if params.include_dp {

47 + let dp = client.data_parallel_size();

48 + ws * dp

49 + } else {

50 + ws

51 + };

52 +

53 + Ok(Json(WorldSizeResponse { world_size }))

54 +}

tests/v1/engine/test_engine_core_client.py

@@ -255,6 +255,8 @@ def test_apply_ready_response_syncs_block_size():

255 dp_stats_address=None,

256 dtype="bfloat16",

257 vllm_version="test",

258 + world_size=1,

259 + data_parallel_size=1,

260 )

261 )

262 client._apply_ready_response(payload)

vllm/v1/engine/__init__.py

@@ -78,6 +78,8 @@ class EngineCoreReadyResponse:

78 dp_stats_address: str | None

79 dtype: str

80 vllm_version: str

81 + world_size: int

82 + data_parallel_size: int

83 # KV cache capacity (None for encoder-only/attention-free models).

84 kv_cache_size_tokens: int | None = None

85 kv_cache_max_concurrency: float | None = None

vllm/v1/engine/core.py

@@ -1527,6 +1527,8 @@ class EngineCoreProc(EngineCore):

1527 dp_stats_address=self.frontend_stats_publish_address,

1528 dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."),

1529 vllm_version=VLLM_VERSION,

1530 + world_size=self.vllm_config.parallel_config.world_size,

1531 + data_parallel_size=self.vllm_config.parallel_config.data_parallel_size,

1532 kv_cache_size_tokens=(

1533 self.vllm_config.cache_config.kv_cache_size_tokens

1534 ),