iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Rust Frontend] Expose local DP size in gRPC Control metadata (#57116)

Browse Source 
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Alec
2026-09-16 11:28:08 -07:00
committed by GitHub
co-authored by Codex
parent fc8132a5e5
commit 9639cbde04
3 changed files with 90 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/proto/control.proto
+1
View File
@@ -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 {}
rust/src/server/src/grpc/control.rs
+1
View File
@@ -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,
}
}
rust/src/server/src/grpc/tests.rs
+88 -73
View File
@@ -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<dyn ChatTextBackend>,
);
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<dyn ChatTextBackend>,
);
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]