[Rust Frontend] Expose local DP size in gRPC Control metadata (#57116)
Signed-off-by: Alec Flowers <aflowers@nvidia.com> Co-authored-by: Codex <noreply@openai.com>
Alec committed
Sep 16, 2026 at 11:28 UTC
9639cbde04f58c185ea6e5122dbd2039908abb7a
3 files changed
+94
-77
rust/proto/control.proto
+1
@@ -68,6 +68,7 @@ message ParallelismInfo {
68
uint32 data_parallel_rank = 4;
69
uint32 decode_context_parallel_size = 5;
70
uint64 world_size = 6;
71
+ uint32 data_parallel_size_local = 7;
72
}
73
74
message GetModelInfoRequest {}
rust/src/server/src/grpc/control.rs
+1
@@ -50,6 +50,7 @@ impl ControlServiceImpl {
50
data_parallel_rank: ready.data_parallel_rank,
51
decode_context_parallel_size: ready.decode_context_parallel_size,
52
world_size: ready.world_size,
53
+ data_parallel_size_local: self.client().engine_count() as u32,
54
}
55
}
56
rust/src/server/src/grpc/tests.rs
+92
-77
@@ -31,7 +31,7 @@ use vllm_chat::{
31
};
32
use vllm_engine_core_client::mock_engine::{
33
DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS,
34
- default_ready_response,
34
+ MockEngineConfig, connect_to_bootstrapped_frontend, default_ready_response,
35
};
36
use vllm_engine_core_client::protocol::decode_value;
37
use vllm_engine_core_client::protocol::dtype::TensorDtype;
@@ -1948,6 +1948,7 @@ async fn control_reports_server_and_model_info() {
1948
assert_eq!(parallelism.tensor_parallel_size, 1);
1949
assert_eq!(parallelism.pipeline_parallel_size, 1);
1950
assert_eq!(parallelism.data_parallel_size, 1);
1951
+ assert_eq!(parallelism.data_parallel_size_local, 1);
1952
assert_eq!(parallelism.data_parallel_rank, 0);
1953
assert_eq!(parallelism.decode_context_parallel_size, 1);
1954
let rl = server.rl_capabilities.expect("RL capabilities");
@@ -2191,84 +2192,98 @@ async fn control_forwards_weight_update_without_pause_guard() {
2192
2193
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2194
async fn control_aggregates_multi_engine_capacity() {
2194
- let ipc = IpcNamespace::new().expect("create ipc namespace");
2195
- let handshake_address = ipc.handshake_endpoint();
2196
-
2197
- let mut ready_0 = default_ready_response();
2198
- ready_0.max_model_len = 8_192;
2199
- ready_0.num_gpu_blocks = 10;
2200
- ready_0.effective_data_parallel_size = 2;
2201
- ready_0.tensor_parallel_size = 2;
2202
- ready_0.pipeline_parallel_size = 3;
2203
- ready_0.world_size = 12;
2204
- ready_0.weight_transfer_backend = Some("nccl".to_string());
2205
- ready_0.enable_sleep_mode = true;
2206
- ready_0.supports_draft_weight_updates = true;
2207
- ready_0.effective_attention_block_size = Some(64);
2208
-
2209
- let mut ready_1 = default_ready_response();
2210
- ready_1.max_model_len = 4_096;
2211
- ready_1.num_gpu_blocks = 20;
2212
- ready_1.effective_data_parallel_size = 2;
2213
- ready_1.tensor_parallel_size = 2;
2214
- ready_1.pipeline_parallel_size = 3;
2215
- ready_1.world_size = 12;
2216
- ready_1.data_parallel_rank = 1;
2217
-
2218
- let engine_tasks = [ready_0, ready_1].map(|ready| {
2219
- let engine_id = EngineId::from_engine_index(
2220
- ready.data_parallel_rank.try_into().expect("test rank fits engine identity"),
2195
+ // Local capacity and ownership must not expand to the deployment-wide DP size.
2196
+ for (start_rank, global_size) in [(0, 2), (0, 8), (4, 8)] {
2197
+ let ipc = IpcNamespace::new().expect("create ipc namespace");
2198
+ let input_address = ipc.input_endpoint();
2199
+ let output_address = ipc.output_endpoint();
2200
+
2201
+ let mut ready_0 = default_ready_response();
2202
+ ready_0.max_model_len = 8_192;
2203
+ ready_0.num_gpu_blocks = 10;
2204
+ ready_0.effective_data_parallel_size = global_size as u64;
2205
+ ready_0.data_parallel_rank = start_rank;
2206
+ ready_0.tensor_parallel_size = 2;
2207
+ ready_0.pipeline_parallel_size = 3;
2208
+ ready_0.world_size = 12;
2209
+ ready_0.weight_transfer_backend = Some("nccl".to_string());
2210
+ ready_0.enable_sleep_mode = true;
2211
+ ready_0.supports_draft_weight_updates = true;
2212
+ ready_0.effective_attention_block_size = Some(64);
2213
+
2214
+ let mut ready_1 = default_ready_response();
2215
+ ready_1.max_model_len = 4_096;
2216
+ ready_1.num_gpu_blocks = 20;
2217
+ ready_1.effective_data_parallel_size = global_size as u64;
2218
+ ready_1.tensor_parallel_size = 2;
2219
+ ready_1.pipeline_parallel_size = 3;
2220
+ ready_1.world_size = 12;
2221
+ ready_1.data_parallel_rank = start_rank + 1;
2222
+
2223
+ let client_config = EngineCoreClientConfig {
2224
+ transport_mode: TransportMode::Bootstrapped {
2225
+ input_address: input_address.clone(),
2226
+ output_address: output_address.clone(),
2227
+ engine_start_index: start_rank,
2228
+ engine_count: 2,
2229
+ data_parallel_size: global_size,
2230
+ ready_timeout: Duration::from_secs(2),
2231
+ },
2232
+ coordinator_mode: None,
2233
+ model_name: "test-model".to_string(),
2234
+ client_index: 0,
2235
+ };
2236
+ let client_task = tokio::spawn(EngineCoreClient::connect(client_config));
2237
+ let mut engine_sockets = Vec::new();
2238
+ for ready in [ready_0, ready_1] {
2239
+ let engine_id = EngineId::from_engine_index(
2240
+ ready.data_parallel_rank.try_into().expect("test rank fits engine identity"),
2241
+ );
2242
+ engine_sockets.push(
2243
+ connect_to_bootstrapped_frontend(
2244
+ &input_address,
2245
+ &output_address,
2246
+ engine_id,
2247
+ MockEngineConfig {
2248
+ ready_response: ready,
2249
+ ..Default::default()
2250
+ },
2251
+ )
2252
+ .await
2253
+ .expect("connect mock engine"),
2254
+ );
2255
+ }
2256
+ let client = client_task.await.unwrap().expect("connect multi-engine client");
2257
+ let chat = ChatLlm::from_shared_backend(
2258
+ Llm::new(client),
2259
+ Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
2260
);
2222
- MockEngineTask::new(spawn_mock_engine_task_with_ready(
2223
- handshake_address.clone(),
2224
- engine_id,
2225
- ready,
2226
- |_, _| boxed_test_future(async {}),
2227
- ))
2228
- });
2261
+ let state = AppState::new(vec!["test-model".to_string()], chat);
2262
+ let service = ControlServiceImpl::new(Arc::new(state));
2263
2230
- let client = EngineCoreClient::connect(EngineCoreClientConfig {
2231
- transport_mode: TransportMode::HandshakeOwner {
2232
- handshake_address,
2233
- advertised_host: "127.0.0.1".to_string(),
2234
- engine_count: 2,
2235
- ready_timeout: Duration::from_secs(2),
2236
- local_input_address: Some(ipc.input_endpoint()),
2237
- local_output_address: Some(ipc.output_endpoint()),
2238
- },
2239
- coordinator_mode: None,
2240
- model_name: "test-model".to_string(),
2241
- client_index: 0,
2242
- })
2243
- .await
2244
- .expect("connect multi-engine client");
2245
- let chat = ChatLlm::from_shared_backend(
2246
- Llm::new(client),
2247
- Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
2248
- );
2249
- let state = AppState::new(vec!["test-model".to_string()], chat);
2250
- let service = ControlServiceImpl::new(Arc::new(state));
2251
-
2252
- let server = pb::control_server::Control::get_server_info(
2253
- &service,
2254
- tonic::Request::new(pb::GetServerInfoRequest {}),
2255
- )
2256
- .await
2257
- .expect("get server info")
2258
- .into_inner();
2259
- assert_eq!(server.max_model_len, 4_096);
2260
- assert_eq!(server.total_kv_blocks, 30);
2261
- assert!(server.effective_attention_block_size.is_none());
2262
- let rl = server.rl_capabilities.expect("RL capabilities");
2263
- assert!(!rl.weight_transfer_enabled);
2264
- assert!(rl.weight_transfer_backend.is_empty());
2265
- assert!(!rl.sleep_mode_enabled);
2266
- assert!(!rl.draft_weight_updates_enabled);
2267
- let parallelism = server.parallelism.unwrap();
2268
- assert_eq!(parallelism.data_parallel_size, 2);
2269
- assert_eq!(parallelism.world_size, 12);
2270
-
2271
- drop(engine_tasks);
2264
+ let server = pb::control_server::Control::get_server_info(
2265
+ &service,
2266
+ tonic::Request::new(pb::GetServerInfoRequest {}),
2267
+ )
2268
+ .await
2269
+ .expect("get server info")
2270
+ .into_inner();
2271
+ assert_eq!(server.max_model_len, 4_096);
2272
+ assert_eq!(server.total_kv_blocks, 30);
2273
+ assert!(server.effective_attention_block_size.is_none());
2274
+ let rl = server.rl_capabilities.expect("RL capabilities");
2275
+ assert!(!rl.weight_transfer_enabled);
2276
+ assert!(rl.weight_transfer_backend.is_empty());
2277
+ assert!(!rl.sleep_mode_enabled);
2278
+ assert!(!rl.draft_weight_updates_enabled);
2279
+ let parallelism = server.parallelism.unwrap();
2280
+ assert_eq!(parallelism.data_parallel_size, global_size as u32);
2281
+ assert_eq!(parallelism.data_parallel_rank, start_rank);
2282
+ assert_eq!(parallelism.data_parallel_size_local, 2);
2283
+ assert_eq!(parallelism.world_size, 12);
2284
+
2285
+ drop(engine_sockets);
2286
+ }
2287
}
2288
2289
#[test]