Compare commits

...
Author SHA1 Message Date
Bugen Zhao 58feec155a observe ready timeout & refine logs
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-28 09:10:12 +00:00
Bugen ZhaoandOpenAI Codex 93dec05d69 Gate EngineCore reattach on ZMQ readiness
Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-27 15:13:05 +00:00
Bugen ZhaoandOpenAI Codex 707fb0de8f Refine EngineCore reattach session workflow
Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-27 15:11:06 +00:00
Bugen ZhaoandOpenAI Codex fe3b526be4 Add EngineCore reattach sessions for frontend development
Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-27 15:11:06 +00:00
18 changed files with 771 additions and 30 deletions
+24 -1
View File
@@ -49,7 +49,7 @@ are started elsewhere and this node should run only the Rust frontend. The front
the global `--data-parallel-size` to determine how many engines it expects to join the shared handshake.
```bash
vllm serve Qwen/Qwen3-0.6B \
vllm-rs serve Qwen/Qwen3-0.6B \
--headless \
--data-parallel-address 127.0.0.1 \
--data-parallel-rpc-port 62100 \
@@ -67,6 +67,29 @@ vllm-rs serve Qwen/Qwen3-0.6B \
--data-parallel-size-local 0
```
For frontend development with a single engine, add `--engine-session` to the
frontend command:
```bash
vllm-rs serve Qwen/Qwen3-0.6B \
--data-parallel-address 127.0.0.1 \
--data-parallel-rpc-port 62100 \
--data-parallel-size 1 \
--engine-session /tmp/vllm-rs-engine-session.json
```
When the session file does not exist, the frontend performs the normal
handshake and writes the engine transport state to it. Stop the frontend only,
then run the same command again to bind the saved endpoints and reconnect to
the already loaded EngineCore. `--engine-session` implies frontend-only
external-engine mode, so `--data-parallel-size-local 0` is not required.
Reattach is currently a development-only DP=1 workflow. It does not preserve
HTTP request state: if the frontend exits with active requests, their clients
disconnect, EngineCore continues them, and the replacement frontend discards
their stale outputs while accepting new work. Prefer restarting while idle to
avoid wasted engine work. Run only one frontend for a session at a time, and
delete the session file whenever EngineCore is restarted.
To build the `vllm-rs` in isolation:
```bash
@@ -66,6 +66,7 @@ async fn main() -> Result<()> {
ready_timeout,
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode: None,
model_name: args.model.clone(),
+60 -13
View File
@@ -423,8 +423,43 @@ impl SharedRuntimeArgs {
engine_count: usize,
local_input_address: Option<String>,
local_output_address: Option<String>,
session_path: Option<PathBuf>,
) -> Config {
let ready_timeout = self.ready_timeout();
self.into_standalone_config(
listener_mode,
TransportMode::HandshakeOwner {
handshake_address,
advertised_host,
engine_count,
ready_timeout,
local_input_address,
local_output_address,
session_path,
},
)
}
fn into_reattach_config(
self,
listener_mode: HttpListenerMode,
session_path: PathBuf,
) -> Config {
let ready_timeout = self.ready_timeout();
self.into_standalone_config(
listener_mode,
TransportMode::Reattach {
path: session_path,
ready_timeout,
},
)
}
fn into_standalone_config(
self,
listener_mode: HttpListenerMode,
transport_mode: TransportMode,
) -> Config {
let shutdown_timeout = self.shutdown_timeout();
let keep_alive_timeout = self.keep_alive_timeout();
let api_server_options = self.api_server_options();
@@ -433,14 +468,7 @@ impl SharedRuntimeArgs {
let profiler = self.profiler();
Config {
transport_mode: TransportMode::HandshakeOwner {
handshake_address,
advertised_host,
engine_count,
ready_timeout,
local_input_address,
local_output_address,
},
transport_mode,
coordinator_mode: CoordinatorMode::MaybeInProc,
model: self.model,
served_model_name: self.served_model_name,
@@ -611,6 +639,12 @@ pub struct ServeArgs {
#[arg(long)]
pub uds: Option<String>,
/// Development session used to reconnect this frontend to a running EngineCore.
///
/// This implies frontend-only external-engine mode.
#[arg(long)]
pub engine_session: Option<PathBuf>,
/// Flag to print debug information about CLI argument parsing and exit.
#[educe(Debug(ignore))]
#[arg(long, hide = true, env = "VLLM_RS_DEBUG_CLI")]
@@ -626,12 +660,14 @@ pub struct ServeArgs {
}
impl ServeArgs {
/// Build the OpenAI-server runtime config used after the managed Python
/// engine starts.
/// Return whether the Rust frontend should connect to an externally owned
/// EngineCore instead of spawning a managed local engine.
pub fn uses_external_engine(&self) -> bool {
self.engine_session.is_some() || self.managed_engine.data_parallel_size_local == Some(0)
}
/// Build the OpenAI-server runtime config for the selected engine transport.
pub fn to_frontend_config(&self, handshake_address: String) -> Config {
// Prefer IPC sockets for local engine input/output.
let (local_input_address, local_output_address) =
self.managed_engine.frontend_local_only().then(frontend_ipc_addresses).unzip();
let listener_mode = match &self.uds {
Some(path) => HttpListenerMode::BindUnix { path: path.clone() },
None => HttpListenerMode::BindTcp {
@@ -640,6 +676,16 @@ impl ServeArgs {
},
};
if let Some(session_path) = self.engine_session.as_ref()
&& session_path.exists()
{
return self.runtime.clone().into_reattach_config(listener_mode, session_path.clone());
}
// Prefer IPC sockets for local engine input/output.
let (local_input_address, local_output_address) =
self.managed_engine.frontend_local_only().then(frontend_ipc_addresses).unzip();
self.runtime.clone().into_managed_config(
listener_mode,
handshake_address,
@@ -647,6 +693,7 @@ impl ServeArgs {
self.managed_engine.data_parallel_size,
local_input_address,
local_output_address,
self.engine_session.clone(),
)
}
+52
View File
@@ -51,6 +51,7 @@ fn serve_args_forward_python_flags_with_separator() {
host: "127.0.0.1",
port: 8000,
uds: None,
engine_session: None,
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
engine_ready_timeout_secs: 600,
@@ -1318,6 +1319,7 @@ fn serve_args_accept_handshake_aliases() {
host: "127.0.0.1",
port: 8000,
uds: None,
engine_session: None,
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
engine_ready_timeout_secs: 600,
@@ -1427,6 +1429,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
ready_timeout,
local_input_address,
local_output_address,
session_path,
} = &config.transport_mode
else {
panic!("expected handshake-owned transport");
@@ -1446,6 +1449,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
.is_some_and(|address| address.starts_with("ipc://"))
);
assert_ne!(local_input_address, local_output_address);
assert!(session_path.is_none());
expect![[r#"
Config {
@@ -1460,6 +1464,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
local_output_address: Some(
"<ipc output>",
),
session_path: None,
},
coordinator_mode: MaybeInProc,
model: "Qwen/Qwen3-0.6B",
@@ -1510,11 +1515,57 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
ready_timeout: *ready_timeout,
local_input_address: Some("<ipc input>".to_string()),
local_output_address: Some("<ipc output>".to_string()),
session_path: None,
},
..config.clone()
});
}
#[test]
fn serve_engine_session_selects_handshake_then_reattach() {
let session_path = std::env::temp_dir().join(format!(
"vllm-rs-engine-session-{}.json",
uuid::Uuid::new_v4()
));
let session_arg = session_path.to_string_lossy().into_owned();
let cli = Cli::try_parse_from([
"vllm-rs",
"serve",
"Qwen/Qwen3-0.6B",
"--engine-session",
&session_arg,
"--engine-ready-timeout-secs",
"42",
])
.unwrap();
let Command::Serve(args) = cli.command else {
panic!("expected serve args");
};
assert!(args.uses_external_engine());
assert_eq!(args.managed_engine.data_parallel_size_local, None);
let initial = args.to_frontend_config("tcp://127.0.0.1:29550".to_string());
assert!(matches!(
initial.transport_mode,
TransportMode::HandshakeOwner {
session_path: Some(ref path),
ready_timeout,
..
} if path == &session_path && ready_timeout == std::time::Duration::from_secs(42)
));
std::fs::write(&session_path, b"{}").unwrap();
let reattach = args.to_frontend_config("tcp://127.0.0.1:29550".to_string());
assert_eq!(
reattach.transport_mode,
TransportMode::Reattach {
path: session_path.clone(),
ready_timeout: std::time::Duration::from_secs(42),
}
);
std::fs::remove_file(session_path).unwrap();
}
#[test]
fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
let cli = Cli::try_parse_from([
@@ -1544,6 +1595,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
ready_timeout: 600s,
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode: MaybeInProc,
model: "Qwen/Qwen3-0.6B",
+13 -2
View File
@@ -113,11 +113,22 @@ async fn async_main(cli: Cli) -> Result<()> {
vllm_bench::run(bench_args).await
}
Command::Serve(args) => {
if args.engine_session.is_some()
&& args
.managed_engine
.data_parallel_size_local
.is_some_and(|local_size| local_size != 0)
{
bail!("--engine-session conflicts with non-zero --data-parallel-size-local");
}
if args.engine_session.is_some() && args.managed_engine.data_parallel_size != 1 {
bail!("--engine-session currently requires --data-parallel-size 1");
}
let handshake_port = args.managed_engine.resolve_handshake_port()?;
if args.managed_engine.data_parallel_size_local == Some(0) {
if args.uses_external_engine() {
if args.headless {
bail!("cannot combine `--headless` with `--data-parallel-size-local 0`");
bail!("cannot combine `--headless` with an external-engine frontend");
}
let handshake_address = args.managed_engine.handshake_address(handshake_port);
+2 -3
View File
@@ -5,7 +5,7 @@ edition.workspace = true
license.workspace = true
[features]
test-util = ["dep:tempfile"]
test-util = []
[dependencies]
arc-swap.workspace = true
@@ -28,7 +28,7 @@ serde_repr.workspace = true
serde_tuple.workspace = true
serde_with.workspace = true
task-local.workspace = true
tempfile = { workspace = true, optional = true }
tempfile.workspace = true
thiserror.workspace = true
thiserror-ext.workspace = true
tokio.workspace = true
@@ -42,7 +42,6 @@ anyhow.workspace = true
clap.workspace = true
expect-test.workspace = true
hex.workspace = true
tempfile.workspace = true
tracing-subscriber.workspace = true
[lints]
@@ -114,6 +114,7 @@ async fn main() -> Result<()> {
ready_timeout,
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode: None,
model_name: args.model.clone(),
@@ -64,6 +64,7 @@ async fn main() -> Result<()> {
ready_timeout: Duration::from_secs(args.ready_timeout_secs),
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode: None,
model_name: args.model.clone(),
+41 -2
View File
@@ -1,6 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
@@ -20,6 +21,7 @@ use crate::protocol::lora::LoraRequest;
use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};
use crate::session;
use crate::transport::{self, ConnectedEngine};
pub(crate) mod imp;
@@ -49,6 +51,8 @@ pub enum TransportMode {
local_input_address: Option<String>,
/// Optional explicit bind address for the output PULL socket.
local_output_address: Option<String>,
/// Optional development session file written after startup completes.
session_path: Option<PathBuf>,
},
/// The Python supervisor has already chosen the frontend transport
@@ -69,6 +73,14 @@ pub enum TransportMode {
/// Maximum time to wait for all expected engines to register.
ready_timeout: Duration,
},
/// Restore a previously handshaken frontend transport from a development session file.
Reattach {
/// Path written by a handshake-owned frontend from the same running engine.
path: PathBuf,
/// Maximum time to wait for both engine data sockets to reconnect.
ready_timeout: Duration,
},
}
/// Which coordinator implementation should be active when one is present for a
@@ -108,6 +120,7 @@ impl EngineCoreClientConfig {
ready_timeout: Duration::from_secs(30),
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode: None,
model_name: String::new(),
@@ -232,6 +245,7 @@ impl EngineCoreClient {
ready_timeout,
local_input_address,
local_output_address,
session_path,
} => {
let enable_inproc_coordinator = match config.coordinator_mode {
None => false,
@@ -241,7 +255,7 @@ impl EngineCoreClient {
}
};
transport::connect_handshake(
let connected = transport::connect_handshake(
handshake_address,
*engine_count,
advertised_host,
@@ -250,7 +264,12 @@ impl EngineCoreClient {
enable_inproc_coordinator,
*ready_timeout,
)
.await?
.await?;
if let Some(path) = session_path {
session::write(path, &connected)?;
info!(path = %path.display(), "wrote engine reattach session");
}
connected
}
TransportMode::Bootstrapped {
@@ -273,6 +292,26 @@ impl EngineCoreClient {
)
.await?
}
TransportMode::Reattach {
path,
ready_timeout,
} => {
if config.coordinator_mode.is_some() {
return Err(Error::EngineSession {
path: path.clone(),
message: "reattach sessions do not support a coordinator".to_string(),
});
}
let (input_address, output_address, engines) = session::read(path)?;
transport::connect_reattach(
path,
&input_address,
&output_address,
engines,
*ready_timeout,
)
.await?
}
};
Self::from_connected(config, connected).await
+11
View File
@@ -1,6 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
@@ -34,6 +35,16 @@ pub enum Error {
Io(#[from] std::io::Error),
#[error("transport error")]
Transport(#[from] zeromq::ZmqError),
#[error("engine session `{path}` failed: {message}")]
EngineSession { path: PathBuf, message: String },
#[error("timed out after {timeout:?} waiting to reattach engine session `{path}`: {message}")]
ReattachTimeout {
path: PathBuf,
timeout: Duration,
message: String,
},
#[error("ZMQ {socket} monitor closed while waiting for engine reattach")]
ReattachMonitorClosed { socket: &'static str },
#[error("ZMQ runtime task failed")]
ZmqRuntimeTask(#[from] tokio::task::JoinError),
#[error("engine core reported fatal failure")]
+1
View File
@@ -8,6 +8,7 @@ mod metrics;
pub mod mock_engine;
pub mod protocol;
pub mod runtime;
mod session;
#[cfg(any(test, feature = "test-util"))]
pub mod test_utils;
mod transport;
+201
View File
@@ -0,0 +1,201 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//! Persistence for development-only EngineCore reattach sessions.
//!
//! A session captures the frontend-owned ZMQ endpoints and the engine metadata
//! produced by the one-shot startup handshake. A replacement frontend can bind
//! the same endpoints and reconstruct its client without asking EngineCore to
//! repeat that handshake.
use std::fs;
use std::path::Path;
use serde::{Deserialize, Serialize};
use tempfile::NamedTempFile;
use thiserror_ext::AsReport as _;
use crate::error::{Error, Result};
use crate::protocol::handshake::EngineCoreReadyResponse;
use crate::transport::{ConnectedEngine, ConnectedTransport, EngineId};
/// Serializable transport state needed to reconnect a frontend to EngineCore.
#[derive(Debug, Serialize, Deserialize)]
struct EngineSession {
input_address: String,
output_address: String,
engines: Vec<SessionEngine>,
}
/// Serializable subset of one [`ConnectedEngine`].
#[derive(Debug, Serialize, Deserialize)]
struct SessionEngine {
engine_id: Vec<u8>,
ready_response: EngineCoreReadyResponse,
}
impl EngineSession {
/// Build a validated session snapshot from a live connected transport.
fn from_connected(path: &Path, connected: &ConnectedTransport) -> Result<Self> {
if connected.coordinator.is_some() {
return Err(Error::EngineSession {
path: path.to_path_buf(),
message: "reattach sessions do not support an in-process coordinator".to_string(),
});
}
if connected.engines.len() != 1 {
return Err(Error::EngineSession {
path: path.to_path_buf(),
message: format!(
"reattach sessions currently require exactly one engine, found {}",
connected.engines.len()
),
});
}
Ok(Self {
input_address: connected.input_address.clone(),
output_address: connected.output_address.clone(),
engines: connected
.engines
.iter()
.map(|engine| SessionEngine {
engine_id: engine.engine_id.to_vec(),
ready_response: engine.ready_response.clone(),
})
.collect(),
})
}
/// Validate a decoded session and reconstruct its transport metadata.
fn into_parts(self, path: &Path) -> Result<(String, String, Vec<ConnectedEngine>)> {
if self.engines.len() != 1 {
return Err(Error::EngineSession {
path: path.to_path_buf(),
message: format!(
"reattach sessions currently require exactly one engine, found {}",
self.engines.len()
),
});
}
Ok((
self.input_address,
self.output_address,
self.engines
.into_iter()
.map(|engine| ConnectedEngine {
engine_id: EngineId::from(engine.engine_id),
ready_response: engine.ready_response,
})
.collect(),
))
}
}
/// Atomically write a connected transport as a reattach session.
///
/// The temporary file is created in the destination directory so
/// [`NamedTempFile::persist`] can replace the session with a same-filesystem
/// rename. Dropping the temporary file cleans it up on serialization or persist
/// failure.
pub(crate) fn write(path: &Path, connected: &ConnectedTransport) -> Result<()> {
let session = EngineSession::from_connected(path, connected)?;
if path.file_name().is_none() {
return Err(Error::EngineSession {
path: path.to_path_buf(),
message: "session path must name a file".to_string(),
});
}
let directory = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
.unwrap_or_else(|| Path::new("."));
let mut temporary_file =
NamedTempFile::new_in(directory).map_err(|error| Error::EngineSession {
path: path.to_path_buf(),
message: error.to_report_string(),
})?;
serde_json::to_writer_pretty(temporary_file.as_file_mut(), &session).map_err(|error| {
Error::EngineSession {
path: path.to_path_buf(),
message: error.to_report_string(),
}
})?;
temporary_file.persist(path).map_err(|error| Error::EngineSession {
path: path.to_path_buf(),
message: error.error.to_report_string(),
})?;
Ok(())
}
/// Read and validate a reattach session from disk.
///
/// Returns the saved input/output endpoints and reconstructed engine metadata
/// used by the reattach transport.
pub(crate) fn read(path: &Path) -> Result<(String, String, Vec<ConnectedEngine>)> {
let bytes = fs::read(path).map_err(|error| Error::EngineSession {
path: path.to_path_buf(),
message: error.to_report_string(),
})?;
let session: EngineSession =
serde_json::from_slice(&bytes).map_err(|error| Error::EngineSession {
path: path.to_path_buf(),
message: error.to_report_string(),
})?;
session.into_parts(path)
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
use crate::mock_engine::default_ready_response;
use crate::test_utils::setup_bootstrapped_mock_engine;
use crate::transport::connect_reattach;
#[tokio::test]
async fn session_round_trip_restores_transport_metadata() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("engine-session.json");
let input_address = format!("ipc://{}", dir.path().join("input.sock").display());
let output_address = format!("ipc://{}", dir.path().join("output.sock").display());
let connect_path = path.clone();
let connect_input = input_address.clone();
let connect_output = output_address.clone();
let connected_task = tokio::spawn(async move {
connect_reattach(
&connect_path,
&connect_input,
&connect_output,
vec![ConnectedEngine {
engine_id: EngineId::from_engine_index(0),
ready_response: default_ready_response(),
}],
Duration::from_secs(2),
)
.await
.unwrap()
});
let (_dealer, _push) = setup_bootstrapped_mock_engine(
input_address.clone(),
output_address.clone(),
EngineId::from_engine_index(0),
)
.await;
let connected = connected_task.await.unwrap();
write(&path, &connected).unwrap();
let (actual_input, actual_output, engines) = read(&path).unwrap();
assert_eq!(actual_input, input_address);
assert_eq!(actual_output, output_address);
assert_eq!(engines.len(), 1);
assert_eq!(engines[0].engine_id.engine_index(), Some(0));
assert_eq!(
engines[0].ready_response.max_model_len,
default_ready_response().max_model_len
);
}
}
@@ -291,6 +291,7 @@ fn handshake_test_config(
ready_timeout,
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode,
model_name: model_name.to_string(),
@@ -1406,6 +1407,161 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() {
client.shutdown().await.unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn reattach_drops_stale_output_and_serves_new_request() {
init_tracing();
let ipc = IpcNamespace::new().unwrap();
let session_dir = tempfile::tempdir().unwrap();
let session_path = session_dir.path().join("engine-session.json");
let handshake_address = ipc.handshake_endpoint();
let input_address = ipc.input_endpoint();
let output_address = ipc.output_endpoint();
let engine_id = EngineId::from_engine_index(0).into_frame().to_vec();
let (old_request_tx, old_request_rx) = oneshot::channel();
let (reconnect_tx, reconnect_rx) = oneshot::channel();
let (new_request_completed_tx, new_request_completed_rx) = oneshot::channel();
let (shutdown_tx, shutdown_rx) = oneshot::channel();
let engine_task = tokio::spawn({
let engine_handshake = handshake_address.clone();
let engine_id = engine_id.clone();
let input_address = input_address.clone();
let output_address = output_address.clone();
async move {
let (_, mut initial_dealer, initial_push) =
setup_mock_engine_with_init(engine_handshake, engine_id.clone()).await;
let old_add = recv_engine_message(&mut initial_dealer).await;
assert_eq!(old_add[0].as_ref(), &[0x00]);
let old_request: EngineCoreRequest = rmp_serde::from_slice(&old_add[1]).unwrap();
assert_eq!(old_request.request_id, "req-before-reattach");
old_request_tx.send(()).unwrap();
reconnect_rx.await.unwrap();
drop((initial_dealer, initial_push));
// The pure-Rust mock sockets do not reconnect automatically after
// their bind peer disappears. Reopen only the data sockets to
// model libzmq's reconnect behavior in Python EngineCore.
let (mut dealer, mut push) =
setup_bootstrapped_mock_engine(input_address, output_address, engine_id).await;
send_outputs(
&mut push,
RequestBatchOutputs {
outputs: vec![request_output(&old_request.request_id, vec![41], None)],
..Default::default()
}
.into(),
)
.await;
let add = recv_engine_message(&mut dealer).await;
assert_eq!(add[0].as_ref(), &[0x00]);
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
assert_eq!(request.client_index, 7);
assert_eq!(request.request_id, "req-after-reattach");
send_outputs(
&mut push,
RequestBatchOutputs {
outputs: vec![request_output(
&request.request_id,
vec![42],
Some(EngineCoreFinishReason::Length),
)],
finished_requests: Some(BTreeSet::from([request.request_id])),
..Default::default()
}
.into(),
)
.await;
new_request_completed_rx.await.unwrap();
// Keep the original request live until the replacement frontend
// has completed its own request, then finish it. Both of its
// outputs are stale from the replacement's point of view.
send_outputs(
&mut push,
RequestBatchOutputs {
outputs: vec![request_output(
&old_request.request_id,
vec![43],
Some(EngineCoreFinishReason::Length),
)],
finished_requests: Some(BTreeSet::from([old_request.request_id])),
..Default::default()
}
.into(),
)
.await;
shutdown_rx.await.unwrap();
}
});
let mut initial_config = handshake_test_config(
handshake_address,
1,
"test-model",
Duration::from_secs(2),
7,
None,
);
let TransportMode::HandshakeOwner {
session_path: path, ..
} = &mut initial_config.transport_mode
else {
unreachable!("handshake_test_config returns handshake-owned transport")
};
*path = Some(session_path.clone());
let initial_client = connect_client_with_ipc(initial_config, &ipc).await;
assert_eq!(initial_client.input_address(), input_address);
assert_eq!(initial_client.output_address(), output_address);
assert!(session_path.exists());
let old_stream = initial_client
.call(sample_request_with_id("req-before-reattach"))
.await
.unwrap();
old_request_rx.await.unwrap();
initial_client.shutdown().await.unwrap();
drop(old_stream);
tokio::time::sleep(Duration::from_millis(50)).await;
let reattach_task = tokio::spawn(EngineCoreClient::connect(EngineCoreClientConfig {
transport_mode: TransportMode::Reattach {
path: session_path,
ready_timeout: Duration::from_secs(2),
},
coordinator_mode: None,
model_name: "test-model".to_string(),
client_index: 7,
}));
tokio::time::sleep(Duration::from_millis(50)).await;
reconnect_tx.send(()).unwrap();
let reattached_client = reattach_task.await.unwrap().unwrap();
assert_eq!(reattached_client.input_address(), input_address);
assert_eq!(reattached_client.output_address(), output_address);
assert_eq!(
reattached_client.engine_identities(),
vec![engine_id.as_slice()]
);
assert!(reattached_client.is_healthy());
let mut stream = reattached_client
.call(sample_request_with_id("req-after-reattach"))
.await
.unwrap();
let output = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap();
assert_eq!(output.new_token_ids, vec![42]);
assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Length));
new_request_completed_tx.send(()).unwrap();
let _ = shutdown_tx.send(());
engine_task.await.unwrap();
reattached_client.shutdown().await.unwrap();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn call_utility_failure_message_surfaces_as_error() {
init_tracing();
+203 -9
View File
@@ -8,13 +8,14 @@ use std::time::Duration;
use bytes::Bytes;
use enum_as_inner::EnumAsInner;
use futures::StreamExt;
use thiserror_ext::AsReport;
use tokio::sync::mpsc;
use tokio::time::timeout;
use tracing::{debug, error, info, trace, warn};
use zeromq::prelude::{Socket, SocketRecv, SocketSend};
use zeromq::util::PeerIdentity;
use zeromq::{PullSocket, RouterSendHalf, RouterSocket, ZmqError, ZmqMessage};
use zeromq::{PullSocket, RouterSendHalf, RouterSocket, SocketEvent, ZmqError, ZmqMessage};
use crate::coordinator::CoordinatorBootstrap;
use crate::error::{Error, Result, bail_unexpected_handshake_message};
@@ -154,17 +155,14 @@ pub async fn connect_handshake(
info!(
engine_count,
handshake_address, "waiting for engines to connect"
handshake_address,
?ready_timeout,
"waiting for engines to connect"
);
// 1. Bind shared local input/output sockets first so every engine receives the same data-plane
// addresses during handshake.
debug!(
local_host,
?ready_timeout,
engine_count,
"binding shared transport sockets"
);
debug!(local_host, engine_count, "binding shared transport sockets");
let (input_address, mut input_socket, output_address, output_socket) =
bind_local_sockets(local_host, local_input_address, local_output_address).await?;
info!(%input_address, %output_address, "bound local transport sockets");
@@ -367,6 +365,109 @@ pub async fn connect_bootstrapped(
})
}
/// Bind a saved frontend transport and wait for both engine data sockets to reconnect.
pub async fn connect_reattach(
session_path: &std::path::Path,
input_address: &str,
output_address: &str,
engines: Vec<ConnectedEngine>,
ready_timeout: Duration,
) -> Result<ConnectedTransport> {
let mut input_socket = RouterSocket::new();
let mut output_socket = PullSocket::new();
// Subscribe before binding either endpoint so a fast reconnect cannot race
// with monitor setup.
let mut input_events = input_socket.monitor();
let mut output_events = output_socket.monitor();
let input_address = input_socket.bind(input_address).await?.to_string();
let output_address = output_socket.bind(output_address).await?.to_string();
let expected_engine = engines.first().ok_or_else(|| Error::EngineSession {
path: session_path.to_path_buf(),
message: "reattach session does not contain an engine".to_string(),
})?;
info!(
path = %session_path.display(),
%input_address,
%output_address,
?ready_timeout,
"waiting for engine transport to reattach"
);
timeout(ready_timeout, async {
tokio::try_join!(
wait_for_accepted_peer(
&mut input_events,
"input ROUTER",
Some(&expected_engine.engine_id),
),
wait_for_accepted_peer(&mut output_events, "output PULL", None),
)
})
.await
.map_err(|_| Error::ReattachTimeout {
path: session_path.to_path_buf(),
timeout: ready_timeout,
message: "ZMQ monitors did not observe both engine data connections".to_string(),
})??;
info!(
%input_address,
%output_address,
"reattached engine transport"
);
// zeromq has no monitor-unregister API. Dropping the receivers closes its
// bounded channels; the sockets retain only a closed sender whose later
// `try_send` calls fail immediately and are ignored.
drop((input_events, output_events));
let (input_send, _) = input_socket.split();
Ok(ConnectedTransport {
input_address,
output_address,
engines,
coordinator: None,
input_send,
output_socket,
})
}
/// Wait until a bound socket accepts a fully handshaken peer.
async fn wait_for_accepted_peer(
events: &mut futures::channel::mpsc::Receiver<SocketEvent>,
socket: &'static str,
expected_engine: Option<&EngineId>,
) -> Result<()> {
while let Some(event) = events.next().await {
match event {
SocketEvent::Accepted(_, peer_id) => {
if let Some(expected_engine) = expected_engine
&& peer_id.as_ref() != &expected_engine[..]
{
return Err(Error::UnexpectedHandshakeIdentity {
expected: expected_engine.to_vec(),
actual: peer_id.as_ref().to_vec(),
});
}
return Ok(());
}
SocketEvent::AcceptFailed(error) => {
debug!(
%socket,
error = %error.as_report(),
"ignored rejected peer while waiting for engine reattach"
);
}
_ => {}
}
}
Err(Error::ReattachMonitorClosed { socket })
}
/// Bind new input and output sockets.
async fn bind_local_sockets(
local_host: &str,
@@ -585,7 +686,10 @@ pub async fn run_output_loop(
#[cfg(test)]
mod tests {
use super::bind_local_sockets;
use super::*;
use crate::mock_engine::default_ready_response;
use crate::test_utils::IpcNamespace;
use zeromq::{DealerSocket, PushSocket, SocketOptions};
#[tokio::test]
async fn bind_local_sockets_resolves_zero_port_bindings() {
@@ -596,4 +700,94 @@ mod tests {
assert!(output_address.starts_with("tcp://127.0.0.1:"));
assert_ne!(input_address, output_address);
}
#[tokio::test]
async fn reattach_returns_after_both_data_sockets_are_connected() {
let ipc = IpcNamespace::new().unwrap();
let input_address = ipc.input_endpoint();
let output_address = ipc.output_endpoint();
let engine_id = EngineId::from_engine_index(0);
let session_path = std::path::PathBuf::from("test-engine-session.json");
let connect_input = input_address.clone();
let connect_output = output_address.clone();
let connect_engine = engine_id.clone();
let mut reattach_task = tokio::spawn(async move {
connect_reattach(
&session_path,
&connect_input,
&connect_output,
vec![ConnectedEngine {
engine_id: connect_engine,
ready_response: default_ready_response(),
}],
Duration::from_secs(5),
)
.await
.unwrap()
});
tokio::task::yield_now().await;
for endpoint in [&input_address, &output_address] {
let socket_path = endpoint.strip_prefix("ipc://").unwrap();
timeout(Duration::from_secs(1), async {
loop {
match tokio::net::UnixStream::connect(socket_path).await {
Ok(stream) => {
drop(stream);
break;
}
Err(_) => tokio::task::yield_now().await,
}
}
})
.await
.expect("reattach endpoint was not bound");
}
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(
!reattach_task.is_finished(),
"a rejected non-ZMTP peer must not fail reattach"
);
let mut options = SocketOptions::default();
options.peer_identity(PeerIdentity::try_from(engine_id.clone()).unwrap());
let mut dealer = DealerSocket::with_options(options);
dealer.connect(&input_address).await.unwrap();
assert!(
!reattach_task.is_finished(),
"reattach must also wait for the output PUSH connection"
);
let mut push = PushSocket::new();
push.connect(&output_address).await.unwrap();
let mut connected = timeout(Duration::from_secs(2), &mut reattach_task)
.await
.expect("reattach timed out")
.unwrap();
send_message(
&mut connected.input_send,
&engine_id,
Bytes::from_static(b"request-type"),
b"request-payload".to_vec(),
)
.await
.unwrap();
let input = timeout(Duration::from_secs(1), dealer.recv())
.await
.unwrap()
.unwrap()
.into_vec();
assert_eq!(input[0].as_ref(), b"request-type");
assert_eq!(input[1].as_ref(), b"request-payload");
push.send(ZmqMessage::from(b"output-payload".to_vec())).await.unwrap();
let output = timeout(Duration::from_secs(1), connected.output_socket.recv())
.await
.unwrap()
.unwrap()
.into_vec();
assert_eq!(output[0].as_ref(), b"output-payload");
}
}
@@ -115,6 +115,7 @@ async fn main() -> Result<()> {
ready_timeout,
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode: None,
model_name: args.model.clone(),
+1
View File
@@ -32,6 +32,7 @@ fn client_config(handshake_address: String, engine_count: usize) -> EngineCoreCl
ready_timeout: Duration::from_secs(5),
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode: None,
model_name: "mock-model".to_string(),
@@ -56,6 +56,7 @@ async fn main() -> Result<()> {
ready_timeout: Duration::from_secs(args.ready_timeout_secs),
local_input_address: None,
local_output_address: None,
session_path: None,
},
coordinator_mode: CoordinatorMode::MaybeInProc,
model: args.model,
+1
View File
@@ -242,6 +242,7 @@ impl Config {
match &self.transport_mode {
TransportMode::HandshakeOwner { engine_count, .. }
| TransportMode::Bootstrapped { engine_count, .. } => *engine_count,
TransportMode::Reattach { .. } => 1,
}
}