From 9639cbde04f58c185ea6e5122dbd2039908abb7a Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:28:08 -0700 Subject: [PATCH] [Rust Frontend] Expose local DP size in gRPC Control metadata (#57116) Signed-off-by: Alec Flowers Co-authored-by: Codex --- rust/proto/control.proto | 1 + rust/src/server/src/grpc/control.rs | 1 + rust/src/server/src/grpc/tests.rs | 161 +++++++++++++++------------- 3 files changed, 90 insertions(+), 73 deletions(-) diff --git a/rust/proto/control.proto b/rust/proto/control.proto index d2816f68d3..ef3218d35c 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -68,6 +68,7 @@ message ParallelismInfo { uint32 data_parallel_rank = 4; uint32 decode_context_parallel_size = 5; uint64 world_size = 6; + uint32 data_parallel_size_local = 7; } message GetModelInfoRequest {} diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index de9c5ae65f..680ebdfbda 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -50,6 +50,7 @@ impl ControlServiceImpl { data_parallel_rank: ready.data_parallel_rank, decode_context_parallel_size: ready.decode_context_parallel_size, world_size: ready.world_size, + data_parallel_size_local: self.client().engine_count() as u32, } } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index a21e013323..f4b35c9df4 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -31,7 +31,7 @@ use vllm_chat::{ }; use vllm_engine_core_client::mock_engine::{ DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS, - default_ready_response, + MockEngineConfig, connect_to_bootstrapped_frontend, default_ready_response, }; use vllm_engine_core_client::protocol::decode_value; use vllm_engine_core_client::protocol::dtype::TensorDtype; @@ -1948,6 +1948,7 @@ async fn control_reports_server_and_model_info() { assert_eq!(parallelism.tensor_parallel_size, 1); assert_eq!(parallelism.pipeline_parallel_size, 1); assert_eq!(parallelism.data_parallel_size, 1); + assert_eq!(parallelism.data_parallel_size_local, 1); assert_eq!(parallelism.data_parallel_rank, 0); assert_eq!(parallelism.decode_context_parallel_size, 1); let rl = server.rl_capabilities.expect("RL capabilities"); @@ -2191,84 +2192,98 @@ async fn control_forwards_weight_update_without_pause_guard() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn control_aggregates_multi_engine_capacity() { - let ipc = IpcNamespace::new().expect("create ipc namespace"); - let handshake_address = ipc.handshake_endpoint(); + // Local capacity and ownership must not expand to the deployment-wide DP size. + for (start_rank, global_size) in [(0, 2), (0, 8), (4, 8)] { + let ipc = IpcNamespace::new().expect("create ipc namespace"); + let input_address = ipc.input_endpoint(); + let output_address = ipc.output_endpoint(); - let mut ready_0 = default_ready_response(); - ready_0.max_model_len = 8_192; - ready_0.num_gpu_blocks = 10; - ready_0.effective_data_parallel_size = 2; - ready_0.tensor_parallel_size = 2; - ready_0.pipeline_parallel_size = 3; - ready_0.world_size = 12; - ready_0.weight_transfer_backend = Some("nccl".to_string()); - ready_0.enable_sleep_mode = true; - ready_0.supports_draft_weight_updates = true; - ready_0.effective_attention_block_size = Some(64); + let mut ready_0 = default_ready_response(); + ready_0.max_model_len = 8_192; + ready_0.num_gpu_blocks = 10; + ready_0.effective_data_parallel_size = global_size as u64; + ready_0.data_parallel_rank = start_rank; + ready_0.tensor_parallel_size = 2; + ready_0.pipeline_parallel_size = 3; + ready_0.world_size = 12; + ready_0.weight_transfer_backend = Some("nccl".to_string()); + ready_0.enable_sleep_mode = true; + ready_0.supports_draft_weight_updates = true; + ready_0.effective_attention_block_size = Some(64); - let mut ready_1 = default_ready_response(); - ready_1.max_model_len = 4_096; - ready_1.num_gpu_blocks = 20; - ready_1.effective_data_parallel_size = 2; - ready_1.tensor_parallel_size = 2; - ready_1.pipeline_parallel_size = 3; - ready_1.world_size = 12; - ready_1.data_parallel_rank = 1; + let mut ready_1 = default_ready_response(); + ready_1.max_model_len = 4_096; + ready_1.num_gpu_blocks = 20; + ready_1.effective_data_parallel_size = global_size as u64; + ready_1.tensor_parallel_size = 2; + ready_1.pipeline_parallel_size = 3; + ready_1.world_size = 12; + ready_1.data_parallel_rank = start_rank + 1; - let engine_tasks = [ready_0, ready_1].map(|ready| { - let engine_id = EngineId::from_engine_index( - ready.data_parallel_rank.try_into().expect("test rank fits engine identity"), + let client_config = EngineCoreClientConfig { + transport_mode: TransportMode::Bootstrapped { + input_address: input_address.clone(), + output_address: output_address.clone(), + engine_start_index: start_rank, + engine_count: 2, + data_parallel_size: global_size, + ready_timeout: Duration::from_secs(2), + }, + coordinator_mode: None, + model_name: "test-model".to_string(), + client_index: 0, + }; + let client_task = tokio::spawn(EngineCoreClient::connect(client_config)); + let mut engine_sockets = Vec::new(); + for ready in [ready_0, ready_1] { + let engine_id = EngineId::from_engine_index( + ready.data_parallel_rank.try_into().expect("test rank fits engine identity"), + ); + engine_sockets.push( + connect_to_bootstrapped_frontend( + &input_address, + &output_address, + engine_id, + MockEngineConfig { + ready_response: ready, + ..Default::default() + }, + ) + .await + .expect("connect mock engine"), + ); + } + let client = client_task.await.unwrap().expect("connect multi-engine client"); + let chat = ChatLlm::from_shared_backend( + Llm::new(client), + Arc::new(FakeTextBackend) as Arc, ); - MockEngineTask::new(spawn_mock_engine_task_with_ready( - handshake_address.clone(), - engine_id, - ready, - |_, _| boxed_test_future(async {}), - )) - }); + let state = AppState::new(vec!["test-model".to_string()], chat); + let service = ControlServiceImpl::new(Arc::new(state)); - let client = EngineCoreClient::connect(EngineCoreClientConfig { - transport_mode: TransportMode::HandshakeOwner { - handshake_address, - advertised_host: "127.0.0.1".to_string(), - engine_count: 2, - ready_timeout: Duration::from_secs(2), - local_input_address: Some(ipc.input_endpoint()), - local_output_address: Some(ipc.output_endpoint()), - }, - coordinator_mode: None, - model_name: "test-model".to_string(), - client_index: 0, - }) - .await - .expect("connect multi-engine client"); - let chat = ChatLlm::from_shared_backend( - Llm::new(client), - Arc::new(FakeTextBackend) as Arc, - ); - let state = AppState::new(vec!["test-model".to_string()], chat); - let service = ControlServiceImpl::new(Arc::new(state)); + let server = pb::control_server::Control::get_server_info( + &service, + tonic::Request::new(pb::GetServerInfoRequest {}), + ) + .await + .expect("get server info") + .into_inner(); + assert_eq!(server.max_model_len, 4_096); + assert_eq!(server.total_kv_blocks, 30); + assert!(server.effective_attention_block_size.is_none()); + let rl = server.rl_capabilities.expect("RL capabilities"); + assert!(!rl.weight_transfer_enabled); + assert!(rl.weight_transfer_backend.is_empty()); + assert!(!rl.sleep_mode_enabled); + assert!(!rl.draft_weight_updates_enabled); + let parallelism = server.parallelism.unwrap(); + assert_eq!(parallelism.data_parallel_size, global_size as u32); + assert_eq!(parallelism.data_parallel_rank, start_rank); + assert_eq!(parallelism.data_parallel_size_local, 2); + assert_eq!(parallelism.world_size, 12); - let server = pb::control_server::Control::get_server_info( - &service, - tonic::Request::new(pb::GetServerInfoRequest {}), - ) - .await - .expect("get server info") - .into_inner(); - assert_eq!(server.max_model_len, 4_096); - assert_eq!(server.total_kv_blocks, 30); - assert!(server.effective_attention_block_size.is_none()); - let rl = server.rl_capabilities.expect("RL capabilities"); - assert!(!rl.weight_transfer_enabled); - assert!(rl.weight_transfer_backend.is_empty()); - assert!(!rl.sleep_mode_enabled); - assert!(!rl.draft_weight_updates_enabled); - let parallelism = server.parallelism.unwrap(); - assert_eq!(parallelism.data_parallel_size, 2); - assert_eq!(parallelism.world_size, 12); - - drop(engine_tasks); + drop(engine_sockets); + } } #[test]