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

[Rust Frontend] Support hybrid/external DP LB in Python supervised bootstrap (#45805)

Browse Source 
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-06-17 07:32:40 +00:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 43fa24e832
commit e2c58570ea
8 changed files with 159 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/rust_frontend.yaml
+4
View File
@@ -99,9 +99,13 @@ steps:
- vllm/v1/engine/
- vllm/v1/worker/
- tests/utils.py
- tests/v1/distributed/test_external_lb_dp.py
- tests/v1/distributed/test_hybrid_lb_dp.py
- tests/v1/distributed/test_internal_lb_dp.py
commands:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export NCCL_CUMEM_HOST_ENABLE=0
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
- 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
+7
View File
@@ -290,6 +290,7 @@ impl SharedRuntimeArgs {
input_address: String,
output_address: String,
coordinator_address: Option<String>,
engine_start_index: u32,
engine_count: usize,
) -> Config {
let ready_timeout = self.ready_timeout();
@@ -301,6 +302,7 @@ impl SharedRuntimeArgs {
transport_mode: TransportMode::Bootstrapped {
input_address,
output_address,
engine_start_index,
engine_count,
ready_timeout,
},
@@ -435,6 +437,10 @@ pub struct FrontendArgs {
/// `stats_update_address`.
#[arg(long)]
pub coordinator_address: Option<String>,
/// First data-parallel engine rank expected to register with this
/// bootstrapped frontend.
#[arg(long, default_value_t = 0)]
pub engine_start_index: u32,
/// Total number of data-parallel engines expected for this frontend.
#[arg(long, default_value_t = 1)]
pub engine_count: usize,
@@ -452,6 +458,7 @@ impl FrontendArgs {
self.input_address,
self.output_address,
self.coordinator_address,
self.engine_start_index,
self.engine_count,
)
}
rust/src/cmd/src/cli/tests.rs
+6 -2
View File
@@ -424,6 +424,7 @@ fn frontend_args_accept_json() {
coordinator_address: Some(
"tcp://127.0.0.1:7000",
),
engine_start_index: 0,
engine_count: 1,
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
@@ -1205,8 +1206,10 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
"ipc:///tmp/output.sock",
"--coordinator-address",
"tcp://127.0.0.1:7000",
"--engine-start-index",
"3",
"--engine-count",
"2",
"1",
"--args-json",
r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
])
@@ -1222,7 +1225,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
transport_mode: Bootstrapped {
input_address: "ipc:///tmp/input.sock",
output_address: "ipc:///tmp/output.sock",
engine_count: 2,
engine_start_index: 3,
engine_count: 1,
ready_timeout: 600s,
},
coordinator_mode: External {
rust/src/engine-core-client/src/client.rs
+5
View File
@@ -56,6 +56,9 @@ pub enum TransportMode {
/// Output PULL socket address that engines will connect to for
/// responses.
output_address: String,
/// First data-parallel engine rank expected to register on this
/// transport.
engine_start_index: u32,
/// Total number of engines expected to register on this transport.
engine_count: usize,
/// Maximum time to wait for all expected engines to register.
@@ -246,6 +249,7 @@ impl EngineCoreClient {
TransportMode::Bootstrapped {
input_address,
output_address,
engine_start_index,
engine_count,
ready_timeout,
} => {
@@ -256,6 +260,7 @@ impl EngineCoreClient {
transport::connect_bootstrapped(
input_address,
output_address,
*engine_start_index,
*engine_count,
*ready_timeout,
)
rust/src/engine-core-client/src/tests/client.rs
+113
View File
@@ -303,6 +303,7 @@ fn bootstrapped_test_config(
transport_mode: TransportMode::Bootstrapped {
input_address,
output_address,
engine_start_index: 0,
engine_count,
ready_timeout,
},
@@ -312,6 +313,34 @@ fn bootstrapped_test_config(
}
}
fn bootstrapped_test_config_with_start_index(
input_address: String,
output_address: String,
engine_start_index: u32,
engine_count: usize,
ready_timeout: Duration,
client_index: u32,
coordinator_mode: Option<CoordinatorMode>,
) -> EngineCoreClientConfig {
let mut config = bootstrapped_test_config(
input_address,
output_address,
engine_count,
ready_timeout,
client_index,
coordinator_mode,
);
let TransportMode::Bootstrapped {
engine_start_index: start,
..
} = &mut config.transport_mode
else {
unreachable!("bootstrapped_test_config returns bootstrapped transport")
};
*start = engine_start_index;
config
}
async fn recv_xpub_message(xpub: &mut XPubSocket) -> Vec<bytes::Bytes> {
xpub.recv().await.unwrap().into_vec()
}
@@ -2689,6 +2718,90 @@ async fn bootstrapped_connects_with_contiguous_engine_ids() {
client.shutdown().await.unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn bootstrapped_connects_with_nonzero_engine_start_index() {
init_tracing();
let ipc = IpcNamespace::new().unwrap();
let input_address = ipc.input_endpoint();
let output_address = ipc.output_endpoint();
let client_task = tokio::spawn({
let input_address = input_address.clone();
let output_address = output_address.clone();
async move {
EngineCoreClient::connect(bootstrapped_test_config_with_start_index(
input_address,
output_address,
3,
1,
Duration::from_secs(2),
0,
None,
))
.await
.unwrap()
}
});
let (_dealer, _push) =
setup_bootstrapped_mock_engine(input_address, output_address, &[0x03, 0x00]).await;
let client = client_task.await.unwrap();
assert_eq!(client.engine_count(), 1);
let engine_ids =
client.engine_identities().into_iter().map(|id| id.to_vec()).collect::<Vec<_>>();
assert_eq!(engine_ids, vec![vec![0x03, 0x00]]);
client.shutdown().await.unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn bootstrapped_rejects_unexpected_engine_id_for_start_index() {
init_tracing();
let ipc = IpcNamespace::new().unwrap();
let input_address = ipc.input_endpoint();
let output_address = ipc.output_endpoint();
let client_task = tokio::spawn({
let input_address = input_address.clone();
let output_address = output_address.clone();
async move {
EngineCoreClient::connect(bootstrapped_test_config_with_start_index(
input_address,
output_address,
3,
1,
Duration::from_secs(2),
0,
None,
))
.await
}
});
let _ = crate::mock_engine::connect_to_bootstrapped_frontend(
input_address,
output_address,
&[0x00, 0x00],
crate::mock_engine::MockEngineConfig {
local: true,
headless: true,
..Default::default()
},
)
.await;
let error = match client_task.await.unwrap() {
Ok(_) => panic!("bootstrapped connect should reject unexpected engine id"),
Err(error) => error,
};
assert!(
error
.to_string()
.contains("received input registration for unexpected engine id")
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn bootstrapped_connect_times_out_without_registration() {
init_tracing();
rust/src/engine-core-client/src/transport.rs
+3 -2
View File
@@ -327,6 +327,7 @@ pub async fn connect_handshake(
pub async fn connect_bootstrapped(
input_address: &str,
output_address: &str,
engine_start_index: u32,
engine_count: usize,
ready_timeout: Duration,
) -> Result<ConnectedTransport> {
@@ -342,8 +343,8 @@ pub async fn connect_bootstrapped(
let engines = wait_for_input_registrations(
&mut input_socket,
// TODO: follow start rank
(0..engine_count).map(|index| EngineId::from((index as u16).to_le_bytes().to_vec())),
(0..engine_count)
.map(|offset| EngineId::from_engine_index(engine_start_index + offset as u32)),
ready_timeout,
)
.await?;
vllm/entrypoints/cli/serve.py
+8 -1
View File
@@ -328,6 +328,12 @@ def run_multi_api_server(args: argparse.Namespace):
)
if rust_frontend_path:
if parallel_config.local_engines_only:
expected_engine_start_index = parallel_config.data_parallel_rank
expected_engine_count = parallel_config.data_parallel_size_local
else:
expected_engine_start_index = 0
expected_engine_count = parallel_config.data_parallel_size
# Start rust front-end process.
api_server_manager = RustFrontendProcessManager(
binary_path=rust_frontend_path,
@@ -335,7 +341,8 @@ def run_multi_api_server(args: argparse.Namespace):
args=args,
input_address=addresses.inputs[0],
output_address=addresses.outputs[0],
engine_count=parallel_config.data_parallel_size,
engine_start_index=expected_engine_start_index,
engine_count=expected_engine_count,
stats_update_address=stats_update_address,
)
else:
vllm/v1/utils.py
+13 -1
View File
@@ -336,6 +336,7 @@ class RustFrontendProcessManager:
args: argparse.Namespace,
input_address: str,
output_address: str,
engine_start_index: int,
engine_count: int,
stats_update_address: str | None = None,
):
@@ -354,6 +355,8 @@ class RustFrontendProcessManager:
input_address,
"--output-address",
output_address,
"--engine-start-index",
str(engine_start_index),
"--engine-count",
str(engine_count),
]
@@ -362,7 +365,16 @@ class RustFrontendProcessManager:
from vllm.entrypoints.serve.utils.api_utils import jsonify_non_default_args
args_json = json.dumps(
jsonify_non_default_args(args, exclude={"api_server_count"}),
jsonify_non_default_args(
args,
exclude={
"api_server_count",
# Python passes the bootstrapped engine range explicitly.
"data_parallel_rank",
"data_parallel_external_lb",
"data_parallel_hybrid_lb",
},
),
sort_keys=True,
)
cmd.extend(["--args-json", args_json])