diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 9dcdb0bdbc..331ebfa8df 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -593,6 +593,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chacha20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.44" @@ -1614,6 +1625,7 @@ dependencies = [ "cfg-if", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -3451,6 +3463,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c9fb96cbc91e3478eaae79a69fcd3f1ae4ad052e471fe6732fff548984b4af" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", +] + [[package]] name = "rand_chacha" version = "0.3.1" @@ -3489,6 +3512,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rand_distr" version = "0.5.1" @@ -3936,9 +3965,9 @@ checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "saa" -version = "5.5.0" +version = "5.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16c7f49c9d5caa3bf4b3106900484b447b9253fe99670ceb81cb6cb5027855e1" +checksum = "bd85992a680de788217d9c70b74aa7685b108cf1aba02f728fb27f1b27199ad1" [[package]] name = "same-file" @@ -3960,12 +3989,12 @@ dependencies = [ [[package]] name = "scc" -version = "3.6.9" +version = "3.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45bb5ce9efd4a6e7b0f86c2697fe4c1d78d1f4e6d988c54b752d577cafe22fe8" +checksum = "5d5f8ca67c4f6208caf824037a6a9df0211b4e0937fccc4f543441f294d04511" dependencies = [ "saa", - "sdd 4.7.3", + "sdd 4.8.11", ] [[package]] @@ -4039,9 +4068,12 @@ checksum = "490dcfcbfef26be6800d11870ff2df8774fa6e86d047e3e8c8a76b25655e41ca" [[package]] name = "sdd" -version = "4.7.3" +version = "4.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b21a75f5913ab130e4b369fb8693be25f29b983e2ecad4279df9bfa5dd8aaf3e" +checksum = "a31877eee03686b905845993b5d9809a220c9dd5a10e9d6021d7ca67384c9070" +dependencies = [ + "saa", +] [[package]] name = "secrecy" @@ -4991,7 +5023,6 @@ dependencies = [ "bytes", "libc", "mio", - "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", @@ -5799,6 +5830,7 @@ dependencies = [ "half", "hex", "itertools 0.14.0", + "libc", "parking_lot", "rmp-serde", "rmpv", @@ -5808,6 +5840,7 @@ dependencies = [ "serde_repr", "serde_tuple", "serde_with", + "socket2", "task-local", "tempfile", "thiserror 2.0.18", @@ -6643,22 +6676,21 @@ checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" [[package]] name = "zeromq" version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efb2c254fd8f366755335c9e43b865f8484fe3bd717d65ffe7c3f28852863030" +source = "git+https://github.com/zeromq/zmq.rs?rev=67dcba97559ebe8e8206dab753710d8b9cb07591#67dcba97559ebe8e8206dab753710d8b9cb07591" dependencies = [ "async-trait", "asynchronous-codec", "bytes", "crossbeam-queue", "futures", + "futures-channel", "log", - "num-traits", - "once_cell", "parking_lot", - "rand 0.9.5", + "rand 0.10.3", "regex", - "scc 3.6.9", - "thiserror 1.0.69", + "scc 3.8.8", + "socket2", + "thiserror 2.0.18", "tokio", "tokio-util", "uuid", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 38436467f8..57f680cc14 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -155,7 +155,7 @@ vllm-tokenizer = { path = "src/tokenizer" } vllm-tracing = { path = "src/tracing" } winnow = { version = "1.0.2", features = ["simd"] } xgrammar-structural-tag = "0.4.1" -zeromq = { version = "0.6.0", default-features = false, features = [ +zeromq = { git = "https://github.com/zeromq/zmq.rs", rev = "67dcba97559ebe8e8206dab753710d8b9cb07591", default-features = false, features = [ "tokio-runtime", "all-transport", ] } diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 1269e0b8af..d3e369bbe3 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -501,8 +501,8 @@ impl SharedRuntimeArgs { self, listen_fd: i32, grpc_listen_fd: Option, - input_address: String, - output_address: String, + input_listener_fd: i32, + output_listener_fd: i32, coordinator_address: Option, engine_start_index: u32, engine_count: usize, @@ -518,8 +518,8 @@ impl SharedRuntimeArgs { Config { transport_mode: TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd, + output_listener_fd, engine_start_index, engine_count, data_parallel_size, @@ -697,14 +697,12 @@ pub struct FrontendArgs { /// supervisor. When not set, no gRPC server is started. #[arg(long)] pub grpc_listen_fd: Option, - /// Frontend input ROUTER socket address that the Python engines will - /// connect to. + /// Inherited frontend input ROUTER listener file descriptor. #[arg(long)] - pub input_address: String, - /// Frontend output PULL socket address that the Python engines will push - /// responses to. + pub input_listener_fd: i32, + /// Inherited frontend output PULL listener file descriptor. #[arg(long)] - pub output_address: String, + pub output_listener_fd: i32, /// Optional Python-owned frontend-side DP coordinator socket address for /// external coordinator mode in the bootstrapped frontend path, i.e., /// `stats_update_address`. @@ -733,8 +731,8 @@ impl FrontendArgs { self.runtime.into_bootstrapped_config( self.listen_fd, self.grpc_listen_fd, - self.input_address, - self.output_address, + self.input_listener_fd, + self.output_listener_fd, 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 944febc36c..aa164bc618 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -668,10 +668,10 @@ fn frontend_args_json_passes_tls_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","ssl_certfile":"/tmp/cert.pem","ssl_keyfile":"/tmp/key.pem"}"#, ]) @@ -693,10 +693,10 @@ fn frontend_args_json_rejects_out_of_range_cert_reqs() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","ssl_certfile":"/tmp/cert.pem","ssl_cert_reqs":5}"#, ]) @@ -718,10 +718,10 @@ fn frontend_args_json_passes_enable_request_id_headers_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#, ]) @@ -807,10 +807,10 @@ fn frontend_args_json_passes_lora_modules_into_config() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","lora_modules":[{"name":"alice","path":"org/alice","base_model_name":null,"is_3d_lora_weight":true}]}"#, ]) @@ -862,10 +862,10 @@ fn frontend_args_json_accepts_api_key_string() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":"secret"}"#, ]) @@ -885,10 +885,10 @@ fn frontend_args_json_accepts_api_key_list() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":["secret-a","secret-b"]}"#, ]) @@ -991,10 +991,10 @@ fn frontend_args_accept_json() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--coordinator-address", "tcp://127.0.0.1:7000", "--args-json", @@ -1008,8 +1008,8 @@ fn frontend_args_accept_json() { FrontendArgs { listen_fd: 3, grpc_listen_fd: None, - input_address: "ipc:///tmp/input.sock", - output_address: "ipc:///tmp/output.sock", + input_listener_fd: 999999, + output_listener_fd: 999998, coordinator_address: Some( "tcp://127.0.0.1:7000", ), @@ -1083,10 +1083,10 @@ fn frontend_args_json_applies_defaults() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--engine-count", "4", "--args-json", @@ -1115,10 +1115,10 @@ fn frontend_args_json_ignores_engine_owned_max_model_len() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","max_model_len":-1}"#, ]) @@ -1165,10 +1165,10 @@ fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", &payload.to_string(), ]) @@ -1204,10 +1204,10 @@ fn frontend_args_json_accepts_supported_non_default_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","revision":"release","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, ]) @@ -1286,10 +1286,10 @@ fn frontend_args_json_ignores_unknown_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","uds":"/tmp/vllm.sock","nested_unknown":{"x":1}}"#, ]) @@ -1308,10 +1308,10 @@ fn frontend_args_json_sets_prompt_tokens_details_flag() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","api_server_count":2,"enable_prompt_tokens_details":true}"#, ]) @@ -1371,10 +1371,10 @@ fn frontend_args_json_parses_cors_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","allowed_origins":["http://a.com"],"allow_credentials":true}"#, ]) @@ -1396,10 +1396,10 @@ fn frontend_args_json_rejects_unsupported_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","root_path":"/prefix"}"#, ]) @@ -1423,10 +1423,10 @@ fn frontend_args_json_aggregates_multiple_unsupported_fields() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","root_path":"/prefix"}"#, ]) @@ -1452,10 +1452,10 @@ fn frontend_args_json_rejects_malformed_json() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B""#, ]) @@ -1999,10 +1999,10 @@ fn frontend_args_reject_legacy_handshake_flags() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, "--handshake-address", @@ -2043,10 +2043,10 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--coordinator-address", "tcp://127.0.0.1:7000", "--engine-start-index", @@ -2068,8 +2068,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present expect![[r#" Config { transport_mode: Bootstrapped { - input_address: "ipc:///tmp/input.sock", - output_address: "ipc:///tmp/output.sock", + input_listener_fd: 999999, + output_listener_fd: 999998, engine_start_index: 3, engine_count: 1, data_parallel_size: 4, @@ -2200,10 +2200,10 @@ fn frontend_config_adopts_inherited_grpc_listener() { "3", "--grpc-listen-fd", "4", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","grpc_port":50051}"#, ]) @@ -2227,10 +2227,10 @@ fn frontend_args_json_enables_profiling_when_profiler_config_set() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}}"#, ]) @@ -2251,10 +2251,10 @@ fn frontend_args_json_disables_profiling_when_profiler_config_absent() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, ]) @@ -2275,10 +2275,10 @@ fn frontend_args_json_disables_profiling_when_profiler_type_is_null() { "frontend", "--listen-fd", "3", - "--input-address", - "ipc:///tmp/input.sock", - "--output-address", - "ipc:///tmp/output.sock", + "--input-listener-fd", + "999999", + "--output-listener-fd", + "999998", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":null}}"#, ]) diff --git a/rust/src/engine-core-client/Cargo.toml b/rust/src/engine-core-client/Cargo.toml index 14b9fe1423..86ffe4c4cd 100644 --- a/rust/src/engine-core-client/Cargo.toml +++ b/rust/src/engine-core-client/Cargo.toml @@ -18,6 +18,7 @@ futures.workspace = true half.workspace = true hex.workspace = true itertools.workspace = true +libc.workspace = true parking_lot.workspace = true rmp-serde.workspace = true rmpv.workspace = true @@ -27,6 +28,7 @@ serde_json.workspace = true serde_repr.workspace = true serde_tuple.workspace = true serde_with.workspace = true +socket2.workspace = true task-local.workspace = true tempfile = { workspace = true, optional = true } thiserror.workspace = true diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 36a4fbffa5..b3516936fe 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -53,16 +53,15 @@ pub enum TransportMode { local_output_address: Option, }, - /// The Python supervisor has already chosen the frontend transport - /// addresses, and the Rust process only needs to bind them and wait for - /// engine registration frames. + /// The Python supervisor has already bound the frontend transport + /// listeners. Rust adopts them and waits for engine registration frames. + /// `EngineCoreClient::connect` consumes each descriptor exactly once; + /// copied integer values do not create additional descriptor ownership. Bootstrapped { - /// Input ROUTER socket address that engines will connect to for - /// requests. - input_address: String, - /// Output PULL socket address that engines will connect to for - /// responses. - output_address: String, + /// Raw input ROUTER listener descriptor inherited from the supervisor. + input_listener_fd: i32, + /// Raw output PULL listener descriptor inherited from the supervisor. + output_listener_fd: i32, /// First data-parallel engine rank expected to register on this /// transport. engine_start_index: u32, @@ -336,8 +335,8 @@ impl EngineCoreClient { } TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd, + output_listener_fd, engine_start_index, engine_count, ready_timeout, @@ -348,8 +347,8 @@ impl EngineCoreClient { } transport::connect_bootstrapped( - input_address, - output_address, + *input_listener_fd, + *output_listener_fd, *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 0ac81764a9..dbd3fa69ce 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -4,6 +4,8 @@ use std::collections::{BTreeMap, BTreeSet}; use std::convert::TryFrom; use std::io::Cursor; +use std::os::fd::IntoRawFd; +use std::os::unix::net::UnixListener; use std::path::PathBuf; use std::process::Command; use std::sync::Once; @@ -311,10 +313,17 @@ fn bootstrapped_test_config( client_index: u32, coordinator_mode: Option, ) -> EngineCoreClientConfig { + fn listener(address: &str) -> i32 { + let path = address.strip_prefix("ipc://").expect("bootstrapped tests use IPC listeners"); + let _ = std::fs::remove_file(path); + let listener = UnixListener::bind(path).expect("bind inherited test listener"); + listener.into_raw_fd() + } + EngineCoreClientConfig { transport_mode: TransportMode::Bootstrapped { - input_address, - output_address, + input_listener_fd: listener(&input_address), + output_listener_fd: listener(&output_address), engine_start_index: 0, engine_count, data_parallel_size: engine_count, @@ -362,15 +371,22 @@ fn bootstrapped_test_config_with_start_index( #[test] fn client_config_validates_bootstrapped_dp_range() { - let mut config = bootstrapped_test_config_with_start_index( - "ipc://unused-input".to_string(), - "ipc://unused-output".to_string(), - 1, - 2, - Duration::from_secs(1), - 0, - None, - ); + // Validation never consumes the descriptors, so placeholders avoid opening + // listeners that nothing would close. + let mut config = EngineCoreClientConfig { + transport_mode: TransportMode::Bootstrapped { + input_listener_fd: -1, + output_listener_fd: -1, + engine_start_index: 1, + engine_count: 2, + data_parallel_size: 3, + ready_timeout: Duration::from_secs(1), + }, + coordinator_mode: None, + model_name: "test-model".to_string(), + client_index: 0, + engine_stats_enabled: true, + }; config.validate().expect("frontend may own a subset of global DP ranks"); let TransportMode::Bootstrapped { diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index 7aac8774ca..e453cafdc8 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -4,6 +4,8 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Debug; use std::ops::Deref; +use std::os::fd::{FromRawFd, OwnedFd, RawFd}; +use std::os::unix::net::UnixListener as StdUnixListener; use std::time::Duration; use bytes::Bytes; @@ -327,8 +329,8 @@ pub async fn connect_handshake( /// identities are synthesized from contiguous rank order instead of being /// discovered through a Rust-owned handshake. pub async fn connect_bootstrapped( - input_address: &str, - output_address: &str, + input_listener_fd: RawFd, + output_listener_fd: RawFd, engine_start_index: u32, engine_count: usize, ready_timeout: Duration, @@ -354,10 +356,16 @@ pub async fn connect_bootstrapped( } let mut input_socket = RouterSocket::new(); - let input_address = input_socket.bind(input_address).await?.to_string(); + let input_address = input_socket + .bind_listener(adopt_listener(input_listener_fd)?) + .await? + .to_string(); let mut output_socket = PullSocket::new(); - let output_address = output_socket.bind(output_address).await?.to_string(); + let output_address = output_socket + .bind_listener(adopt_listener(output_listener_fd)?) + .await? + .to_string(); let engines = wait_for_input_registrations( &mut input_socket, @@ -385,6 +393,34 @@ pub async fn connect_bootstrapped( }) } +fn adopt_listener(fd: RawFd) -> Result { + if fd < 0 || unsafe { libc::fcntl(fd, libc::F_GETFD) } < 0 { + return Err(Error::InvalidClientConfig { + message: format!("inherited ZMQ listener fd {fd} is not open"), + }); + } + + // SAFETY: Python's Popen(pass_fds) gives the frontend its own descriptor. + // The bootstrapped startup path calls this function once per descriptor, + // transferring ownership to the listener constructed below. + let fd = unsafe { OwnedFd::from_raw_fd(fd) }; + let socket: socket2::Socket = fd.into(); + let domain = socket.domain()?; + + if domain == socket2::Domain::IPV4 || domain == socket2::Domain::IPV6 { + let listener: std::net::TcpListener = socket.into(); + return Ok(listener.into()); + } + if domain == socket2::Domain::UNIX { + let listener: StdUnixListener = socket.into(); + return Ok(listener.into()); + } + + Err(Error::InvalidClientConfig { + message: format!("unsupported inherited listener domain {domain:?}"), + }) +} + /// Bind new input and output sockets. async fn bind_local_sockets( local_host: &str, @@ -603,7 +639,12 @@ pub async fn run_output_loop( #[cfg(test)] mod tests { - use super::bind_local_sockets; + use std::os::fd::IntoRawFd; + + use zeromq::RouterSocket; + use zeromq::prelude::Socket; + + use super::{adopt_listener, bind_local_sockets}; #[tokio::test] async fn bind_local_sockets_resolves_zero_port_bindings() { @@ -614,4 +655,21 @@ mod tests { assert!(output_address.starts_with("tcp://127.0.0.1:")); assert_ne!(input_address, output_address); } + + #[tokio::test] + async fn inherited_tcp_listener_is_adopted_without_rebinding() { + let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).expect("bind listener"); + let expected = listener.local_addr().expect("listener address"); + let mut router = RouterSocket::new(); + + let endpoint = router + .bind_listener( + adopt_listener(listener.into_raw_fd()).expect("adopt inherited listener fd"), + ) + .await + .expect("bind inherited listener"); + + assert_eq!(endpoint.to_string(), format!("tcp://{expected}")); + assert!(router.close().await.is_empty()); + } } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index d8d0fb4e6c..631eb76b36 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -5,6 +5,8 @@ use std::collections::BTreeMap; use std::fs; use std::future::Future; use std::io; +use std::os::fd::IntoRawFd; +use std::os::unix::net::UnixListener; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -2230,10 +2232,14 @@ async fn control_aggregates_multi_engine_capacity() { ready_1.world_size = 12; ready_1.data_parallel_rank = start_rank + 1; + let listener_fd = |address: &str| { + let path = address.strip_prefix("ipc://").expect("IPC test endpoint"); + UnixListener::bind(path).expect("bind inherited test listener").into_raw_fd() + }; let client_config = EngineCoreClientConfig { transport_mode: TransportMode::Bootstrapped { - input_address: input_address.clone(), - output_address: output_address.clone(), + input_listener_fd: listener_fd(&input_address), + output_listener_fd: listener_fd(&output_address), engine_start_index: start_rank, engine_count: 2, data_parallel_size: global_size, diff --git a/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py b/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py deleted file mode 100644 index c8ff969a8a..0000000000 --- a/tests/entrypoints/launchers/api_server/_api_server_spawn_workers.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Spawn worker targets kept free of heavy imports. - -``multiprocessing`` with the ``spawn`` start method re-imports the module that -defines a process target in the child. Housing these stubs in a stdlib-only -module keeps child startup fast and deterministic, instead of paying a multi- -second ``import vllm`` before the child can run. -""" - - -def exit_before_report_worker(listen_address, sock, args, client_config=None): - """Exit immediately without touching ``actual_address_pipe``.""" - return diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 6c579df838..318f57b86f 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing +import os import socket import threading import time @@ -11,14 +12,10 @@ from unittest.mock import patch import pytest import zmq -from tests.entrypoints.launchers.api_server._api_server_spawn_workers import ( - exit_before_report_worker, -) -from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.utils.network_utils import make_zmq_listener, make_zmq_socket from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.utils import ( APIServerProcessManager, - get_engine_client_zmq_addr, wait_for_completion_or_failure, ) @@ -46,36 +43,30 @@ def update_admission_stats_worker(listen_address, sock, args, client_config): args.results.put(stats.get_num_requests()) -# Module-level stub for the gather_actual_addresses test. Must be -# importable by `multiprocessing.spawn` (no closures, no nesting). -def defer_addresses_stub_worker(listen_address, sock, args, client_config): - """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via the pipe, then exit.""" +def _adopt_zmq_listener_worker(listen_address, sock, args, client_config): ctx = zmq.Context() + router = make_zmq_socket( + ctx, + client_config["input_address"], + zmq.ROUTER, + bind=True, + listener=client_config["input_listener"], + ) + pull = make_zmq_socket( + ctx, + client_config["output_address"], + zmq.PULL, + bind=True, + listener=client_config["output_listener"], + ) try: - in_sock = make_zmq_socket( - ctx, client_config["input_address"], zmq.ROUTER, bind=True - ) - out_sock = make_zmq_socket( - ctx, client_config["output_address"], zmq.PULL, bind=True - ) - try: - pipe = client_config["actual_address_pipe"] - try: - pipe.send( - { - "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(), - "output_address": out_sock.getsockopt( - zmq.LAST_ENDPOINT - ).decode(), - } - ) - finally: - pipe.close() - finally: - in_sock.close(linger=0) - out_sock.close(linger=0) + identity, payload = router.recv_multipart() + assert payload == b"input" + assert pull.recv() == b"output" + router.send_multipart([identity, b"ok"]) finally: + router.close(linger=0) + pull.close(linger=0) ctx.term() @@ -89,15 +80,11 @@ def api_server_args(): "sock": sock, "args": "test_args", # Simple string to avoid pickling issues "num_servers": 3, - "input_addresses": [ - "tcp://127.0.0.1:5001", - "tcp://127.0.0.1:5002", - "tcp://127.0.0.1:5003", + "input_listeners": [ + make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) for _ in range(3) ], - "output_addresses": [ - "tcp://127.0.0.1:6001", - "tcp://127.0.0.1:6002", - "tcp://127.0.0.1:6003", + "output_listeners": [ + make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) for _ in range(3) ], "stats_update_address": "tcp://127.0.0.1:7000", } @@ -120,6 +107,13 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): try: # Verify the manager was initialized correctly assert len(manager.processes) == 3 + assert all( + listener.socket.fileno() == -1 + for listener in ( + *args["input_listeners"], + *args["output_listeners"], + ) + ) # Verify all processes are running for proc in manager.processes: @@ -145,6 +139,70 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): assert not proc.is_alive() +def _assert_endpoint_reserved(address: str) -> None: + """A competing bind must fail while the endpoint has a listener owner.""" + if address.startswith("ipc://"): + competitor = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + bind_address: str | tuple[str, int] = address.removeprefix("ipc://") + else: + competitor = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + bind_address = ("127.0.0.1", int(address.rsplit(":", 1)[1])) + with competitor, pytest.raises(OSError, match="Address already in use"): + competitor.bind(bind_address) + + +@pytest.mark.parametrize("scheme", ["tcp", "ipc"]) +def test_api_server_child_adopts_inherited_zmq_listeners(tmp_path, scheme): + if scheme == "tcp": + input_path = output_path = "tcp://127.0.0.1:0" + else: + input_path = f"ipc://{tmp_path}/input.sock" + output_path = f"ipc://{tmp_path}/output.sock" + input_listener = make_zmq_listener(input_path, zmq.ROUTER) + output_listener = make_zmq_listener(output_path, zmq.PULL) + input_address = input_listener.address + output_address = output_listener.address + # The supervisor reserves each endpoint as soon as it allocates it. + _assert_endpoint_reserved(input_address) + _assert_endpoint_reserved(output_address) + http_socket = socket.socket() + manager = APIServerProcessManager( + target_server_fn=_adopt_zmq_listener_worker, + listen_address="localhost:8000", + sock=http_socket, + args="test_args", + num_servers=1, + input_listeners=[input_listener], + output_listeners=[output_listener], + ) + + ctx = zmq.Context() + dealer = make_zmq_socket( + ctx, input_address, zmq.DEALER, bind=False, identity=b"test" + ) + push = make_zmq_socket(ctx, output_address, zmq.PUSH, bind=False) + dealer.setsockopt(zmq.RCVTIMEO, 30_000) + try: + assert input_listener.socket.fileno() == -1 + assert output_listener.socket.fileno() == -1 + # After the parent closes its copies, the child's inherited descriptors + # keep both endpoints reserved; closing must not unlink IPC paths. + assert manager.processes[0].is_alive() + _assert_endpoint_reserved(input_address) + _assert_endpoint_reserved(output_address) + dealer.send(b"input") + push.send(b"output") + assert dealer.recv() == b"ok" + manager.processes[0].join(timeout=10) + assert manager.processes[0].exitcode == 0 + finally: + manager.shutdown() + dealer.close(linger=0) + push.close(linger=0) + ctx.term() + http_socket.close() + + @pytest.mark.skip_global_cleanup def test_api_server_processes_share_admission_stats(api_server_args): args = api_server_args.copy() @@ -354,96 +412,7 @@ def test_external_process_monitoring(api_server_args): WORKER_RUNTIME_SECONDS = prev_worker_runtime -@pytest.mark.timeout(60) -def test_gather_actual_addresses_end_to_end(): - """Each child binds ROUTER/PULL with a kernel-picked port and reports - the bound endpoints back via its per-child pipe; the manager surfaces - them via :py:meth:`gather_actual_addresses`.""" - host = "127.0.0.1" - num_servers = 4 - - placeholder_inputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - placeholder_outputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - for addr in placeholder_inputs + placeholder_outputs: - assert addr == f"tcp://{host}:0", addr - - sock = socket.socket() - manager = APIServerProcessManager( - listen_address=f"tcp://{host}:0", - sock=sock, - args="test_args", - num_servers=num_servers, - input_addresses=placeholder_inputs, - output_addresses=placeholder_outputs, - target_server_fn=defer_addresses_stub_worker, - ) - - try: - assert len(manager.processes) == num_servers - actual_inputs, actual_outputs = manager.gather_actual_addresses(timeout=15.0) - finally: - manager.shutdown() - time.sleep(0.2) - sock.close() - - assert len(actual_inputs) == num_servers - assert len(actual_outputs) == num_servers - - for addr in actual_inputs + actual_outputs: - scheme, parsed_host, port = split_zmq_path(addr) - assert scheme == "tcp", addr - assert parsed_host == host, addr - assert port and int(port) > 0, addr - - all_addrs = actual_inputs + actual_outputs - assert len(set(all_addrs)) == len(all_addrs), all_addrs - - -@pytest.mark.timeout(30) -def test_gather_actual_addresses_child_crash_before_report(): - """A child that exits before sending its endpoints must surface a - clear ``RuntimeError`` rather than hang or return ``None`` slots.""" - host = "127.0.0.1" - num_servers = 2 - placeholder_inputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - placeholder_outputs = [ - get_engine_client_zmq_addr(local_only=False, host=host) - for _ in range(num_servers) - ] - - sock = socket.socket() - manager = APIServerProcessManager( - listen_address=f"tcp://{host}:0", - sock=sock, - args="test_args", - num_servers=num_servers, - input_addresses=placeholder_inputs, - output_addresses=placeholder_outputs, - # exit_before_report_worker exits without touching - # ``actual_address_pipe`` — simulates a child that dies before - # reporting its bound addresses. - target_server_fn=exit_before_report_worker, - ) - try: - # Sentinel-first vs pipe-EOF-first both produce "reporting". - with pytest.raises(RuntimeError, match="reporting"): - manager.gather_actual_addresses(timeout=10.0) - finally: - manager.shutdown() - time.sleep(0.2) - sock.close() - - -def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog): +def test_rust_frontend_launch_log_redacts_credentials(monkeypatch): """The Rust frontend command carries every non-default arg as JSON. Credentials must not reach the log line.""" import subprocess as subprocess_mod @@ -469,15 +438,17 @@ def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog): return 0 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) try: monkeypatch.setattr(subprocess_mod, "Popen", lambda *a, **kw: _FakeProc()) - with caplog.at_level("INFO", logger="vllm.v1.utils"): + with patch("vllm.v1.utils.logger.info") as log_info: RustFrontendProcessManager( binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -485,11 +456,13 @@ def test_rust_frontend_launch_log_redacts_credentials(monkeypatch, caplog): finally: sock.close() - message = caplog.text + log_format, *log_args = log_info.call_args.args + message = log_format % tuple(log_args) assert "Launching Rust frontend:" in message assert hf_token not in message assert api_key not in message assert '"hf_token": "***"' in message + assert '"api_key": "***"' in message def test_rust_frontend_uses_config_model_as_model_tag(monkeypatch, caplog, tmp_path): @@ -519,6 +492,8 @@ def test_rust_frontend_uses_config_model_as_model_tag(monkeypatch, caplog, tmp_p return 0 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) try: monkeypatch.setattr(subprocess_mod, "Popen", lambda *a, **kw: _FakeProc()) with caplog.at_level("INFO", logger="vllm.v1.utils"): @@ -526,8 +501,8 @@ def test_rust_frontend_uses_config_model_as_model_tag(monkeypatch, caplog, tmp_p binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -568,14 +543,18 @@ def test_rust_frontend_inherits_grpc_listener(monkeypatch): sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) grpc_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + input_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.ROUTER) + output_listener = make_zmq_listener("tcp://127.0.0.1:0", zmq.PULL) + input_fd = input_listener.socket.fileno() + output_fd = output_listener.socket.fileno() try: monkeypatch.setattr(subprocess_mod, "Popen", fake_popen) RustFrontendProcessManager( binary_path="/nonexistent/vllm-rs", sock=sock, args=args, - input_address="ipc:///tmp/in", - output_address="ipc:///tmp/out", + input_listener=input_listener, + output_listener=output_listener, engine_start_index=0, engine_count=1, data_parallel_size=1, @@ -585,7 +564,44 @@ def test_rust_frontend_inherits_grpc_listener(monkeypatch): grpc_fd = grpc_sock.fileno() flag_index = cmd.index("--grpc-listen-fd") assert cmd[flag_index + 1] == str(grpc_fd) - assert list(kwargs["pass_fds"]) == [sock.fileno(), grpc_fd] + assert list(kwargs["pass_fds"]) == [ + sock.fileno(), + input_fd, + output_fd, + grpc_fd, + ] finally: sock.close() grpc_sock.close() + + +def test_rust_frontend_spawn_failure_cleans_up_listeners(tmp_path): + """No child owns the listeners when Popen fails, so the parent removes them.""" + from vllm.entrypoints.launchers.cli_args import make_arg_parser + from vllm.utils.argparse_utils import FlexibleArgumentParser + from vllm.v1.utils import RustFrontendProcessManager + + args = make_arg_parser(FlexibleArgumentParser()).parse_args( + ["--model", "org/model"] + ) + input_listener = make_zmq_listener(f"ipc://{tmp_path}/input.sock", zmq.ROUTER) + output_listener = make_zmq_listener(f"ipc://{tmp_path}/output.sock", zmq.PULL) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + with pytest.raises(FileNotFoundError): + RustFrontendProcessManager( + binary_path=str(tmp_path / "missing-vllm-rs"), + sock=sock, + args=args, + input_listener=input_listener, + output_listener=output_listener, + engine_start_index=0, + engine_count=1, + data_parallel_size=1, + ) + finally: + sock.close() + + for listener in (input_listener, output_listener): + assert listener.socket.fileno() == -1 + assert not os.path.exists(listener.address.removeprefix("ipc://")) diff --git a/tests/utils_/test_network_utils.py b/tests/utils_/test_network_utils.py index a6c359e269..7ba3d6ae3e 100644 --- a/tests/utils_/test_network_utils.py +++ b/tests/utils_/test_network_utils.py @@ -13,6 +13,7 @@ from vllm.utils.network_utils import ( get_open_ports_list, get_tcp_uri, join_host_port, + make_zmq_listener, make_zmq_path, make_zmq_socket, split_host_port, @@ -205,6 +206,52 @@ def test_make_zmq_socket_ipv6(): ctx.term() +@pytest.mark.parametrize("scheme", ["tcp", "ipc"]) +def test_make_zmq_socket_adopts_listener(tmp_path, scheme): + path = ( + "tcp://127.0.0.1:0" if scheme == "tcp" else f"ipc://{tmp_path}/inherited.sock" + ) + listener = make_zmq_listener(path, zmq.ROUTER) + ctx = zmq.Context() + router = dealer = None + try: + router = make_zmq_socket( + ctx, + listener.address, + zmq.ROUTER, + bind=True, + listener=listener.socket, + ) + assert listener.socket.fileno() == -1 + dealer = make_zmq_socket( + ctx, listener.address, zmq.DEALER, bind=False, identity=b"client" + ) + dealer.send(b"ping") + assert router.recv_multipart() == [b"client", b"ping"] + finally: + if router is not None: + router.close(linger=0) + if dealer is not None: + dealer.close(linger=0) + ctx.term() + listener.cleanup() + + +def test_make_zmq_listener_failed_ipc_bind_keeps_existing_path(tmp_path): + path = tmp_path / "taken.sock" + owner = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + owner.bind(str(path)) + owner.listen() + try: + with pytest.raises(OSError, match="Address already in use"): + make_zmq_listener(f"ipc://{path}", zmq.ROUTER) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(path)) + finally: + owner.close() + + def test_make_zmq_path(): assert make_zmq_path("tcp", "127.0.0.1", 5555) == "tcp://127.0.0.1:5555" assert make_zmq_path("tcp", "::1", 5555) == "tcp://[::1]:5555" diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index 768f51f437..1b23356134 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -23,7 +23,7 @@ from vllm.v1.engine.utils import ( EngineZmqAddresses, _dp_nodes_master_first, _node_ip_from_resources, - get_engine_zmq_addresses, + bind_engine_zmq_listeners, launch_core_engines, ) from vllm.v1.utils import APIServerProcessManager @@ -67,35 +67,27 @@ class _StubEngineCoreActor(EngineCoreActorMixin): # Module-level stub worker for the Ray-DP regression test. Must be importable -# by ``multiprocessing.spawn`` (no closures, no nesting). Mirrors the worker -# in ``tests/entrypoints/test_api_server_process_manager.py``. -def _bind_and_report_worker(listen_address, sock, args, client_config): - """Bind ROUTER/PULL with a kernel-assigned port, report the actual - endpoints back via ``actual_address_pipe``, then exit.""" +# by ``multiprocessing.spawn`` (no closures, no nesting). +def _adopt_listener_worker(listen_address, sock, args, client_config): + """Adopt the supervisor-bound ROUTER/PULL listeners, then exit.""" ctx = zmq.Context() try: in_sock = make_zmq_socket( - ctx, client_config["input_address"], zmq.ROUTER, bind=True + ctx, + client_config["input_address"], + zmq.ROUTER, + bind=True, + listener=client_config["input_listener"], ) out_sock = make_zmq_socket( - ctx, client_config["output_address"], zmq.PULL, bind=True + ctx, + client_config["output_address"], + zmq.PULL, + bind=True, + listener=client_config["output_listener"], ) - try: - pipe = client_config["actual_address_pipe"] - try: - pipe.send( - { - "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(), - "output_address": out_sock.getsockopt( - zmq.LAST_ENDPOINT - ).decode(), - } - ) - finally: - pipe.close() - finally: - in_sock.close(linger=0) - out_sock.close(linger=0) + in_sock.close(linger=0) + out_sock.close(linger=0) finally: ctx.term() @@ -327,17 +319,9 @@ def test_ray_dp_addresses_resolved_before_actor_creation( ``launch_core_engines`` Ray branch pickles ``addresses`` into each engine actor at ``.remote()`` time, and ``EngineCoreActorMixin._perform_handshakes`` is a no-op, so the actor uses that pickled snapshot for the rest of its - life. If ``run_multi_api_server`` allocates ``addresses`` as - ``tcp://host:0`` placeholders (its default), the actors hold placeholders - forever and DEALER-connect to port 0 — ZMQ ``connect`` is async and does - not raise, so the failure mode is a deterministic hang. - - The Ray-DP carve-out in ``run_multi_api_server`` forces - ``defer_api_server_ports=False`` when ``data_parallel_backend == "ray"`` - so addresses are pre-allocated in the driver and Ray pickles real ports - into each actor. This test mirrors that call-site logic and asserts the - actors hold real (non-placeholder) endpoints. If the carve-out is - removed without an alternative fix, the test fails. + life. The supervisor therefore binds the frontend listeners before actor + construction and Ray pickles their resolved endpoints. This test asserts + that each actor holds real (non-placeholder) endpoints. """ created_placement_groups: list[Any] = [] @@ -356,15 +340,10 @@ def test_ray_dp_addresses_resolved_before_actor_creation( vllm_config = _make_vllm_config_ray_dp_multinode() - # Mirror run_multi_api_server's address-allocation logic. The Ray DP - # carve-out forces pre-allocation so the addresses pickled into engine - # actors at .remote() time are real, not ``tcp://host:0``. - is_ray_dp = vllm_config.parallel_config.data_parallel_backend == "ray" - addresses = get_engine_zmq_addresses( - vllm_config, - num_api_servers=2, - defer_api_server_ports=not is_ray_dp, - ) + # Mirror run_multi_api_server: bind listeners in the supervisor so the + # addresses pickled into Ray actors contain kernel-assigned ports. + listeners = bind_engine_zmq_listeners(vllm_config, num_api_servers=2) + addresses = listeners.addresses sock = socket.socket() engine_manager: CoreEngineActorManager | None = None @@ -381,26 +360,17 @@ def test_ray_dp_addresses_resolved_before_actor_creation( engine_manager = engine_launch.engine_manager assert isinstance(engine_manager, CoreEngineActorManager) - # API-server children bind to the pre-allocated ports. + # API-server children adopt the already-bound listener FDs. api_server_manager = APIServerProcessManager( listen_address="tcp://127.0.0.1:0", sock=sock, args="test_args", num_servers=2, - input_addresses=addresses.inputs, - output_addresses=addresses.outputs, - target_server_fn=_bind_and_report_worker, + input_listeners=listeners.inputs, + output_listeners=listeners.outputs, + target_server_fn=_adopt_listener_worker, ) - # run_multi_api_server skips ``gather_actual_addresses`` for - # Ray DP (addresses are already real). Mirror that. - if not is_ray_dp: - actual_inputs, actual_outputs = ( - api_server_manager.gather_actual_addresses(timeout=15.0) - ) - addresses.inputs = actual_inputs - addresses.outputs = actual_outputs - # Snapshot what each Ray actor actually holds. actors = ( engine_manager.local_engine_actors + engine_manager.remote_engine_actors @@ -427,9 +397,9 @@ def test_ray_dp_addresses_resolved_before_actor_creation( assert scheme == "tcp", url assert port and int(port) > 0, ( f"Ray actor was pickled with placeholder address {url!r}; " - "``run_multi_api_server`` must pre-allocate ports for the " - "Ray DP backend so the actors hold real endpoints by the " - "time they DEALER-connect. See PR #42585 / Ray-DP " + "``run_multi_api_server`` must bind frontend listeners before " + "constructing Ray DP actors so they hold real endpoints when " + "they DEALER-connect. See PR #42585 / Ray-DP " "multi-API-server regression." ) diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index 4c9410488b..94e6545c08 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -328,19 +328,10 @@ def run_multi_api_server(args: argparse.Namespace): None ) - from vllm.v1.engine.utils import get_engine_zmq_addresses + from vllm.v1.engine.utils import bind_engine_zmq_listeners - # Defer port allocation to the child's bind() to avoid TOCTOU, except - # for Rust front-end and Ray DP, which can't see the post-bind rebind - # (CLI-arg subprocess / pickled-into-actor snapshot respectively) and - # so pre-allocate driver-side -- reintroducing the original race only - # there. - is_ray_dp = parallel_config.data_parallel_backend == "ray" - addresses = get_engine_zmq_addresses( - vllm_config, - num_api_servers, - defer_api_server_ports=not (rust_frontend_path or is_ray_dp), - ) + zmq_listeners = bind_engine_zmq_listeners(vllm_config, num_api_servers) + addresses = zmq_listeners.addresses with launch_core_engines( vllm_config, executor_class, log_stats, addresses @@ -365,8 +356,8 @@ def run_multi_api_server(args: argparse.Namespace): sock=sock, grpc_sock=grpc_sock, args=args, - input_address=addresses.inputs[0], - output_address=addresses.outputs[0], + input_listener=zmq_listeners.inputs[0], + output_listener=zmq_listeners.outputs[0], engine_start_index=expected_engine_start_index, engine_count=expected_engine_count, data_parallel_size=parallel_config.data_parallel_size, @@ -379,22 +370,12 @@ def run_multi_api_server(args: argparse.Namespace): sock=sock, args=args, num_servers=num_api_servers, - input_addresses=addresses.inputs, - output_addresses=addresses.outputs, + input_listeners=zmq_listeners.inputs, + output_listeners=zmq_listeners.outputs, stats_update_address=stats_update_address, tensor_queue=engine_launch.tensor_queue, ) - if not is_ray_dp: - # Forward each child's bound endpoints to the engine handshake - # (runs on ``with`` exit). Skipped for Ray DP, where addresses - # are pre-allocated above and Ray actors already hold them. - actual_inputs, actual_outputs = ( - api_server_manager.gather_actual_addresses() - ) - addresses.inputs = actual_inputs - addresses.outputs = actual_outputs - # Set frontend processes to watch during engine startup. # If any of these processes exit before the engines are up, the engine startup # will be aborted with an error. diff --git a/vllm/utils/network_utils.py b/vllm/utils/network_utils.py index fa24d4d5f9..6bc4f9c9fd 100644 --- a/vllm/utils/network_utils.py +++ b/vllm/utils/network_utils.py @@ -11,6 +11,7 @@ from collections.abc import ( Iterator, Sequence, ) +from dataclasses import dataclass from typing import Any from uuid import uuid4 @@ -25,6 +26,23 @@ from vllm.logger import init_logger logger = init_logger(__name__) +@dataclass +class ZmqListener: + """A supervisor-bound listener that a ZMQ socket can adopt.""" + + address: str + socket: socket.socket + + def close(self) -> None: + self.socket.close() + + def cleanup(self) -> None: + self.close() + if self.address.startswith("ipc://"): + with contextlib.suppress(FileNotFoundError): + os.unlink(self.address.removeprefix("ipc://")) + + def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]): for sock in sockets: if sock is not None: @@ -307,6 +325,27 @@ def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str: # Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501 +# Apply the existing make_zmq_socket memory policy to every socket owner. +_ZMQ_LARGE_BUFFER_SIZE = 512 * 1024**2 +_ZMQ_LARGE_BUFFER_MIN_TOTAL_MEMORY = 32 * 1024**3 +_ZMQ_LARGE_BUFFER_MIN_AVAILABLE_MEMORY = 16 * 1024**3 + + +def _get_zmq_socket_buffer_size() -> int: + """Choose the shared libzmq and inherited-listener buffer policy.""" + mem = psutil.virtual_memory() + # For systems with substantial memory (>32GB total, >16GB available): + # - Set a large 0.5GB buffer to improve throughput + # For systems with less memory: + # - Use system default (-1) to avoid excessive memory consumption + if ( + mem.total > _ZMQ_LARGE_BUFFER_MIN_TOTAL_MEMORY + and mem.available > _ZMQ_LARGE_BUFFER_MIN_AVAILABLE_MEMORY + ): + return _ZMQ_LARGE_BUFFER_SIZE + return -1 + + def make_zmq_socket( ctx: zmq.asyncio.Context | zmq.Context, # type: ignore[name-defined] path: str, @@ -315,22 +354,19 @@ def make_zmq_socket( identity: bytes | None = None, linger: int | None = None, router_handover: bool = False, + listener: socket.socket | None = None, ) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined] - """Make a ZMQ socket with the proper bind/connect semantics.""" - mem = psutil.virtual_memory() - socket = ctx.socket(socket_type) + """Make a ZMQ socket with the proper bind/connect semantics. - # Calculate buffer size based on system memory - total_mem = mem.total / 1024**3 - available_mem = mem.available / 1024**3 - # For systems with substantial memory (>32GB total, >16GB available): - # - Set a large 0.5GB buffer to improve throughput - # For systems with less memory: - # - Use system default (-1) to avoid excessive memory consumption - buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1 + When supplied, ``listener`` is detached and its fd moves to libzmq. + """ + socket = ctx.socket(socket_type) + buf_size = _get_zmq_socket_buffer_size() if bind is None: bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB) + if listener is not None and not bind: + raise ValueError("An inherited ZMQ listener requires bind=True") if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER): socket.setsockopt(zmq.RCVHWM, 0) @@ -353,6 +389,9 @@ def make_zmq_socket( if socket_type == zmq.XPUB: socket.setsockopt(zmq.XPUB_VERBOSE, True) + if listener is not None: + socket.setsockopt(zmq.USE_FD, listener.detach()) + # Determine if the path is a TCP socket with an IPv6 address. # Enable IPv6 on the zmq socket if so. scheme, host, _ = split_zmq_path(path) @@ -367,6 +406,48 @@ def make_zmq_socket( return socket +def make_zmq_listener(path: str, socket_type: Any) -> ZmqListener: + """Bind and listen on a raw socket for later ZMQ adoption.""" + scheme, host, port = split_zmq_path(path) + if scheme == "tcp": + family = socket.AF_INET6 if is_valid_ipv6_address(host) else socket.AF_INET + bind_address: str | tuple[str, int] = (host, int(port)) + elif scheme == "ipc": + family = socket.AF_UNIX + bind_address = path.removeprefix("ipc://") + os.makedirs(os.path.dirname(bind_address), exist_ok=True) + else: + raise ValueError(f"Cannot inherit a {scheme} ZMQ listener") + + listener = socket.socket(family, socket.SOCK_STREAM) + bound = False + try: + buf_size = _get_zmq_socket_buffer_size() + if buf_size >= 0: + # Accepted stream sockets inherit these settings from the listener. + # ZMQ_USE_FD adopts this pre-created socket, so the policy belongs + # here as well as on the eventual libzmq socket. + if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER): + listener.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, buf_size) + if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER): + listener.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, buf_size) + + listener.bind(bind_address) + bound = True + listener.listen() + if scheme == "tcp": + path = get_tcp_uri(host, listener.getsockname()[1]) + return ZmqListener(address=path, socket=listener) + except BaseException: + listener.close() + # A failed bind leaves any existing pathname owned by its listener. + if scheme == "ipc" and bound: + assert isinstance(bind_address, str) + with contextlib.suppress(FileNotFoundError): + os.unlink(bind_address) + raise + + @contextlib.contextmanager def zmq_socket_ctx( path: str, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 88fdc1d1e7..74520036a7 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -11,7 +11,6 @@ from collections import Counter, defaultdict from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import Future from dataclasses import dataclass -from multiprocessing.connection import Connection from multiprocessing.queues import Queue from threading import Thread from typing import Any, TypeAlias, TypeVar @@ -56,7 +55,8 @@ from vllm.v1.engine.tensor_ipc import TensorIpcSender from vllm.v1.engine.utils import ( CoreEngineActorManager, CoreEngineProcManager, - get_engine_zmq_addresses, + EngineZmqAddresses, + get_engine_zmq_bind_addresses, launch_core_engines, ) from vllm.v1.executor import Executor @@ -634,63 +634,50 @@ class MPClient(EngineCoreClient): self.stats_update_address = client_addresses.get("stats_update_address") # Tensor queues passed via client_addresses for multi-API-server case tensor_queue = client_addresses.get("tensor_queue") + input_listener = client_addresses.get("input_listener") + output_listener = client_addresses.get("output_listener") self.input_socket = self.resources.input_socket = make_zmq_socket( self.ctx, input_address, zmq.ROUTER, bind=True, router_handover=enable_input_socket_handover, + listener=input_listener, ) self.resources.output_socket = make_zmq_socket( - self.ctx, output_address, zmq.PULL + self.ctx, + output_address, + zmq.PULL, + listener=output_listener, ) - # Report bound endpoints back so the parent can forward - # them to engines (mirrors the DPCoordinator pattern). - actual_address_pipe: Connection | None = client_addresses.get( - "actual_address_pipe" - ) - if actual_address_pipe is not None: - try: - actual_input = self.input_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - actual_output = self.resources.output_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - actual_address_pipe.send( - { - "input_address": actual_input, - "output_address": actual_output, - } - ) - finally: - actual_address_pipe.close() # Engines are managed externally: this process does not fork # them, so there is no fork race; start the MM warmup now. self._start_mm_warmup() else: # Engines are managed by this client. - addresses = get_engine_zmq_addresses(vllm_config) + bind_addresses = get_engine_zmq_bind_addresses(vllm_config) self.input_socket = self.resources.input_socket = make_zmq_socket( self.ctx, - addresses.inputs[0], + bind_addresses.inputs[0], zmq.ROUTER, bind=True, router_handover=enable_input_socket_handover, ) self.resources.output_socket = make_zmq_socket( - self.ctx, addresses.outputs[0], zmq.PULL + self.ctx, bind_addresses.outputs[0], zmq.PULL ) - # Resolve ``tcp://host:0`` placeholders to bound endpoints - # before engines DEALER-connect. No-op for IPC. - addresses.inputs[0] = self.input_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() - addresses.outputs[0] = self.resources.output_socket.getsockopt( - zmq.LAST_ENDPOINT - ).decode() + # EngineZmqAddresses contains only endpoints resolved by the + # sockets that own the bind operation. + addresses = EngineZmqAddresses( + inputs=[self.input_socket.getsockopt(zmq.LAST_ENDPOINT).decode()], + outputs=[ + self.resources.output_socket.getsockopt( + zmq.LAST_ENDPOINT + ).decode() + ], + ) with launch_core_engines( vllm_config, executor_class, log_stats, addresses diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 49e6884ac9..4f1be3bacd 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -23,9 +23,10 @@ from vllm.platforms import current_platform from vllm.ray.ray_env import get_env_vars_to_copy from vllm.utils import numa_utils from vllm.utils.network_utils import ( - get_open_port, + ZmqListener, get_open_zmq_ipc_path, get_tcp_uri, + make_zmq_listener, zmq_socket_ctx, ) from vllm.utils.system_utils import get_mp_context, set_env_var @@ -82,8 +83,18 @@ class CoreEngine: self.state = CoreEngineState.NEW +@dataclass +class EngineZmqBindAddresses: + """Frontend endpoints before their owning sockets bind.""" + + inputs: list[str] + outputs: list[str] + + @dataclass class EngineZmqAddresses: + """Resolved frontend endpoints advertised to engine processes.""" + # ZMQ input socket addresses for each front-end client (requests) inputs: list[str] # ZMQ output socket addresses for each front-end client (responses) @@ -98,6 +109,27 @@ class EngineZmqAddresses: frontend_stats_publish_address: str | None = None +@dataclass +class EngineZmqListeners: + inputs: list[ZmqListener] + outputs: list[ZmqListener] + + @property + def addresses(self) -> EngineZmqAddresses: + return EngineZmqAddresses( + inputs=[listener.address for listener in self.inputs], + outputs=[listener.address for listener in self.outputs], + ) + + def cleanup(self) -> None: + for listener in (*self.inputs, *self.outputs): + listener.cleanup() + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.cleanup() + + @dataclass class EngineHandshakeMetadata: """Metadata sent to each engine process during startup handshake, @@ -1079,22 +1111,16 @@ class CoreEngineActorManager: ray.util.remove_placement_group(pg) -def get_engine_zmq_addresses( +def get_engine_zmq_bind_addresses( vllm_config: VllmConfig, num_api_servers: int = 1, - *, - defer_api_server_ports: bool = True, -) -> EngineZmqAddresses: - """Allocate ZMQ addresses for engine-client communication. +) -> EngineZmqBindAddresses: + """Create ZMQ endpoint templates for engine-client communication. - By default each TCP address is a ``tcp://host:0`` placeholder; the - consumer (API-server child or single-process ``MPClient``) binds, then - recovers the kernel-assigned port via ``getsockopt(zmq.LAST_ENDPOINT)`` - and writes it back into ``addresses`` before the engine handshake. - - Set ``defer_api_server_ports=False`` only when the consumer cannot - report a bound port back (e.g. the Rust front-end). IPC paths are - unaffected.""" + Each TCP address is a ``tcp://host:0`` placeholder. A single-process + ``MPClient`` binds its own ZMQ sockets. Cross-process frontends use + ``bind_engine_zmq_listeners`` so the supervisor binds raw listeners before + launching engines and transfers listener ownership to each frontend.""" parallel_config = vllm_config.parallel_config local_engine_count = parallel_config.data_parallel_size_local local_start_index = parallel_config.data_parallel_rank_local @@ -1119,14 +1145,39 @@ def get_engine_zmq_addresses( def _addr() -> str: if client_local_only: return get_open_zmq_ipc_path() - return get_tcp_uri(host, 0 if defer_api_server_ports else get_open_port()) + return get_tcp_uri(host, 0) - return EngineZmqAddresses( + return EngineZmqBindAddresses( inputs=[_addr() for _ in range(num_api_servers)], outputs=[_addr() for _ in range(num_api_servers)], ) +def bind_engine_zmq_listeners( + vllm_config: VllmConfig, + num_api_servers: int = 1, +) -> EngineZmqListeners: + """Bind frontend transport listeners before launching engines.""" + bind_addresses = get_engine_zmq_bind_addresses( + vllm_config, + num_api_servers, + ) + inputs: list[ZmqListener] = [] + outputs: list[ZmqListener] = [] + try: + inputs.extend( + make_zmq_listener(address, zmq.ROUTER) for address in bind_addresses.inputs + ) + outputs.extend( + make_zmq_listener(address, zmq.PULL) for address in bind_addresses.outputs + ) + return EngineZmqListeners(inputs=inputs, outputs=outputs) + except BaseException: + for listener in (*inputs, *outputs): + listener.cleanup() + raise + + FrontendProcess = BaseProcess | _SubprocessWrapper diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index ac877263a5..868864f0ca 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -29,7 +29,7 @@ from torch.autograd.profiler import record_function import vllm.envs as envs from vllm.logger import init_logger from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message -from vllm.utils.network_utils import get_open_zmq_ipc_path, get_tcp_uri +from vllm.utils.network_utils import ZmqListener, get_open_zmq_ipc_path, get_tcp_uri from vllm.utils.system_utils import decorate_logs, kill_process_tree, set_process_title from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.core.sched.output import SchedulerOutput @@ -177,28 +177,22 @@ class APIServerProcessManager: sock: Any, args: argparse.Namespace, num_servers: int, - input_addresses: list[str], - output_addresses: list[str], + input_listeners: list[ZmqListener], + output_listeners: list[ZmqListener], target_server_fn: Callable | None = None, stats_update_address: str | None = None, tensor_queue: Queue | None = None, ): """Initialize and start API server worker processes. - ``input_addresses``/``output_addresses`` may contain - ``tcp://host:0`` placeholders; each child must report the actual - bound endpoint over its ``actual_address_pipe`` in ``client_config`` - and the parent collects them via - :py:meth:`gather_actual_addresses`. - Args: target_server_fn: Override function to call for each API server process listen_address: Address to listen for client connections sock: Socket for client connections args: Command line arguments num_servers: Number of API server processes to start - input_addresses: Input addresses for each API server - output_addresses: Output addresses for each API server + input_listeners: Input listeners for each API server + output_listeners: Output listeners for each API server stats_update_address: Optional stats update address tensor_queue: Optional tensor IPC queue for sharing MM tensors @@ -209,7 +203,6 @@ class APIServerProcessManager: spawn_context = multiprocessing.get_context("spawn") self.processes: list[BaseProcess] = [] - self._address_pipes: list[connection.Connection] = [] admission_counters = None if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None: @@ -220,126 +213,65 @@ class APIServerProcessManager: SharedAdmissionStats.num_counters(num_servers), ) - for i, in_addr, out_addr in zip( - range(num_servers), input_addresses, output_addresses - ): - client_config: dict[str, Any] = { - "input_address": in_addr, - "output_address": out_addr, - "client_count": num_servers, - "client_index": i, - } - if admission_counters is not None: - client_config["mp_admission_counters"] = admission_counters - if stats_update_address is not None: - client_config["stats_update_address"] = stats_update_address - if tensor_queue is not None: - client_config["tensor_queue"] = tensor_queue - - parent_recv, child_send = spawn_context.Pipe(duplex=False) - self._address_pipes.append(parent_recv) - client_config["actual_address_pipe"] = child_send - - proc = spawn_context.Process( - target=target_server_fn or run_api_server_worker_proc, - name=f"ApiServer_{i}", - args=(listen_address, sock, args, client_config), + listeners = [*input_listeners, *output_listeners] + self._listeners = listeners + try: + listener_pairs = zip( + range(num_servers), input_listeners, output_listeners, strict=True ) - self.processes.append(proc) - proc.start() + for i, input_listener, output_listener in listener_pairs: + client_config: dict[str, Any] = { + "input_address": input_listener.address, + "output_address": output_listener.address, + "input_listener": input_listener.socket, + "output_listener": output_listener.socket, + "client_count": num_servers, + "client_index": i, + } + if admission_counters is not None: + client_config["mp_admission_counters"] = admission_counters + if stats_update_address is not None: + client_config["stats_update_address"] = stats_update_address + if tensor_queue is not None: + client_config["tensor_queue"] = tensor_queue - # Drop parent's write end so reader sees EOF on child death. - child_send.close() + proc = spawn_context.Process( + target=target_server_fn or run_api_server_worker_proc, + name=f"ApiServer_{i}", + args=(listen_address, sock, args, client_config), + ) + self.processes.append(proc) + proc.start() + input_listener.close() + output_listener.close() + except BaseException: + _shutdown_processes_and_cleanup_listeners(self.processes, listeners) + raise logger.info("Started %d API server processes", len(self.processes)) # Shutdown only the API server processes on garbage collection # The extra processes are managed by their owners - self._finalizer = weakref.finalize(self, shutdown, self.processes) - - def gather_actual_addresses( - self, - timeout: float = envs.VLLM_ENGINE_READY_TIMEOUT_S, - ) -> tuple[list[str], list[str]]: - """Return (inputs, outputs) reported by each child, indexed by - ``client_index``. Raises ``RuntimeError`` on timeout or premature - child exit.""" - n = len(self._address_pipes) - inputs: list[str | None] = [None] * n - outputs: list[str | None] = [None] * n - pending: dict[connection.Connection, int] = { - pipe: i for i, pipe in enumerate(self._address_pipes) - } - sentinel_to_idx: dict[Any, int] = { - proc.sentinel: i for i, proc in enumerate(self.processes) - } - - deadline = time.monotonic() + timeout - try: - while pending: - remaining = deadline - time.monotonic() - if remaining <= 0: - missing = [self.processes[i].name for i in pending.values()] - raise RuntimeError( - f"Timed out after {timeout:.1f}s waiting for " - f"API server(s) to report bound ZMQ addresses: " - f"{missing}" - ) - waitables: list[Any] = list(pending.keys()) + list( - sentinel_to_idx.keys() - ) - ready = connection.wait(waitables, timeout=remaining) - # Drain pipes before checking sentinels: a child that sent - # its message and then exited can surface both events in - # the same poll, and we must record the success first. - for item in ready: - if isinstance(item, connection.Connection) and item in pending: - idx = pending.pop(item) - try: - msg: dict[str, str] = item.recv() - except EOFError as e: - raise RuntimeError( - f"API server {self.processes[idx].name} " - f"closed its address pipe without " - f"reporting its bound ZMQ addresses" - ) from e - inputs[idx] = msg["input_address"] - outputs[idx] = msg["output_address"] - item.close() - for item in ready: - if item in sentinel_to_idx: - idx = sentinel_to_idx.pop(item) - pipe = self._address_pipes[idx] - if pipe in pending: - proc = self.processes[idx] - raise RuntimeError( - f"API server process {proc.name} exited " - f"(code={proc.exitcode}) before reporting " - f"its bound ZMQ addresses" - ) - finally: - for pipe in pending: - with contextlib.suppress(Exception): - pipe.close() - - return inputs, outputs # type: ignore[return-value] + self._finalizer = weakref.finalize( + self, + _shutdown_processes_and_cleanup_listeners, + self.processes, + listeners, + ) def shutdown(self, timeout: float | None = None) -> None: """Shutdown API server processes with configurable timeout""" - for pipe in self._address_pipes: - with contextlib.suppress(Exception): - pipe.close() - self._address_pipes = [] - if self._finalizer.detach() is not None: - shutdown(self.processes, timeout=timeout) + _shutdown_processes_and_cleanup_listeners( + self.processes, self._listeners, timeout=timeout + ) class RustFrontendProcessManager: """Manages a single Rust frontend subprocess. Launches the Rust vllm-rs binary in 'frontend' mode, passing the - listening socket fd and ZMQ transport addresses. Provides the same + HTTP and ZMQ listener file descriptors. Provides the same interface as APIServerProcessManager for process monitoring. """ @@ -348,8 +280,8 @@ class RustFrontendProcessManager: binary_path: str, sock: Any, args: argparse.Namespace, - input_address: str, - output_address: str, + input_listener: ZmqListener, + output_listener: ZmqListener, engine_start_index: int, engine_count: int, data_parallel_size: int, @@ -359,18 +291,22 @@ class RustFrontendProcessManager: import os import subprocess - fd = sock.fileno() - os.set_inheritable(fd, True) + listen_fd = sock.fileno() + input_fd = input_listener.socket.fileno() + output_fd = output_listener.socket.fileno() + inherited_fds = [listen_fd, input_fd, output_fd] + for fd in inherited_fds: + os.set_inheritable(fd, True) cmd = [ binary_path, "frontend", "--listen-fd", - str(fd), - "--input-address", - input_address, - "--output-address", - output_address, + str(listen_fd), + "--input-listener-fd", + str(input_fd), + "--output-listener-fd", + str(output_fd), "--engine-start-index", str(engine_start_index), "--engine-count", @@ -378,12 +314,11 @@ class RustFrontendProcessManager: "--data-parallel-size", str(data_parallel_size), ] - pass_fds = [fd] if grpc_sock is not None: grpc_fd = grpc_sock.fileno() os.set_inheritable(grpc_fd, True) cmd.extend(["--grpc-listen-fd", str(grpc_fd)]) - pass_fds.append(grpc_fd) + inherited_fds.append(grpc_fd) if stats_update_address is not None: cmd.extend(["--coordinator-address", stats_update_address]) from vllm.entrypoints.serve.utils.api_utils import jsonify_non_default_args @@ -425,18 +360,34 @@ class RustFrontendProcessManager: redacted_json = json.dumps(redact_sensitive_args(args_dict), sort_keys=True) logger.info("Launching Rust frontend: %s", " ".join(cmd[:-1] + [redacted_json])) - self._proc = subprocess.Popen(cmd, pass_fds=pass_fds) + self._listeners = [input_listener, output_listener] + try: + self._proc = subprocess.Popen(cmd, pass_fds=inherited_fds) + except BaseException: + for listener in self._listeners: + listener.cleanup() + raise + else: + input_listener.close() + output_listener.close() # Create a process wrapper with a sentinel fd for monitoring self.processes: list[_SubprocessWrapper] = [ _SubprocessWrapper(self._proc, "RustFrontend") ] - self._finalizer = weakref.finalize(self, _shutdown_subprocesses, self.processes) + self._finalizer = weakref.finalize( + self, + _shutdown_subprocesses_and_cleanup_listeners, + self.processes, + self._listeners, + ) def shutdown(self, timeout: float | None = None) -> None: if self._finalizer.detach() is not None: - _shutdown_subprocesses(self.processes, timeout=timeout) + _shutdown_subprocesses_and_cleanup_listeners( + self.processes, self._listeners, timeout=timeout + ) class _SubprocessWrapper: @@ -533,6 +484,18 @@ def _shutdown_subprocesses( logger.debug_once("[shutdown] Subprocess manager: complete") +def _shutdown_subprocesses_and_cleanup_listeners( + procs: list[_SubprocessWrapper], + listeners: list[ZmqListener], + timeout: float | None = None, +) -> None: + try: + _shutdown_subprocesses(procs, timeout=timeout) + finally: + for listener in listeners: + listener.cleanup() + + def run_api_server_worker_proc( listen_address, sock, args, client_config=None, **uvicorn_kwargs ) -> None: @@ -687,6 +650,18 @@ def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None: logger.debug_once("[shutdown] Process manager: complete") +def _shutdown_processes_and_cleanup_listeners( + procs: list[BaseProcess], + listeners: list[ZmqListener], + timeout: float | None = None, +) -> None: + try: + shutdown(procs, timeout=timeout) + finally: + for listener in listeners: + listener.cleanup() + + def copy_slice( from_tensor: torch.Tensor, to_tensor: torch.Tensor, length: int ) -> torch.Tensor: