diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 49910d80c7..a9bcc01959 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -657,12 +657,9 @@ impl EngineCoreClient { ); let result: Result<()> = async { + // The engine that receives it wakes the other ranks if they are idle. if let Some(coordinator) = self.coordinator.as_ref() { - let snapshot = coordinator.snapshot(); - req.current_wave = snapshot.current_wave; - if !snapshot.engines_running { - coordinator.notify_first_request(engine_id.clone())?; - } + req.current_wave = coordinator.snapshot().current_wave; } debug!( diff --git a/rust/src/engine-core-client/src/coordinator/external.rs b/rust/src/engine-core-client/src/coordinator/external.rs index 606224ea9e..d0418352a6 100644 --- a/rust/src/engine-core-client/src/coordinator/external.rs +++ b/rust/src/engine-core-client/src/coordinator/external.rs @@ -3,34 +3,16 @@ use std::sync::Arc; -use serde_tuple::{Deserialize_tuple, Serialize_tuple}; +use serde_tuple::Deserialize_tuple; use thiserror_ext::AsReport; -use tokio::sync::mpsc; use tracing::{debug, warn}; -use zeromq::prelude::{SocketRecv, SocketSend}; +use zeromq::prelude::SocketRecv; use zeromq::{XSubSocket, ZmqMessage}; use crate::client::imp::ClientInner; -use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState}; +use crate::coordinator::handle::CoordinatorState; use crate::error::{Error, Result, bail_unexpected_coordinator_output}; -use crate::protocol::{OpaqueValue, decode_msgpack, encode_msgpack}; - -/// Frontend-to-coordinator wakeup message sent when the first request arrives -/// while all engines are paused. -/// -/// This matches the frontend-side msgpack tuple sent by Python -/// `DPAsyncMPClient._ensure_stats_update_task` to the coordinator front socket. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)] -struct CoordinatorWakeupMessage { - /// Engine index that already has the triggering request and should be - /// excluded from the coordinator's `START_DP_WAVE` rebroadcast. - exclude_engine_index: u32, - /// DP wave number observed by the frontend when the request was admitted. - wave: u32, -} +use crate::protocol::{OpaqueValue, decode_msgpack}; /// Coordinator-to-frontend state publish received on the front-side coordinator /// socket. @@ -56,56 +38,17 @@ struct CoordinatorStateUpdate { /// Background half of an external Python-owned coordinator connection. /// -/// This owns the command receiver and one frontend-facing XSUB socket. It -/// mirrors the subset of Python's coordinator protocol needed by the Rust -/// bootstrapped frontend: receive `(counts, wave, running)` publishes, ignore -/// `counts`, and send `(exclude_engine_index, wave)` wakeup messages when the -/// first request arrives while engines are paused. +/// This owns one frontend-facing XSUB socket. It mirrors the subset of Python's +/// coordinator protocol needed by the Rust bootstrapped frontend: receive +/// `(counts, wave, running)` publishes and ignore `counts`. pub(crate) struct ExternalCoordinatorService { state: Arc, - command_rx: mpsc::UnboundedReceiver, socket: XSubSocket, } impl ExternalCoordinatorService { - pub(super) fn new( - state: Arc, - command_rx: mpsc::UnboundedReceiver, - socket: XSubSocket, - ) -> Self { - Self { - state, - command_rx, - socket, - } - } - - /// Apply one frontend-originated command to the external coordinator state - /// machine. - async fn handle_command(&mut self, command: CoordinatorCommand) -> Result<()> { - match command { - CoordinatorCommand::FirstRequest { - target_engine_id, - wave, - } => { - let target_engine_index = target_engine_id.engine_index().ok_or_else(|| { - Error::UnsupportedCoordinatorEngineId { - engine_id: target_engine_id.to_vec(), - } - })?; - debug!( - wave, - exclude_engine_index = target_engine_index, - "notifying external coordinator about first request while engines were paused" - ); - let payload = encode_msgpack(&CoordinatorWakeupMessage { - exclude_engine_index: target_engine_index, - wave, - })?; - self.socket.send(ZmqMessage::from(payload)).await?; - } - } - Ok(()) + pub(super) fn new(state: Arc, socket: XSubSocket) -> Self { + Self { state, socket } } /// Apply one publish received from the xsub socket containing a coordinator @@ -136,26 +79,14 @@ impl ExternalCoordinatorService { Ok(()) } - /// Drive the coordinator event loop until either side of the control plane - /// is closed or a fatal error is observed. + /// Drive the coordinator event loop until the task is aborted or a fatal + /// error is observed. pub(crate) async fn run(mut self, inner: Arc) { let result: Result<()> = async { loop { - tokio::select! { - // Received frontend-originated command from the handle. - command = self.command_rx.recv() => { - let Some(command) = command else { - warn!("external coordinator command channel closed, shutting down service"); - return Ok(()); - }; - self.handle_command(command).await?; - } - // Received publish from the external coordinator socket. - publish = self.socket.recv() => { - let publish = publish.map_err(Error::from)?; - self.handle_publish(publish).await?; - } - } + // Received publish from the external coordinator socket. + let publish = self.socket.recv().await.map_err(Error::from)?; + self.handle_publish(publish).await?; } } .await; diff --git a/rust/src/engine-core-client/src/coordinator/handle.rs b/rust/src/engine-core-client/src/coordinator/handle.rs index b68c76ddec..9ebc1717e3 100644 --- a/rust/src/engine-core-client/src/coordinator/handle.rs +++ b/rust/src/engine-core-client/src/coordinator/handle.rs @@ -4,101 +4,56 @@ use std::sync::Arc; use parking_lot::Mutex; -use tokio::sync::mpsc; use zeromq::prelude::Socket; use zeromq::{XPubSocket, XSubSocket}; use crate::coordinator::external::ExternalCoordinatorService; use crate::coordinator::inproc::InProcCoordinatorRunner; -use crate::error::{Error, Result, bail_control_closed}; -use crate::transport::EngineId; +use crate::error::Result; /// Snapshot to the coordinator state for request routing and stamping. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct CoordinatorStateSnapshot { /// The current DP wave, which will be stamped on outgoing requests. pub current_wave: u32, - /// Whether the engines are currently running or paused, which determines if - /// the frontend must trigger a new wave on the next request. + /// Whether the engines are currently running or paused, as last reported + /// by the engines themselves. pub engines_running: bool, } -impl CoordinatorStateSnapshot { - /// Resume the engines for a `FirstRequest` and return the wave to broadcast - /// and the engine to exclude from the wakeup. - /// - /// The request may have been stamped with a `request_wave` older than - /// `current_wave` if a `WaveComplete` advanced it after the command was - /// enqueued. Such a request still needs serving, so the current wave is - /// broadcast to every engine (`exclude = None`); the wave is never rewound. - /// A non-stale request excludes the engine that already received it. Mirrors - /// the Python coordinator's front-end path. - pub(crate) fn start_wave_for_first_request( - &mut self, - request_wave: u32, - target_engine_index: u32, - ) -> (u32, Option) { - self.engines_running = true; - let exclude = (request_wave >= self.current_wave).then_some(target_engine_index); - (self.current_wave, exclude) - } -} - /// Shared in-process coordinator state. pub(crate) type CoordinatorState = Mutex; -/// Commands sent from the frontend request path into the background runner. -#[derive(Debug)] -pub(crate) enum CoordinatorCommand { - /// The first request arrived while all engines were paused. - /// - /// The coordinator should broadcast `START_DP_WAVE` with the current wave - /// and the target engine index as the excluded engine. - FirstRequest { - target_engine_id: EngineId, - wave: u32, - }, -} - /// Frontend-facing coordinator handle used by `EngineCoreClient::call()`. /// -/// This side stays intentionally small: it can read the latest wave snapshot -/// and enqueue a `FirstRequest` transition when the request path observes the -/// system in the paused state. +/// The frontend only reads the latest wave snapshot: waves are started by the +/// engines, which notify the coordinator when an idle engine receives a request. #[derive(Clone)] pub(crate) struct CoordinatorHandle { state: Arc, - command_tx: mpsc::UnboundedSender, } impl CoordinatorHandle { - fn new_parts() -> ( - Self, - Arc, - mpsc::UnboundedReceiver, - ) { + fn new_parts() -> (Self, Arc) { let state = Arc::new(Mutex::new(CoordinatorStateSnapshot { current_wave: 0, engines_running: false, })); - let (command_tx, command_rx) = mpsc::unbounded_channel(); ( Self { state: state.clone(), - command_tx, }, state, - command_rx, ) } /// Build the paired frontend handle and background runner around one /// engine-facing coordinator broadcast socket. pub(crate) fn new_inproc(coordinator_input: XPubSocket) -> (Self, InProcCoordinatorRunner) { - let (handle, state, command_rx) = Self::new_parts(); + let (handle, state) = Self::new_parts(); ( handle, - InProcCoordinatorRunner::new(state, command_rx, coordinator_input), + InProcCoordinatorRunner::new(state, coordinator_input), ) } @@ -107,41 +62,15 @@ impl CoordinatorHandle { pub(crate) async fn connect_external( coordinator_address: &str, ) -> Result<(Self, ExternalCoordinatorService)> { - let (handle, state, command_rx) = Self::new_parts(); + let (handle, state) = Self::new_parts(); let mut socket = XSubSocket::new(); socket.connect(coordinator_address).await?; socket.subscribe("").await?; - Ok(( - handle, - ExternalCoordinatorService::new(state, command_rx, socket), - )) + Ok((handle, ExternalCoordinatorService::new(state, socket))) } /// Snapshot the coordinator state for request routing and stamping. pub(crate) fn snapshot(&self) -> CoordinatorStateSnapshot { *self.state.lock() } - - /// Notify the runner that a new request arrived while engines were paused. - /// - /// The handle flips `engines_running` optimistically so concurrent request - /// submissions coalesce behind one `START_DP_WAVE` broadcast instead of all - /// trying to trigger the wave independently. - pub(crate) fn notify_first_request(&self, target_engine_id: EngineId) -> Result<()> { - let mut state = self.state.lock(); - if state.engines_running { - return Ok(()); - } - - let command = CoordinatorCommand::FirstRequest { - target_engine_id, - wave: state.current_wave, - }; - if self.command_tx.send(command).is_err() { - bail_control_closed!("in-process coordinator command channel already shut down"); - } - - state.engines_running = true; - Ok(()) - } } diff --git a/rust/src/engine-core-client/src/coordinator/inproc.rs b/rust/src/engine-core-client/src/coordinator/inproc.rs index cbd69095f8..d780b64e99 100644 --- a/rust/src/engine-core-client/src/coordinator/inproc.rs +++ b/rust/src/engine-core-client/src/coordinator/inproc.rs @@ -12,7 +12,7 @@ use zeromq::{XPubSocket, ZmqError, ZmqMessage}; use crate::client::imp::ClientInner; use crate::coordinator::bootstrap::is_engine_subscription; -use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState}; +use crate::coordinator::handle::CoordinatorState; use crate::error::{Error, Result, bail_unexpected_coordinator_output}; use crate::protocol::encode_msgpack; use crate::protocol::output::{DpControlMessage, DpControlOutput, EngineCoreOutputs}; @@ -30,32 +30,25 @@ use crate::protocol::request::EngineCoreRequestType; struct StartDpWaveMessage { /// DP wave number that all engines should start processing. wave: u32, - /// Engine index that already received the triggering request and so does not - /// need an extra wakeup. `None` wakes every engine (used when the triggering - /// request was for a stale wave). + /// Engine index that started the wave and so does not need an extra wakeup. + /// `None` wakes every engine. exclude_engine_index: Option, } /// Background half of the in-process coordinator. /// -/// This owns the command receiver and the engine-facing coordinator input -/// socket. It is the single place where wave transitions are serialized and -/// where `START_DP_WAVE` broadcasts are emitted. +/// This owns the engine-facing coordinator input socket. It is the single place +/// where wave transitions are serialized and where `START_DP_WAVE` broadcasts +/// are emitted. pub(crate) struct InProcCoordinatorRunner { state: Arc, - command_rx: mpsc::UnboundedReceiver, coordinator_input: XPubSocket, } impl InProcCoordinatorRunner { - pub(super) fn new( - state: Arc, - command_rx: mpsc::UnboundedReceiver, - coordinator_input: XPubSocket, - ) -> Self { + pub(super) fn new(state: Arc, coordinator_input: XPubSocket) -> Self { Self { state, - command_rx, coordinator_input, } } @@ -82,34 +75,6 @@ impl InProcCoordinatorRunner { Ok(()) } - /// Apply one frontend-originated command to the coordinator state machine. - async fn handle_command(&mut self, command: CoordinatorCommand) -> Result<()> { - match command { - CoordinatorCommand::FirstRequest { - target_engine_id, - wave, - } => { - let target_engine_index = target_engine_id.engine_index().ok_or_else(|| { - Error::UnsupportedCoordinatorEngineId { - engine_id: target_engine_id.to_vec(), - } - })?; - let (current_wave, exclude) = { - let mut state = self.state.lock(); - state.start_wave_for_first_request(wave, target_engine_index) - }; - debug!( - current_wave, - request_wave = wave, - ?exclude, - "starting DP wave after first request while engines were paused" - ); - self.broadcast_start_wave(current_wave, exclude).await?; - } - } - Ok(()) - } - /// Apply one engine-originated control output to the coordinator state /// machine. async fn handle_outputs(&mut self, outputs: EngineCoreOutputs) -> Result<()> { @@ -160,7 +125,7 @@ impl InProcCoordinatorRunner { debug!( wave, exclude_engine_index = engine_index, - "starting DP wave after stale-wave notification from engine" + "starting DP wave on notification from engine" ); self.broadcast_start_wave(wave, Some(engine_index)).await?; } @@ -175,8 +140,8 @@ impl InProcCoordinatorRunner { Ok(()) } - /// Drive the coordinator event loop until either side of the control plane - /// is closed or a fatal error is observed. + /// Drive the coordinator event loop until the engine control output is + /// closed or a fatal error is observed. /// /// Any fatal error closes the main client registries so request streams and /// future calls observe a stable shutdown cause. @@ -216,14 +181,6 @@ impl InProcCoordinatorRunner { Err(error) => return Err(error.into()), } } - // Received frontend-originated command from the handle. - command = self.command_rx.recv() => { - let Some(command) = command else { - warn!("coordinator command channel closed, shutting down coordinator runner"); - return Ok(()); - }; - self.handle_command(command).await?; - } // Received engine-originated control output from the coordinator socket. outputs = output_rx.recv() => { let Some(outputs) = outputs else { @@ -242,48 +199,3 @@ impl InProcCoordinatorRunner { inner.close_registries(Arc::new(error)); } } - -#[cfg(test)] -mod tests { - use crate::coordinator::handle::CoordinatorStateSnapshot; - - /// A `FirstRequest` for the current wave starts that wave and excludes the - /// engine that already received the triggering request. - #[test] - fn first_request_for_current_wave_excludes_target() { - let mut state = CoordinatorStateSnapshot { - current_wave: 3, - engines_running: false, - }; - - let (wave, exclude) = state.start_wave_for_first_request(3, 2); - - assert_eq!(wave, 3); - assert_eq!(exclude, Some(2)); - assert!(state.engines_running); - assert_eq!(state.current_wave, 3); - } - - /// A `FirstRequest` whose wave was superseded by a racing `WaveComplete` - /// (`request_wave < current_wave`) must still start the request's wave: it - /// broadcasts the current wave and wakes every engine (`exclude = None`) - /// rather than rewinding the wave or dropping the request. - #[test] - fn stale_first_request_starts_current_wave_for_all_engines() { - let mut state = CoordinatorStateSnapshot { - current_wave: 4, - engines_running: false, - }; - - // Request stamped with wave 3 while the coordinator already advanced to 4. - let (wave, exclude) = state.start_wave_for_first_request(3, 2); - - assert_eq!( - wave, 4, - "must broadcast the current wave, not the stale one" - ); - assert_eq!(exclude, None, "a stale request must wake every engine"); - assert!(state.engines_running); - assert_eq!(state.current_wave, 4, "wave must not be rewound"); - } -} diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index d55f2cbbcb..6217208e3b 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/rust/src/engine-core-client/src/error.rs @@ -55,8 +55,6 @@ pub enum Error { UnexpectedCoordinatorOutput { message: String }, #[error("unexpected output on main dispatcher path: {message}")] UnexpectedDispatcherOutput { message: String }, - #[error("coordinator requires a Python-compatible two-byte engine id, got {engine_id:?}")] - UnsupportedCoordinatorEngineId { engine_id: Vec }, #[error("unsupported auxiliary frame(s): expected 1 frame, got {frame_count}")] UnsupportedAuxFrames { frame_count: usize }, #[error("external coordinator mode is not implemented yet")] @@ -66,8 +64,6 @@ pub enum Error { context: &'static str, field: &'static str, }, - #[error("engine control channel closed unexpectedly: {message}")] - ControlClosed { message: String }, #[error("request `{request_id}` is already in flight")] DuplicateRequestId { request_id: String }, #[error( diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 7b7a880bde..c2216168b8 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -267,6 +267,19 @@ async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, Option) { rmp_serde::from_slice(&frames[1]).expect("decode START_DP_WAVE payload") } +async fn send_start_wave(push: &mut PushSocket, engine_index: u32, wave: u32) { + send_outputs( + push, + DpControlOutput { + engine_index, + timestamp: 0.0, + control: DpControlMessage::StartWave(wave), + } + .into(), + ) + .await; +} + async fn connect_client_with_ipc( config: EngineCoreClientConfig, ipc: &IpcNamespace, @@ -413,12 +426,6 @@ async fn recv_xpub_subscription(xpub: &mut XPubSocket) { assert_eq!(frames[0].as_ref(), b"\x01"); } -async fn recv_external_coordinator_wakeup(xpub: &mut XPubSocket) -> (u32, u32) { - let frames = recv_xpub_message(xpub).await; - assert_eq!(frames.len(), 1); - rmp_serde::from_slice(&frames[0]).expect("decode external coordinator wakeup") -} - async fn send_external_coordinator_publish( xpub: &mut XPubSocket, payload: &T, @@ -617,6 +624,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { engine.coordinator.take().expect("coordinator sockets should be present"); let data_socket = engine.data_sockets.first_mut().expect("data socket"); + // An idle engine handed a request starts the wave itself. + let add = recv_engine_message(&mut data_socket.dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + assert_eq!(request.request_id, "req-1"); + assert_eq!(request.current_wave, 0); + send_start_wave(&mut coordinator.output_push, 0, 0).await; + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; assert_eq!((wave, exclude_engine), (0, Some(0))); @@ -629,12 +644,6 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let replay = recv_start_dp_wave(&mut coordinator.input_sub).await; assert_eq!(replay, (0, None)); - let add = recv_engine_message(&mut data_socket.dealer).await; - assert_eq!(add[0].as_ref(), &[0x00]); - let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); - assert_eq!(request.request_id, "req-1"); - assert_eq!(request.current_wave, 0); - assert!( timeout( Duration::from_millis(200), @@ -670,14 +679,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { ) .await; - let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (1, Some(0))); - let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); assert_eq!(request.request_id, "req-3"); assert_eq!(request.current_wave, 1); + send_start_wave(&mut coordinator.output_push, 0, 1).await; + + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; + assert_eq!((wave, exclude_engine), (1, Some(0))); send_outputs( &mut data_socket.push, @@ -896,9 +906,6 @@ async fn coordinator_accepts_stats_only_outputs() { engine.coordinator.take().expect("coordinator sockets should be present"); let data_socket = engine.data_sockets.first_mut().expect("data socket"); - let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (0, Some(0))); - send_outputs( &mut coordinator.output_push, RequestBatchOutputs { @@ -3248,7 +3255,7 @@ async fn bootstrapped_external_coordinator_connects_and_subscribes() { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends_one_wakeup() { +async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends_no_wakeup() { init_tracing(); let ipc = IpcNamespace::new().unwrap(); let input_address = ipc.input_endpoint(); @@ -3289,18 +3296,11 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); - let wakeup = timeout( - Duration::from_secs(1), - recv_external_coordinator_wakeup(&mut stats_socket), - ) - .await - .unwrap(); - assert_eq!(wakeup, (0, 7)); - + // The engine that receives the request starts the wave, not the frontend. assert!( timeout( Duration::from_millis(200), - recv_external_coordinator_wakeup(&mut stats_socket) + recv_xpub_message(&mut stats_socket) ) .await .is_err() @@ -3333,78 +3333,3 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends client.shutdown().await.unwrap(); } - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() { - init_tracing(); - let ipc = IpcNamespace::new().unwrap(); - let input_address = ipc.input_endpoint(); - let output_address = ipc.output_endpoint(); - let coordinator_address = ipc.endpoint("stats.sock"); - - let mut stats_socket = XPubSocket::new(); - stats_socket.bind(&coordinator_address).await.unwrap(); - - let client_task = tokio::spawn({ - let input_address = input_address.clone(); - let output_address = output_address.clone(); - let coordinator_address = coordinator_address.clone(); - async move { - EngineCoreClient::connect(bootstrapped_test_config( - input_address, - output_address, - 1, - Duration::from_secs(2), - 0, - Some(CoordinatorMode::External { - address: coordinator_address, - }), - )) - .await - .unwrap() - } - }); - - let (mut dealer, mut push) = - setup_bootstrapped_mock_engine(input_address, output_address, &[0x00, 0x00]).await; - let client = client_task.await.unwrap(); - recv_xpub_subscription(&mut stats_socket).await; - - send_external_coordinator_publish(&mut stats_socket, &(Value::Nil, 5_u32, true)).await; - tokio::time::sleep(Duration::from_millis(50)).await; - - let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); - - assert!( - timeout( - Duration::from_millis(200), - recv_external_coordinator_wakeup(&mut stats_socket) - ) - .await - .is_err() - ); - - let add = recv_engine_message(&mut dealer).await; - let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); - assert_eq!(request.current_wave, 5); - - send_outputs( - &mut push, - RequestBatchOutputs { - outputs: vec![request_output( - "req-1", - vec![], - Some(EngineCoreFinishReason::Length), - )], - finished_requests: Some(BTreeSet::from(["req-1".to_string()])), - ..Default::default() - } - .into(), - ) - .await; - - let final_output = timeout(Duration::from_secs(1), stream.next()).await.unwrap(); - assert!(final_output.is_some()); - - client.shutdown().await.unwrap(); -} diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index 7b350fc55e..d5beea2f63 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -18,6 +18,7 @@ from vllm.inputs import PromptType from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.sampling_params import RequestOutputKind +from vllm.v1.engine import EngineCoreRequestType from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.core_client import DPLBAsyncMPClient from vllm.v1.metrics.loggers import StatLoggerBase @@ -336,12 +337,12 @@ async def _poll_flag(engine: AsyncLLM, want: bool, timeout: float) -> bool: @pytest.mark.asyncio async def test_dp_pause_late_request_does_not_block_drain(): - """A request arriving after pause must not leave the coordinator believing + """A request arriving after pause must not make the coordinator believe the engines are running. - Paused engines discard START_DP_WAVE, so nothing can report wave - completion afterwards; if forwarding the wake also marks the engines as - running, drain has no way back and can only end in a timeout. + Paused engines start no wave for it, so nothing could report wave + completion afterwards; a running flag set on its behalf would leave drain + no way back but a timeout. MoE only: wave coordination is enabled iff the model is MoE, so a dense model never reaches the coordinator state this exercises. @@ -365,11 +366,9 @@ async def test_dp_pause_late_request_does_not_block_drain(): # A design that never reports the engines as running would satisfy # every drain assertion below while destroying the signal, so pin it - # down first. The front-end sets its own copy optimistically when it - # forwards the wake, so sample only after several coordinator - # publishes (every 100ms while stats change) have overwritten it. - # Rank 0 reports the wave only after its first step, which can outlast - # the 2s on a cold device, so wait for the report rather than sample it. + # down first. Rank 0 reports the wave only after its first step, which + # can outlast the 2s on a cold device, so wait for the report rather + # than sample it. await asyncio.sleep(2) assert not long_request.done(), "the warmup request was too short to sample" assert await _poll_flag(engine, True, timeout=30), ( @@ -382,24 +381,12 @@ async def test_dp_pause_late_request_does_not_block_drain(): await engine.pause_generation(mode="abort") await engine.wait_for_requests_to_drain(drain_timeout=30) - # Awaiting add_request guarantees the new-request notification has been - # sent to the coordinator - the message that used to latch the flag. collector = await engine.add_request( request_id="late", prompt=DP_PAUSE_PROMPT, params=SamplingParams(max_tokens=5), ) - - # The front-end marks the engines running off the back of that - # notification. This is what makes the test non-vacuous: it is the - # path that used to leave the coordinator stuck. - assert await _poll_flag(engine, True, timeout=5), ( - "the late request did not notify the coordinator" - ) - - # It must settle back by itself. Unfixed it never does, because the - # paused engines discard the wake and so never report wave completion. - assert await _poll_flag(engine, False, timeout=60) + assert not await _poll_flag(engine, True, timeout=5) # The late request was held rather than dropped: it completes on resume. await engine.resume_generation() @@ -409,6 +396,46 @@ async def test_dp_pause_late_request_does_not_block_drain(): break +@pytest.mark.asyncio +async def test_dp_request_reaching_rank_after_resume_wakes_peers(): + """A request can reach its rank only after resume (e.g. from another API + server), with no other work anywhere. That rank must wake its peers itself, + or it steps alone. Modelled here by holding the request's ADD on the client + side across pause and resume.""" + with ExitStack() as after: + engine = AsyncLLM.from_engine_args( + _get_dp_pause_engine_args(expert_parallel=True) + ) + after.callback(engine.shutdown) + client = engine.engine_core + params = SamplingParams(max_tokens=5) + + await _consume(engine.generate(DP_PAUSE_PROMPT, params, "warmup")) + await engine.pause_generation(mode="abort") + assert await _poll_flag(engine, False, timeout=30) + + held: list[tuple] = [] + send_input = client._send_input + + def hold_add(request_type, request, engine_id=None): + if request_type == EngineCoreRequestType.ADD: + held.append((request_type, request, engine_id)) + return asyncio.sleep(0) + return send_input(request_type, request, engine_id) + + client._send_input = hold_add + task = asyncio.create_task( + _consume(engine.generate(DP_PAUSE_PROMPT, params, "late")) + ) + await asyncio.sleep(2) + assert held + + await engine.resume_generation() + client._send_input = send_input + await send_input(*held[0]) + await asyncio.wait_for(task, timeout=60) + + @pytest.mark.asyncio async def test_dp_sleep_late_request_does_not_block_drain(): """The same latch, reached through sleep rather than pause. @@ -441,14 +468,7 @@ async def test_dp_sleep_late_request_does_not_block_drain(): prompt=DP_PAUSE_PROMPT, params=SamplingParams(max_tokens=5), ) - - assert await _poll_flag(engine, True, timeout=5), ( - "the request did not notify the coordinator" - ) - - # Sleeping engines cannot report wave completion, so a coordinator that - # marked them running when it forwarded the wake never hears otherwise. - assert await _poll_flag(engine, False, timeout=60) + assert not await _poll_flag(engine, True, timeout=5) await engine.wake_up() assert not await engine.is_sleeping() @@ -628,7 +648,7 @@ async def test_dp_pause_keep_race_staggered_engines(): @pytest.mark.asyncio async def test_dp_pause_barrier_request_deadlock(): - """Test that start_dp_wave is ignored while paused. + """Test that a request sent while paused does not wake any engine. Sequence: 1. Pause all engines (PAUSED_ALL). @@ -697,7 +717,7 @@ async def test_dp_pause_barrier_request_deadlock(): # Yield so generate() preprocessing completes and # add_request_async is called (which, in buggy code, - # would send FIRST_REQ and wake engine 1). + # would wake engine 1). for _ in range(200): await asyncio.sleep(0) @@ -725,7 +745,7 @@ async def test_dp_pause_barrier_request_deadlock(): for t in mid_barrier_tasks: t.cancel() pytest.fail( - "Staggered barrier deadlocked — FIRST_REQ sent while " + "Staggered barrier deadlocked — a request sent while " "paused caused collective-op mismatch between engines" ) diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index 2d9f7e8632..df65b9f481 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -25,7 +25,7 @@ from vllm.engine.arg_utils import EngineArgs from vllm.platforms import current_platform from vllm.utils.torch_utils import set_default_torch_num_threads from vllm.v1.core.sched.interface import PauseState -from vllm.v1.engine import EngineCoreRequest +from vllm.v1.engine import EngineCoreRequest, EngineCoreRequestType from vllm.v1.engine.core import DPEngineCoreProc, EngineCore, EngineCoreProc from vllm.v1.executor.abstract import Executor from vllm.v1.executor.uniproc_executor import UniProcExecutor @@ -689,6 +689,54 @@ def test_dp_sync_interval_idle_pause_consensus_on_first_step(monkeypatch): assert not core.pending_pause +@pytest.mark.parametrize( + "request_wave,engines_running,pause_state,announces", + [ + pytest.param(3, False, PauseState.UNPAUSED, True, id="current-wave"), + pytest.param(2, False, PauseState.UNPAUSED, True, id="stale-wave"), + pytest.param(4, False, PauseState.UNPAUSED, True, id="newer-wave"), + pytest.param(3, True, PauseState.UNPAUSED, False, id="running"), + pytest.param(3, False, PauseState.PAUSED_ALL, False, id="paused"), + ], +) +def test_idle_dp_rank_announces_wave_for_new_work( + request_wave, engines_running, pause_state, announces +): + """An idle rank handed work announces the wave itself, since only engines + start waves; a paused rank waits for resume.""" + core = object.__new__(DPEngineCoreProc) + core.has_coordinator = True + core.current_wave = 3 + core.engines_running = engines_running + core.scheduler = MagicMock(pause_state=pause_state) + core.output_queue = MagicMock() + + with patch.object(EngineCore, "add_request"): + core.add_request(MagicMock(), request_wave) + + if announces: + assert core.engines_running + _, outputs = core.output_queue.put_nowait.call_args.args[0] + assert outputs.start_wave == max(request_wave, 3) + else: + core.output_queue.put_nowait.assert_not_called() + + +@pytest.mark.parametrize("paused", [False, True]) +def test_paused_dp_rank_ignores_start_wave(paused): + """A START_DP_WAVE still in flight when the pause completes must not wake + the paused rank.""" + core = object.__new__(DPEngineCoreProc) + core.ignore_start_dp_wave = paused + core.engine_index = 0 + core.current_wave = 3 + core.engines_running = False + + core._handle_client_request(EngineCoreRequestType.START_DP_WAVE, (3, 1)) + + assert core.engines_running != paused + + def _pausable_engine_core_proc() -> EngineCoreProc: """A bare EngineCoreProc holding just the state pause_scheduler touches.""" core = object.__new__(EngineCoreProc) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index b58aed8547..e2804e9aa5 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -314,7 +314,7 @@ async def test_dplb_scale_down_routes_after_stale_stats_snapshot(): client.current_wave = 0 client.resources = SimpleNamespace(stats_update_task=None) client.stats_update_address = "inproc://scale-down-stats" - client.first_req_sock_addr = "inproc://scale-down-first-request" + client.scale_sock_addr = "inproc://scale-down-notify" pause_started = asyncio.Event() async def pause_scheduler(*args, **kwargs): @@ -326,11 +326,11 @@ async def test_dplb_scale_down_routes_after_stale_stats_snapshot(): with ( zmq.asyncio.Context() as ctx, ctx.socket(zmq.XPUB) as coordinator, - ctx.socket(zmq.PAIR) as first_req_socket, + ctx.socket(zmq.PAIR) as scale_socket, ): client.ctx = ctx coordinator.bind(client.stats_update_address) - first_req_socket.bind(client.first_req_sock_addr) + scale_socket.bind(client.scale_sock_addr) client._ensure_stats_update_task() scale_task = None try: diff --git a/vllm/distributed/elastic_ep/external_elastic_ep.py b/vllm/distributed/elastic_ep/external_elastic_ep.py index 6db6e5f4b1..66d94b4371 100644 --- a/vllm/distributed/elastic_ep/external_elastic_ep.py +++ b/vllm/distributed/elastic_ep/external_elastic_ep.py @@ -838,7 +838,7 @@ class ExternalElasticEPScaleCoordinator: if prepared.dp_rank == 0: self.client._coord_store = self.reconfig_store_ref self.client._ensure_stats_update_task() - await self.client.first_req_send_socket.send( + await self.client.scale_send_socket.send( msgspec.msgpack.encode( ("SCALE_ELASTIC_EP", bootstrap.new_data_parallel_size) ) diff --git a/vllm/v1/engine/coordinator.py b/vllm/v1/engine/coordinator.py index 28a4c7af6d..524794dac0 100644 --- a/vllm/v1/engine/coordinator.py +++ b/vllm/v1/engine/coordinator.py @@ -41,12 +41,9 @@ class DPCoordinator: DPEngineCoreProc._has_global_unfinished_reqs method. * Broadcasts the START_DP_WAVE message to engines to move them from paused - to running state when one engine receives a new request. This can happen - in two cases: - 1) A front-end sending a new request while the engines are paused will - concurrently notify the coordinator. - 2) An engine receiving a request for a stale request wave while in paused - state will notify the coordinator. + to running state when an idle engine receives a new request and notifies + the coordinator. Only the engines start waves: a front-end's view of the + running state can be stale, and a paused engine discards START_DP_WAVE. Engines will move into running state when receiving a new request or START_DP_WAVE message. @@ -353,27 +350,9 @@ class DPCoordinatorProc: ) continue # Skip normal engine notification processing - # Wave coordination: handle new-request messages from front-end. - # Only process these when wave coordination is enabled - if self.enable_wave_coordination: - # We received a message on the front-end XPUB socket, - # from an API server sending a new request while the - # engines are paused, so that we can wake the other - # engines. - engine_to_exclude, wave = decoded - if not engines_running: - if wave < current_wave: - # If the wave number is stale, ensure the message - # is handled by all the engines. - engine_to_exclude = None - - # engines_running is only set from the engines' - # own notifications; a paused engine discards - # START_DP_WAVE, so sending it is not evidence - # that the engines are running. - self._send_start_wave( - publish_back, current_wave, engine_to_exclude - ) + logger.error( + "DP Coordinator received unexpected message from front-end" + ) if output_back in events: # We received a message from one of the engines. @@ -448,13 +427,12 @@ class DPCoordinatorProc: wave > current_wave or (wave == current_wave and not engines_running) ): - # 3. The engine received request for a non-current wave - # so we must ensure that other engines progress to the - # next wave (race condition handling). + # 3. An engine started a wave (e.g. an idle engine + # received a request), so the others must join it. logger.debug( - "Starting wave %d after notification of " - "stale wave request from engine.", + "Starting wave %d on notification from engine %d.", wave, + eng_index, ) current_wave = wave engines_running = True diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index c9843dce96..2b275b65c5 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -2269,15 +2269,14 @@ class DPEngineCoreProc(EngineCoreProc): def add_request(self, request: Request, request_wave: int = 0): super().add_request(request, request_wave) - if self.has_coordinator and request_wave != self.current_wave: + if self.has_coordinator: if request_wave > self.current_wave: self.current_wave = request_wave - elif ( + if ( not self.engines_running and self.scheduler.pause_state == PauseState.UNPAUSED ): - # Request received for an already-completed wave, notify - # front-end that we need to start the next one. + # Engines alone start waves: have the coordinator wake the others. self.engines_running = True self.output_queue.put_nowait( (-1, EngineCoreOutputs(start_wave=self.current_wave)) diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index cfe8976668..0c87c9a052 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -506,8 +506,8 @@ class BackgroundResources: coordinator: DPCoordinator | None = None output_socket: zmq.Socket | zmq.asyncio.Socket | None = None input_socket: zmq.Socket | zmq.asyncio.Socket | None = None - first_req_send_socket: zmq.asyncio.Socket | None = None - first_req_rcv_socket: zmq.asyncio.Socket | None = None + scale_send_socket: zmq.asyncio.Socket | None = None + scale_rcv_socket: zmq.asyncio.Socket | None = None stats_update_socket: zmq.asyncio.Socket | None = None output_queue_task: asyncio.Task | None = None stats_update_task: asyncio.Task | None = None @@ -535,8 +535,8 @@ class BackgroundResources: sockets = ( self.output_socket, self.input_socket, - self.first_req_send_socket, - self.first_req_rcv_socket, + self.scale_send_socket, + self.scale_rcv_socket, self.stats_update_socket, ) @@ -1493,9 +1493,9 @@ class DPAsyncMPClient(AsyncMPClient): self.eep_scaling_cache: ElasticScalingCache | None = None - self.first_req_sock_addr = get_open_zmq_inproc_path() - self.first_req_send_socket = self.resources.first_req_send_socket = ( - make_zmq_socket(self.ctx, self.first_req_sock_addr, zmq.PAIR, bind=True) + self.scale_sock_addr = get_open_zmq_inproc_path() + self.scale_send_socket = self.resources.scale_send_socket = make_zmq_socket( + self.ctx, self.scale_sock_addr, zmq.PAIR, bind=True ) try: # If we are running in an asyncio event loop, start the stats task. @@ -1518,75 +1518,48 @@ class DPAsyncMPClient(AsyncMPClient): with ( make_zmq_socket(self.ctx, stats_addr, zmq.XSUB, linger=0) as socket, make_zmq_socket( - self.ctx, self.first_req_sock_addr, zmq.PAIR, bind=False, linger=0 - ) as first_req_rcv_socket, + self.ctx, self.scale_sock_addr, zmq.PAIR, bind=False, linger=0 + ) as scale_rcv_socket, ): assert isinstance(socket, zmq.asyncio.Socket) - assert isinstance(first_req_rcv_socket, zmq.asyncio.Socket) + assert isinstance(scale_rcv_socket, zmq.asyncio.Socket) self.resources.stats_update_socket = socket - self.resources.first_req_rcv_socket = first_req_rcv_socket + self.resources.scale_rcv_socket = scale_rcv_socket # Send subscription message. await socket.send(b"\x01") poller = zmq.asyncio.Poller() poller.register(socket, zmq.POLLIN) - poller.register(first_req_rcv_socket, zmq.POLLIN) + poller.register(scale_rcv_socket, zmq.POLLIN) while True: - events = await poller.poll() - if ( - not self.engines_running - and len(events) == 2 - or (events[0][0] == first_req_rcv_socket) - ): - # Check if this is a regular request notification or - # scale up notification - buf = first_req_rcv_socket.recv(flags=zmq.NOBLOCK).result() - - decoded = msgspec.msgpack.decode(buf) - if ( - isinstance(decoded, (list, tuple)) - and len(decoded) == 2 - and decoded[0] == "SCALE_ELASTIC_EP" - ): - # Extract new engine count from the decoded message - new_engine_count = decoded[1] - # Update engine_ranks_managed and count_slice - parallel_config = self.vllm_config.parallel_config - dp_size = parallel_config.data_parallel_size - dp_rank = parallel_config.data_parallel_rank - assert dp_rank == 0 - assert dp_size == new_engine_count - assert not parallel_config.data_parallel_hybrid_lb - # External LB clients continue to manage only local ranks. - if not parallel_config.data_parallel_external_lb: - self.engine_ranks_managed = list(range(dp_size)) - if len(self.lb_engines) < new_engine_count: - self.lb_engines += [ - [0, 0, 0.0] - for _ in range( - new_engine_count - len(self.lb_engines) - ) - ] - else: - self.lb_engines = self.lb_engines[:new_engine_count] - # Send scale up notification to coordinator - scale_msg = msgspec.msgpack.encode( - ("SCALE_ELASTIC_EP", new_engine_count) - ) - await socket.send(scale_msg) - continue - - # we're sending a request while the engines are - # paused, so that it can wake the others up - # (to run dummy EP loop). - assert decoded[0] == "FIRST_REQ" - target_eng_index = decoded[1] - self.engines_running = True - msg = msgspec.msgpack.encode( - (target_eng_index, self.current_wave) - ) - await socket.send(msg) + events = dict(await poller.poll()) + if scale_rcv_socket in events: + buf = scale_rcv_socket.recv(flags=zmq.NOBLOCK).result() + msg_type, new_engine_count = msgspec.msgpack.decode(buf) + assert msg_type == "SCALE_ELASTIC_EP" + # Update engine_ranks_managed and count_slice + parallel_config = self.vllm_config.parallel_config + dp_size = parallel_config.data_parallel_size + dp_rank = parallel_config.data_parallel_rank + assert dp_rank == 0 + assert dp_size == new_engine_count + assert not parallel_config.data_parallel_hybrid_lb + # External LB clients continue to manage only local ranks. + if not parallel_config.data_parallel_external_lb: + self.engine_ranks_managed = list(range(dp_size)) + if len(self.lb_engines) < new_engine_count: + self.lb_engines += [ + [0, 0, 0.0] + for _ in range( + new_engine_count - len(self.lb_engines) + ) + ] + else: + self.lb_engines = self.lb_engines[:new_engine_count] + # Forward the scale notification to the coordinator. + await socket.send(buf) + continue buf = None while True: @@ -1624,14 +1597,9 @@ class DPAsyncMPClient(AsyncMPClient): request.current_wave = self.current_wave request.client_index = self.client_index + # The engine that receives it wakes the other ranks if they are idle. chosen_engine = self.get_core_engine_for_request(request) - to_await = self._send_input(EngineCoreRequestType.ADD, request, chosen_engine) - if not self.engines_running: - # Notify coordinator that we're sending a request - req_msg = msgspec.msgpack.encode(("FIRST_REQ", chosen_engine)) - await self.first_req_send_socket.send(req_msg) - - await to_await + await self._send_input(EngineCoreRequestType.ADD, request, chosen_engine) self._ensure_output_queue_task() @@ -2100,7 +2068,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient): scale_up_marker = msgspec.msgpack.encode( ("SCALE_ELASTIC_EP", new_data_parallel_size) ) - await self.first_req_send_socket.send(scale_up_marker) + await self.scale_send_socket.send(scale_up_marker) logger.info( "[Elastic EP] Scale up completed, new data parallel size: %s", @@ -2176,7 +2144,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient): scale_down_marker = msgspec.msgpack.encode( ("SCALE_ELASTIC_EP", new_data_parallel_size) ) - await self.first_req_send_socket.send(scale_down_marker) + await self.scale_send_socket.send(scale_down_marker) await wait_future await self.resume_scheduler_async() except Exception: