forked from Karylab-cklius/vllm
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
58feec155a | ||
|
|
93dec05d69 | ||
|
|
707fb0de8f | ||
|
|
fe3b526be4 |
+24
-1
@@ -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
@@ -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(),
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")]
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -242,6 +242,7 @@ impl Config {
|
||||
match &self.transport_mode {
|
||||
TransportMode::HandshakeOwner { engine_count, .. }
|
||||
| TransportMode::Bootstrapped { engine_count, .. } => *engine_count,
|
||||
TransportMode::Reattach { .. } => 1,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user