diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 16d69f7734..f9abac2004 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -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" diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 14291d7de0..003d96fa92 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -290,6 +290,7 @@ impl SharedRuntimeArgs { input_address: String, output_address: String, coordinator_address: Option, + 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, + /// 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, ) } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 31ace0a77d..c57c23e017 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -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 { diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index c646de567d..b4357f77c7 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -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, ) diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 7d24299dea..c00a422685 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -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, +) -> 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 { 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::>(); + 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(); diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index 360f94eda1..d0d9b4efe3 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -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 { @@ -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?; diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index 415ab549cc..8491e98216 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -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: diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index 9fd13a941e..ba66358c66 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -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])