iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[DP] Engines own wave wake-ups; drop the front-end FIRST_REQ (#59559)

Browse Source 
Signed-off-by: fh3 <a@b>
Co-authored-by: fh3 <a@b>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
aoshen02
2026-10-16 19:41:54 -04:00
committed by GitHub
co-authored by fh3 Claude Opus 5.5
parent b544587408
commit a94c315711
13 changed files with 226 additions and 523 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/engine-core-client/src/client.rs
+2 -5
View File
@@ -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!(
rust/src/engine-core-client/src/coordinator/external.rs
+14 -83
View File
@@ -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:
/// <https://github.com/vllm-project/vllm/blob/694449050f8dac3d9853e97e518b4a43ec52106a/vllm/v1/engine/core_client.py#L1230-L1236>
#[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<CoordinatorState>,
command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
socket: XSubSocket,
}
impl ExternalCoordinatorService {
pub(super) fn new(
state: Arc<CoordinatorState>,
command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
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<CoordinatorState>, 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<ClientInner>) {
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;
rust/src/engine-core-client/src/coordinator/handle.rs
+10 -81
View File
@@ -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<u32>) {
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<CoordinatorStateSnapshot>;
/// 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<CoordinatorState>,
command_tx: mpsc::UnboundedSender<CoordinatorCommand>,
}
impl CoordinatorHandle {
fn new_parts() -> (
Self,
Arc<CoordinatorState>,
mpsc::UnboundedReceiver<CoordinatorCommand>,
) {
fn new_parts() -> (Self, Arc<CoordinatorState>) {
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(())
}
}
rust/src/engine-core-client/src/coordinator/inproc.rs
+10 -98
View File
@@ -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<u32>,
}
/// 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<CoordinatorState>,
command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
coordinator_input: XPubSocket,
}
impl InProcCoordinatorRunner {
pub(super) fn new(
state: Arc<CoordinatorState>,
command_rx: mpsc::UnboundedReceiver<CoordinatorCommand>,
coordinator_input: XPubSocket,
) -> Self {
pub(super) fn new(state: Arc<CoordinatorState>, 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");
}
}
rust/src/engine-core-client/src/error.rs
-4
View File
@@ -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<u8> },
#[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(
rust/src/engine-core-client/src/tests/client.rs
+28 -103
View File
@@ -267,6 +267,19 @@ async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, Option<u32>) {
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<T: serde::Serialize>(
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();
}
tests/v1/distributed/test_async_llm_dp.py
+53 -33
View File
@@ -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"
)
tests/v1/engine/test_engine_core.py
+49 -1
View File
@@ -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)
tests/v1/engine/test_engine_core_client.py
+3 -3
View File
@@ -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:
vllm/distributed/elastic_ep/external_elastic_ep.py
+1 -1
View File
@@ -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)
)
vllm/v1/engine/coordinator.py
+10 -32
View File
@@ -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
vllm/v1/engine/core.py
+3 -4
View File
@@ -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))
vllm/v1/engine/core_client.py
+43 -75
View File
@@ -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: