diff --git a/crates/engine_zmq_client/src/connector.rs b/crates/engine_zmq_client/src/connector.rs index 13ce8ef7e..1af0f4adf 100644 --- a/crates/engine_zmq_client/src/connector.rs +++ b/crates/engine_zmq_client/src/connector.rs @@ -5,7 +5,8 @@ // - No client-side DP load balancing. SMG routing picks the rank and stamps // `data_parallel_rank`; the connector consumes the piggybacked per-rank // `scheduler_stats` as the load signal. -// - No DP coordinator / wave protocol yet (single shared transport, DP=1). +// - No DP coordinator process: for a lockstep engine group the connector plays +// the wake role itself, over the input socket each rank already listens on. // - No utility RPC (deferred). use std::{collections::HashMap, sync::Arc, time::Duration}; @@ -21,7 +22,7 @@ use crate::{ error::{Error, Result}, protocol::{ tokenspeed::TokenSpeedProtocol, vllm::VllmProtocol, EngineBatch, EngineLoad, EngineOutput, - EngineProtocol, + EngineProtocol, WaveEvent, }, transport::{run_output_loop, send_message, ConnectedEngine, ConnectedTransport, EngineId}, }; @@ -100,6 +101,15 @@ impl RequestRegistry { } } +/// Per-group wave bookkeeping (see [`WaveEvent`] for what a wave is), mirroring +/// the state vLLM's DP coordinator keeps: which wave the group is on, and +/// whether it is stepping. +#[derive(Default)] +struct WaveState { + current: u32, + running: bool, +} + struct ClientInner { /// Shared input ROUTER send half (serialized across concurrent submits). input_send: tokio::sync::Mutex, @@ -107,6 +117,9 @@ struct ClientInner { registry: Mutex>, /// Latest per-rank load, keyed by engine index. SMG's DP load signal. load: Mutex>, + /// Wave state for a lockstep engine group; `None` when the ranks step + /// independently (dense DP, DP=1) and so never pause together. + wave: Option>, /// Auto-abort channel fed by dropped streams. abort_tx: mpsc::UnboundedSender<(EngineId, String)>, } @@ -157,6 +170,85 @@ impl ClientInner

{ .await } + /// Tell every rank but `exclude_index` to start `wave`. Does nothing for a + /// protocol without a wave protocol. + async fn broadcast_start_wave(&self, wave: u32, exclude_index: u32) -> Result<()> { + let Some((frame, payload)) = P::encode_start_wave(wave, exclude_index)? else { + return Ok(()); + }; + for engine in &self.engines { + if engine.engine_id.engine_index() == Some(exclude_index) { + continue; + } + self.send_to_engine( + &engine.engine_id, + frame.clone(), + payload.clone(), + Vec::new(), + ) + .await?; + } + Ok(()) + } + + /// Wake a paused lockstep group so the rank now holding a request can make + /// progress: its peers must be stepping too, because every rank joins the + /// same all-reduce. Upstream vLLM has a DP coordinator process broadcast the + /// restart; SMG owns rank selection, so it sends the same signal over the + /// input socket each rank already listens on. A no-op for independent ranks + /// and for a group that is already running. + async fn wake_group(&self, holder_index: u32) -> Result<()> { + let Some(state) = self.wave.as_ref() else { + return Ok(()); + }; + let wave = { + let mut state = state.lock(); + if state.running { + return Ok(()); + } + state.running = true; + state.current + }; + if let Err(error) = self.broadcast_start_wave(wave, holder_index).await { + // The group is still asleep, so let the next submit try again. + state.lock().running = false; + return Err(error); + } + Ok(()) + } + + /// Fold a wave notification from an engine into the group's state. + async fn observe_wave(&self, event: WaveEvent, engine_index: u32) { + let Some(state) = self.wave.as_ref() else { + warn!(?event, "wave notification from independent ranks; ignoring"); + return; + }; + match event { + // The group drained the wave and parked itself; the engines have + // already moved on to the next one. The next submit wakes them. + WaveEvent::Complete(wave) => { + let mut state = state.lock(); + if wave >= state.current { + state.current = wave.saturating_add(1); + state.running = false; + } + } + // A rank took a request for an already-drained wave and is asking + // for the rest of the group to catch up. + WaveEvent::Start(wave) => { + { + let mut state = state.lock(); + state.current = state.current.max(wave); + state.running = true; + } + if let Err(error) = self.broadcast_start_wave(wave, engine_index).await { + warn!(%error, wave, "failed to start the requested wave"); + state.lock().running = false; + } + } + } + } + /// Send an Abort for one request id to its engine. async fn abort(&self, engine_id: &EngineId, request_id: &str) -> Result<()> { let payload = P::encode_abort(request_id)?; @@ -183,12 +275,21 @@ impl Client

{ .. } = transport; + // Only a group whose ranks step in lockstep pauses as a unit and needs + // waking. The engines report which they are at handshake: a lockstep + // group keeps its data-parallel size, while independent ranks are + // reconfigured to a size of one before they answer. + let lockstep = engines + .iter() + .any(|engine| engine.ready_response.data_parallel_size > 1); + let (abort_tx, abort_rx) = mpsc::unbounded_channel(); let inner = Arc::new(ClientInner { input_send: tokio::sync::Mutex::new(input_send), engines, registry: Mutex::new(RequestRegistry::default()), load: Mutex::new(HashMap::new()), + wave: lockstep.then(|| Mutex::new(WaveState::default())), abort_tx, }); @@ -250,6 +351,21 @@ impl Client

{ return Err(error); } + // The rank now holds the request; its peers must be awake for it to + // step. Waking after the send keeps the group parked (and so unable to + // report another drained wave) for the whole window. + if let Err(error) = self + .inner + .wake_group(engine_id.engine_index().unwrap_or(u32::MAX)) + .await + { + // The request can never make progress with its peers asleep, so + // take it back rather than leaving it stranded on the engine. + self.inner.registry.lock().remove_all([&request_id]); + let _ = self.inner.abort(&engine_id, &request_id).await; + return Err(error); + } + Ok(RequestStream { request_id, engine_id, @@ -300,6 +416,9 @@ async fn run_dispatcher( if let Some(load) = batch.load { inner.load.lock().insert(batch.engine_index, load); } + if let Some(event) = batch.wave { + inner.observe_wave(event, batch.engine_index).await; + } let mut registry = inner.registry.lock(); for output in batch.outputs { registry.route(output); @@ -419,13 +538,19 @@ mod tests { use super::*; use crate::{ codec::{decode_msgpack, encode_msgpack}, - mock_engine::{connect_to_frontend, default_ready_response, IpcNamespace, MockEngine}, - protocol::vllm::{ - output::{ - EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, + mock_engine::{ + connect_to_frontend, default_ready_response, EngineInbound, IpcNamespace, MockEngine, + }, + protocol::{ + handshake::EngineCoreReadyResponse, + vllm::{ + output::{ + DpControlMessage, DpControlOutput, EngineCoreFinishReason, EngineCoreOutput, + EngineCoreOutputs, RequestBatchOutputs, + }, + request::EngineCoreRequest, + stats::SchedulerStats, }, - request::EngineCoreRequest, - stats::SchedulerStats, }, transport::{connect_handshake, ENGINE_CORE_DEAD_SENTINEL}, }; @@ -463,6 +588,62 @@ mod tests { ) } + /// Connect `engine_count` mock ranks behind one transport. `lockstep` + /// controls the data-parallel size they report at handshake, which is how + /// the client tells a group that pauses as a unit (vLLM MoE DP) from ranks + /// that step independently. + async fn connect_ranks( + engine_count: usize, + lockstep: bool, + ) -> (EngineCoreClient, Vec, IpcNamespace) { + let ns = IpcNamespace::new().unwrap(); + let (handshake, input, output) = ( + ns.handshake_endpoint(), + ns.input_endpoint(), + ns.output_endpoint(), + ); + let ready = |rank: u32| EngineCoreReadyResponse { + data_parallel_size: if lockstep { engine_count as u64 } else { 1 }, + data_parallel_rank: rank, + ..default_ready_response() + }; + let engines = (0..engine_count as u32).map(|rank| { + connect_to_frontend(&handshake, EngineId::from_engine_index(rank), ready(rank)) + }); + let (transport, engines) = tokio::join!( + connect_handshake( + &handshake, + engine_count, + "127.0.0.1", + Some(&input), + Some(&output), + TIMEOUT + ), + futures::future::join_all(engines), + ); + let engines = engines.into_iter().map(Result::unwrap).collect(); + (EngineCoreClient::new(transport.unwrap()), engines, ns) + } + + /// An add-request pinned to one rank. + fn request_for(request_id: &str, rank: u32) -> EngineCoreRequest { + EngineCoreRequest { + request_id: request_id.to_string(), + prompt_token_ids: Some(vec![1, 2, 3]), + data_parallel_rank: Some(rank), + ..EngineCoreRequest::default() + } + } + + fn wave_control(engine_index: u32, control: DpControlMessage) -> Vec { + let outputs = EngineCoreOutputs::DpControl(DpControlOutput { + engine_index, + timestamp: 0.0, + control, + }); + vec![Bytes::from(encode_msgpack(&outputs).unwrap())] + } + fn batch(engine_index: u32, output: EngineCoreOutput) -> Vec { let finished = output.finish_reason.map(|_| { let mut set = std::collections::BTreeSet::new(); @@ -709,4 +890,84 @@ mod tests { // Terminal output ends the stream. assert!(stream.next().await.is_none()); } + + /// A lockstep group is paused when the client connects, so the first submit + /// must wake every rank except the one taking the request. + #[tokio::test] + async fn first_submit_wakes_the_paused_lockstep_group() { + let (client, mut engines, _ns) = connect_ranks(2, true).await; + + let _stream = client.submit(request_for("req-1", 1)).await.unwrap(); + + match engines[1].recv().await.unwrap() { + EngineInbound::Add(request) => assert_eq!(request.request_id, "req-1"), + other => panic!("rank 1 expected the Add, got {other:?}"), + } + assert!(matches!( + engines[0].recv().await.unwrap(), + EngineInbound::StartDpWave { + wave: 0, + exclude_engine_index: 1, + } + )); + } + + /// While the group runs, submits carry no wake; once it reports the wave + /// drained, the next submit wakes it again on the following wave. + #[tokio::test] + async fn a_drained_wave_re_arms_the_wake() { + let (client, mut engines, _ns) = connect_ranks(2, true).await; + + let _first = client.submit(request_for("req-1", 1)).await.unwrap(); + assert!(matches!( + engines[0].recv().await.unwrap(), + EngineInbound::StartDpWave { wave: 0, .. } + )); + + // Running group: the second submit only sends the Add. + let _second = client.submit(request_for("req-2", 0)).await.unwrap(); + match engines[0].recv().await.unwrap() { + EngineInbound::Add(request) => assert_eq!(request.request_id, "req-2"), + other => panic!("expected the Add with no wake, got {other:?}"), + } + + // The group drains wave 0 and parks itself. + engines[0] + .send_output(wave_control(0, DpControlMessage::WaveComplete(0))) + .await + .unwrap(); + // The notification travels through the dispatcher, so wait for it to + // land before submitting against the parked group. + tokio::time::timeout(TIMEOUT, async { + while client.inner.wave.as_ref().unwrap().lock().running { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("the drained wave never parked the group"); + + let _third = client.submit(request_for("req-3", 1)).await.unwrap(); + match engines[0].recv().await.unwrap() { + // The engines moved on to wave 1 as they paused, so that is the + // wave the wake must name. + EngineInbound::StartDpWave { wave, .. } => assert_eq!(wave, 1), + other => panic!("expected a wake for the next wave, got {other:?}"), + } + } + + /// Independent ranks (dense DP) never pause as a group, so they are never + /// woken — each rank only ever sees its own requests. + #[tokio::test] + async fn independent_ranks_are_never_woken() { + let (client, mut engines, _ns) = connect_ranks(2, false).await; + + let _first = client.submit(request_for("req-1", 1)).await.unwrap(); + let _second = client.submit(request_for("req-2", 0)).await.unwrap(); + + // Rank 0's first message is its own Add, not a wake for rank 1's. + match engines[0].recv().await.unwrap() { + EngineInbound::Add(request) => assert_eq!(request.request_id, "req-2"), + other => panic!("independent rank expected only its Add, got {other:?}"), + } + } } diff --git a/crates/engine_zmq_client/src/mock_engine.rs b/crates/engine_zmq_client/src/mock_engine.rs index 17de2b1e3..24a53fc31 100644 --- a/crates/engine_zmq_client/src/mock_engine.rs +++ b/crates/engine_zmq_client/src/mock_engine.rs @@ -94,7 +94,13 @@ pub enum EngineInbound { Add(Box), /// An abort for the given request ids (`EngineCoreRequestType::Abort`). Abort(Vec), - /// Any other request type byte (StartDpWave / Utility), unhandled here. + /// A lockstep-group wake (`EngineCoreRequestType::StartDpWave`): start this + /// wave unless this engine is the excluded one. + StartDpWave { + wave: u32, + exclude_engine_index: u32, + }, + /// Any other request type byte (Utility), unhandled here. Other(u8), } @@ -125,6 +131,13 @@ impl MockEngineInput { Some(EngineCoreRequestType::Abort) => { Ok(EngineInbound::Abort(decode_msgpack(payload)?)) } + Some(EngineCoreRequestType::StartDpWave) => { + let (wave, exclude_engine_index) = decode_msgpack(payload)?; + Ok(EngineInbound::StartDpWave { + wave, + exclude_engine_index, + }) + } Some(other) => Ok(EngineInbound::Other(other as u8)), None => Err(Error::UnexpectedHandshakeMessage { message: format!("unknown request type frame {:?}", type_frame.as_ref()), @@ -182,6 +195,12 @@ impl MockEngine { self.input.recv_frames().await } + /// Receive and classify the next request. Convenience for sequential test + /// drivers that care about the request type rather than the raw frames. + pub async fn recv(&mut self) -> Result { + self.input.recv().await + } + /// Push a raw multi-frame output. Convenience for sequential test drivers. pub async fn send_output(&mut self, frames: Vec) -> Result<()> { self.output.send_frames(frames).await diff --git a/crates/engine_zmq_client/src/protocol/mod.rs b/crates/engine_zmq_client/src/protocol/mod.rs index 23ca47533..6cd81055b 100644 --- a/crates/engine_zmq_client/src/protocol/mod.rs +++ b/crates/engine_zmq_client/src/protocol/mod.rs @@ -30,6 +30,22 @@ pub struct EngineLoad { pub kv_cache_usage: f64, } +/// A *wave* is vLLM's term (not ours) for one round of work by a data-parallel +/// group whose ranks step in lockstep: they all-reduce every step, so they drain +/// the wave together and then all park. A request handed to one parked rank +/// leaves the others asleep until someone starts the next wave. Upstream names: +/// `current_wave`, `wave_complete`, `START_DP_WAVE` in +/// `vllm/v1/engine/core.py`. Only vLLM's MoE data parallelism steps this way; +/// dense DP ranks are independent and report no waves. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WaveEvent { + /// The group drained this wave and is now parked (vLLM `wave_complete`). + Complete(u32), + /// A rank took a request for an already-drained wave and needs the rest of + /// the group started on this one (vLLM `start_wave`). + Start(u32), +} + /// A single per-request output decoded from one engine tick. Lets the /// engine-neutral connector route outputs to their request streams without /// knowing the concrete protocol payload. @@ -52,6 +68,9 @@ pub struct EngineBatch { pub finished_request_ids: Vec, /// Per-rank load signal, when the protocol carries it (`None` otherwise). pub load: Option, + /// Wave-control notification from a lockstep engine group (`None` on every + /// ordinary tick). + pub wave: Option, } impl Default for EngineBatch { @@ -61,6 +80,7 @@ impl Default for EngineBatch { outputs: Vec::new(), finished_request_ids: Vec::new(), load: None, + wave: None, } } } @@ -92,6 +112,11 @@ pub trait EngineProtocol: Send + Sync + 'static { fn encode_add(request: &Self::Request) -> Result<(Vec, Vec)>; /// Encode the abort payload for one request id. fn encode_abort(request_id: &str) -> Result>; + /// Encode "start `wave` on every rank but `exclude_engine_index`" (the one + /// already holding the request) as `(request-type frame, payload)`. + /// `Ok(None)` when the protocol has no wave protocol — its ranks run + /// independently, so there is nothing to wake. + fn encode_start_wave(wave: u32, exclude_engine_index: u32) -> Result)>>; /// Decode one output message (frame 0 plus ordered aux frames) into a batch. fn decode_batch(frames: &[Bytes]) -> Result>; } diff --git a/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs b/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs index ed5968247..6264c778c 100644 --- a/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs +++ b/crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs @@ -121,6 +121,15 @@ impl EngineProtocol for TokenSpeedProtocol { encode_msgpack(&[request_id.to_string()]) } + fn encode_start_wave( + _wave: u32, + _exclude_engine_index: u32, + ) -> Result)>> { + // TokenSpeed has no wave protocol: its scheduler never pauses a group + // of ranks, so there is nothing to wake. + Ok(None) + } + fn decode_batch(frames: &[Bytes]) -> Result> { // Output messages are `[payload, aux...]`. The slim batch carries no // tensor fields today, so aux frames are unexpected but not fatal. @@ -146,6 +155,7 @@ impl EngineProtocol for TokenSpeedProtocol { outputs, finished_request_ids, load: None, + wave: None, }) } } diff --git a/crates/engine_zmq_client/src/protocol/vllm/mod.rs b/crates/engine_zmq_client/src/protocol/vllm/mod.rs index 933a0b32a..618d5e681 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/mod.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/mod.rs @@ -28,11 +28,13 @@ use crate::{ error::Result, protocol::{ vllm::{ - output::{decode_engine_core_outputs, EngineCoreOutput, EngineCoreOutputs}, + output::{ + decode_engine_core_outputs, DpControlMessage, EngineCoreOutput, EngineCoreOutputs, + }, request::{EngineCoreRequest, EngineCoreRequestType}, stats::SchedulerStats, }, - EngineBatch, EngineLoad, EngineOutput, EngineProtocol, + EngineBatch, EngineLoad, EngineOutput, EngineProtocol, WaveEvent, }, }; @@ -94,10 +96,19 @@ impl EngineProtocol for VllmProtocol { encode_msgpack(&[request_id.to_string()]) } + fn encode_start_wave(wave: u32, exclude_engine_index: u32) -> Result)>> { + // Python decodes this with the generic msgpack decoder and unpacks it + // positionally as `new_wave, exclude_eng_index`. + Ok(Some(( + EngineCoreRequestType::StartDpWave.to_frame(), + encode_msgpack(&(wave, exclude_engine_index))?, + ))) + } + fn decode_batch(frames: &[Bytes]) -> Result> { // vLLM multiplexes request batches, utility RPCs, and DP control on one - // wire struct; only request batches carry per-request outputs (the - // others surface as an empty batch the dispatcher ignores). + // wire struct; only request batches carry per-request outputs (utility + // results surface as an empty batch the dispatcher ignores). match decode_engine_core_outputs(frames)? { EngineCoreOutputs::RequestBatch(batch) => Ok(EngineBatch { engine_index: batch.engine_index, @@ -107,10 +118,17 @@ impl EngineProtocol for VllmProtocol { .map(|ids| ids.into_iter().collect()) .unwrap_or_default(), load: batch.scheduler_stats.map(|stats| EngineLoad::from(*stats)), + wave: None, + }), + EngineCoreOutputs::DpControl(control) => Ok(EngineBatch { + engine_index: control.engine_index, + wave: Some(match control.control { + DpControlMessage::WaveComplete(wave) => WaveEvent::Complete(wave), + DpControlMessage::StartWave(wave) => WaveEvent::Start(wave), + }), + ..EngineBatch::default() }), - EngineCoreOutputs::Utility(_) | EngineCoreOutputs::DpControl(_) => { - Ok(EngineBatch::default()) - } + EngineCoreOutputs::Utility(_) => Ok(EngineBatch::default()), } } } diff --git a/crates/mock_worker/src/zmq.rs b/crates/mock_worker/src/zmq.rs index 8c3c81c63..0a38c6239 100644 --- a/crates/mock_worker/src/zmq.rs +++ b/crates/mock_worker/src/zmq.rs @@ -128,6 +128,11 @@ pub async fn serve(cfg: Arc, handshake_address: String, engine_index: u3 )); } } + Ok(EngineInbound::StartDpWave { wave, .. }) => { + // The mock is a single independent engine, never a lockstep + // group, so it never pauses and has no wave to start. + tracing::debug!("zmq engine {engine_index} ignoring start of wave {wave}"); + } Ok(EngineInbound::Other(byte)) => { tracing::debug!("zmq engine {engine_index} ignoring request type {byte}"); }