@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e2c58570ea
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

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

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 17, 2026 at 15:32 UTC e2c58570eaae740fae1a5e50a01ceb18aa45c687
8 files changed +159 -6
.buildkite/test_areas/rust_frontend.yaml
+4
@@ -99,9 +99,13 @@ steps:
99 - vllm/v1/engine/
100 - vllm/v1/worker/
101 - tests/utils.py
102 + - tests/v1/distributed/test_external_lb_dp.py
103 + - tests/v1/distributed/test_hybrid_lb_dp.py
104 - tests/v1/distributed/test_internal_lb_dp.py
105 commands:
106 - export VLLM_USE_RUST_FRONTEND=1
107 - export VLLM_WORKER_MULTIPROC_METHOD=spawn
108 - export NCCL_CUMEM_HOST_ENABLE=0
109 - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
110 + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
111 + - 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
@@ -290,6 +290,7 @@ impl SharedRuntimeArgs {
290 input_address: String,
291 output_address: String,
292 coordinator_address: Option<String>,
293 + engine_start_index: u32,
294 engine_count: usize,
295 ) -> Config {
296 let ready_timeout = self.ready_timeout();
@@ -301,6 +302,7 @@ impl SharedRuntimeArgs {
302 transport_mode: TransportMode::Bootstrapped {
303 input_address,
304 output_address,
305 + engine_start_index,
306 engine_count,
307 ready_timeout,
308 },
@@ -435,6 +437,10 @@ pub struct FrontendArgs {
437 /// `stats_update_address`.
438 #[arg(long)]
439 pub coordinator_address: Option<String>,
440 + /// First data-parallel engine rank expected to register with this
441 + /// bootstrapped frontend.
442 + #[arg(long, default_value_t = 0)]
443 + pub engine_start_index: u32,
444 /// Total number of data-parallel engines expected for this frontend.
445 #[arg(long, default_value_t = 1)]
446 pub engine_count: usize,
@@ -452,6 +458,7 @@ impl FrontendArgs {
458 self.input_address,
459 self.output_address,
460 self.coordinator_address,
461 + self.engine_start_index,
462 self.engine_count,
463 )
464 }
rust/src/cmd/src/cli/tests.rs
+6 -2
@@ -424,6 +424,7 @@ fn frontend_args_accept_json() {
424 coordinator_address: Some(
425 "tcp://127.0.0.1:7000",
426 ),
427 + engine_start_index: 0,
428 engine_count: 1,
429 runtime: SharedRuntimeArgs {
430 model: "Qwen/Qwen3-0.6B",
@@ -1205,8 +1206,10 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
1206 "ipc:///tmp/output.sock",
1207 "--coordinator-address",
1208 "tcp://127.0.0.1:7000",
1209 + "--engine-start-index",
1210 + "3",
1211 "--engine-count",
1209 - "2",
1212 + "1",
1213 "--args-json",
1214 r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
1215 ])
@@ -1222,7 +1225,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
1225 transport_mode: Bootstrapped {
1226 input_address: "ipc:///tmp/input.sock",
1227 output_address: "ipc:///tmp/output.sock",
1225 - engine_count: 2,
1228 + engine_start_index: 3,
1229 + engine_count: 1,
1230 ready_timeout: 600s,
1231 },
1232 coordinator_mode: External {
rust/src/engine-core-client/src/client.rs
+5
@@ -56,6 +56,9 @@ pub enum TransportMode {
56 /// Output PULL socket address that engines will connect to for
57 /// responses.
58 output_address: String,
59 + /// First data-parallel engine rank expected to register on this
60 + /// transport.
61 + engine_start_index: u32,
62 /// Total number of engines expected to register on this transport.
63 engine_count: usize,
64 /// Maximum time to wait for all expected engines to register.
@@ -246,6 +249,7 @@ impl EngineCoreClient {
249 TransportMode::Bootstrapped {
250 input_address,
251 output_address,
252 + engine_start_index,
253 engine_count,
254 ready_timeout,
255 } => {
@@ -256,6 +260,7 @@ impl EngineCoreClient {
260 transport::connect_bootstrapped(
261 input_address,
262 output_address,
263 + *engine_start_index,
264 *engine_count,
265 *ready_timeout,
266 )
rust/src/engine-core-client/src/tests/client.rs
+113
@@ -303,6 +303,7 @@ fn bootstrapped_test_config(
303 transport_mode: TransportMode::Bootstrapped {
304 input_address,
305 output_address,
306 + engine_start_index: 0,
307 engine_count,
308 ready_timeout,
309 },
@@ -312,6 +313,34 @@ fn bootstrapped_test_config(
313 }
314 }
315
316 +fn bootstrapped_test_config_with_start_index(
317 + input_address: String,
318 + output_address: String,
319 + engine_start_index: u32,
320 + engine_count: usize,
321 + ready_timeout: Duration,
322 + client_index: u32,
323 + coordinator_mode: Option<CoordinatorMode>,
324 +) -> EngineCoreClientConfig {
325 + let mut config = bootstrapped_test_config(
326 + input_address,
327 + output_address,
328 + engine_count,
329 + ready_timeout,
330 + client_index,
331 + coordinator_mode,
332 + );
333 + let TransportMode::Bootstrapped {
334 + engine_start_index: start,
335 + ..
336 + } = &mut config.transport_mode
337 + else {
338 + unreachable!("bootstrapped_test_config returns bootstrapped transport")
339 + };
340 + *start = engine_start_index;
341 + config
342 +}
343 +
344 async fn recv_xpub_message(xpub: &mut XPubSocket) -> Vec<bytes::Bytes> {
345 xpub.recv().await.unwrap().into_vec()
346 }
@@ -2689,6 +2718,90 @@ async fn bootstrapped_connects_with_contiguous_engine_ids() {
2718 client.shutdown().await.unwrap();
2719 }
2720
2721 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2722 +async fn bootstrapped_connects_with_nonzero_engine_start_index() {
2723 + init_tracing();
2724 + let ipc = IpcNamespace::new().unwrap();
2725 + let input_address = ipc.input_endpoint();
2726 + let output_address = ipc.output_endpoint();
2727 +
2728 + let client_task = tokio::spawn({
2729 + let input_address = input_address.clone();
2730 + let output_address = output_address.clone();
2731 + async move {
2732 + EngineCoreClient::connect(bootstrapped_test_config_with_start_index(
2733 + input_address,
2734 + output_address,
2735 + 3,
2736 + 1,
2737 + Duration::from_secs(2),
2738 + 0,
2739 + None,
2740 + ))
2741 + .await
2742 + .unwrap()
2743 + }
2744 + });
2745 +
2746 + let (_dealer, _push) =
2747 + setup_bootstrapped_mock_engine(input_address, output_address, &[0x03, 0x00]).await;
2748 + let client = client_task.await.unwrap();
2749 +
2750 + assert_eq!(client.engine_count(), 1);
2751 + let engine_ids =
2752 + client.engine_identities().into_iter().map(|id| id.to_vec()).collect::<Vec<_>>();
2753 + assert_eq!(engine_ids, vec![vec![0x03, 0x00]]);
2754 +
2755 + client.shutdown().await.unwrap();
2756 +}
2757 +
2758 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2759 +async fn bootstrapped_rejects_unexpected_engine_id_for_start_index() {
2760 + init_tracing();
2761 + let ipc = IpcNamespace::new().unwrap();
2762 + let input_address = ipc.input_endpoint();
2763 + let output_address = ipc.output_endpoint();
2764 +
2765 + let client_task = tokio::spawn({
2766 + let input_address = input_address.clone();
2767 + let output_address = output_address.clone();
2768 + async move {
2769 + EngineCoreClient::connect(bootstrapped_test_config_with_start_index(
2770 + input_address,
2771 + output_address,
2772 + 3,
2773 + 1,
2774 + Duration::from_secs(2),
2775 + 0,
2776 + None,
2777 + ))
2778 + .await
2779 + }
2780 + });
2781 +
2782 + let _ = crate::mock_engine::connect_to_bootstrapped_frontend(
2783 + input_address,
2784 + output_address,
2785 + &[0x00, 0x00],
2786 + crate::mock_engine::MockEngineConfig {
2787 + local: true,
2788 + headless: true,
2789 + ..Default::default()
2790 + },
2791 + )
2792 + .await;
2793 + let error = match client_task.await.unwrap() {
2794 + Ok(_) => panic!("bootstrapped connect should reject unexpected engine id"),
2795 + Err(error) => error,
2796 + };
2797 +
2798 + assert!(
2799 + error
2800 + .to_string()
2801 + .contains("received input registration for unexpected engine id")
2802 + );
2803 +}
2804 +
2805 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2806 async fn bootstrapped_connect_times_out_without_registration() {
2807 init_tracing();
rust/src/engine-core-client/src/transport.rs
+3 -2
@@ -327,6 +327,7 @@ pub async fn connect_handshake(
327 pub async fn connect_bootstrapped(
328 input_address: &str,
329 output_address: &str,
330 + engine_start_index: u32,
331 engine_count: usize,
332 ready_timeout: Duration,
333 ) -> Result<ConnectedTransport> {
@@ -342,8 +343,8 @@ pub async fn connect_bootstrapped(
343
344 let engines = wait_for_input_registrations(
345 &mut input_socket,
345 - // TODO: follow start rank
346 - (0..engine_count).map(|index| EngineId::from((index as u16).to_le_bytes().to_vec())),
346 + (0..engine_count)
347 + .map(|offset| EngineId::from_engine_index(engine_start_index + offset as u32)),
348 ready_timeout,
349 )
350 .await?;
vllm/entrypoints/cli/serve.py
+8 -1
@@ -328,6 +328,12 @@ def run_multi_api_server(args: argparse.Namespace):
328 )
329
330 if rust_frontend_path:
331 + if parallel_config.local_engines_only:
332 + expected_engine_start_index = parallel_config.data_parallel_rank
333 + expected_engine_count = parallel_config.data_parallel_size_local
334 + else:
335 + expected_engine_start_index = 0
336 + expected_engine_count = parallel_config.data_parallel_size
337 # Start rust front-end process.
338 api_server_manager = RustFrontendProcessManager(
339 binary_path=rust_frontend_path,
@@ -335,7 +341,8 @@ def run_multi_api_server(args: argparse.Namespace):
341 args=args,
342 input_address=addresses.inputs[0],
343 output_address=addresses.outputs[0],
338 - engine_count=parallel_config.data_parallel_size,
344 + engine_start_index=expected_engine_start_index,
345 + engine_count=expected_engine_count,
346 stats_update_address=stats_update_address,
347 )
348 else:
vllm/v1/utils.py
+13 -1
@@ -336,6 +336,7 @@ class RustFrontendProcessManager:
336 args: argparse.Namespace,
337 input_address: str,
338 output_address: str,
339 + engine_start_index: int,
340 engine_count: int,
341 stats_update_address: str | None = None,
342 ):
@@ -354,6 +355,8 @@ class RustFrontendProcessManager:
355 input_address,
356 "--output-address",
357 output_address,
358 + "--engine-start-index",
359 + str(engine_start_index),
360 "--engine-count",
361 str(engine_count),
362 ]
@@ -362,7 +365,16 @@ class RustFrontendProcessManager:
365 from vllm.entrypoints.serve.utils.api_utils import jsonify_non_default_args
366
367 args_json = json.dumps(
365 - jsonify_non_default_args(args, exclude={"api_server_count"}),
368 + jsonify_non_default_args(
369 + args,
370 + exclude={
371 + "api_server_count",
372 + # Python passes the bootstrapped engine range explicitly.
373 + "data_parallel_rank",
374 + "data_parallel_external_lb",
375 + "data_parallel_hybrid_lb",
376 + },
377 + ),
378 sort_keys=True,
379 )
380 cmd.extend(["--args-json", args_json])