[Rust Frontend] Wait for mock engine endpoints before ZMQ connect (#47965)
Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Reid committed
Jul 16, 2026 at 17:20 UTC
f44f3d6f798b1ab2e69ecaf82fab684ead5ba7f0
1 file changed
+34
-24
rust/src/engine-core-client/src/mock_engine.rs
+34
-24
@@ -1,10 +1,9 @@
1
// SPDX-License-Identifier: Apache-2.0
2
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-use std::path::Path;
4
use std::time::Duration;
5
7
-use tokio::time::timeout;
6
+use tokio::time::{sleep, timeout};
7
use zeromq::prelude::{Socket, SocketRecv, SocketSend};
8
use zeromq::util::PeerIdentity;
9
use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
@@ -32,7 +31,7 @@ pub struct MockEngineConfig {
31
/// Engine-ready payload reported after INIT, including max model length,
32
/// KV block count, and dtype.
33
pub ready_response: EngineCoreReadyResponse,
35
- /// Maximum time to wait for IPC endpoints to appear before connecting.
34
+ /// Maximum time to wait for endpoints to accept connections.
35
pub connect_timeout: Duration,
36
}
37
@@ -121,22 +120,33 @@ fn peer_identity(engine_id: impl Into<EngineId>) -> Result<PeerIdentity> {
120
})
121
}
122
124
-/// Wait for an IPC endpoint path to appear before attempting to connect.
125
-async fn wait_for_ipc_endpoint(endpoint: &str, connect_timeout: Duration) -> Result<()> {
126
- let Some(socket_path) = endpoint.strip_prefix("ipc://") else {
127
- return Ok(());
128
- };
129
-
130
- timeout(connect_timeout, async {
131
- while !Path::new(socket_path).exists() {
132
- tokio::time::sleep(Duration::from_millis(20)).await;
133
- }
134
- })
135
- .await
136
- .map_err(|_| Error::HandshakeTimeout {
137
- stage: "mock engine IPC endpoint",
138
- timeout: connect_timeout,
139
- })
123
+/// Wait for an endpoint to accept connections before attempting the ZMQ connect.
124
+async fn wait_for_endpoint(endpoint: &str, connect_timeout: Duration) -> Result<()> {
125
+ if let Some(socket_path) = endpoint.strip_prefix("ipc://") {
126
+ timeout(connect_timeout, async {
127
+ while tokio::net::UnixStream::connect(socket_path).await.is_err() {
128
+ sleep(Duration::from_millis(20)).await;
129
+ }
130
+ })
131
+ .await
132
+ .map_err(|_| Error::HandshakeTimeout {
133
+ stage: "mock engine IPC endpoint",
134
+ timeout: connect_timeout,
135
+ })
136
+ } else if let Some(address) = endpoint.strip_prefix("tcp://") {
137
+ timeout(connect_timeout, async {
138
+ while tokio::net::TcpStream::connect(address).await.is_err() {
139
+ sleep(Duration::from_millis(20)).await;
140
+ }
141
+ })
142
+ .await
143
+ .map_err(|_| Error::HandshakeTimeout {
144
+ stage: "mock engine TCP endpoint",
145
+ timeout: connect_timeout,
146
+ })
147
+ } else {
148
+ Ok(())
149
+ }
150
}
151
152
/// Encode the engine-ready response sent on input socket registration.
@@ -151,7 +161,7 @@ pub async fn connect_to_frontend(
161
config: MockEngineConfig,
162
) -> Result<MockEngineSockets> {
163
let engine_handshake = engine_handshake.as_ref();
154
- wait_for_ipc_endpoint(engine_handshake, config.connect_timeout).await?;
164
+ wait_for_endpoint(engine_handshake, config.connect_timeout).await?;
165
166
let peer_identity = peer_identity(engine_id)?;
167
let mut options = SocketOptions::default();
@@ -192,8 +202,8 @@ pub async fn connect_to_frontend(
202
for (input_address, output_address) in
203
init.addresses.inputs.iter().zip(init.addresses.outputs.iter())
204
{
195
- wait_for_ipc_endpoint(input_address, config.connect_timeout).await?;
196
- wait_for_ipc_endpoint(output_address, config.connect_timeout).await?;
205
+ wait_for_endpoint(input_address, config.connect_timeout).await?;
206
+ wait_for_endpoint(output_address, config.connect_timeout).await?;
207
208
let mut input_options = SocketOptions::default();
209
input_options.peer_identity(peer_identity.clone());
@@ -260,8 +270,8 @@ pub async fn connect_to_bootstrapped_frontend(
270
) -> Result<(DealerSocket, PushSocket)> {
271
let input_address = input_address.as_ref();
272
let output_address = output_address.as_ref();
263
- wait_for_ipc_endpoint(input_address, config.connect_timeout).await?;
264
- wait_for_ipc_endpoint(output_address, config.connect_timeout).await?;
273
+ wait_for_endpoint(input_address, config.connect_timeout).await?;
274
+ wait_for_endpoint(output_address, config.connect_timeout).await?;
275
276
let peer_identity = peer_identity(engine_id)?;
277
let mut input_options = SocketOptions::default();