feat(zmq): wake a paused lockstep engine group from the connector - #2078
Conversation
vLLM's MoE data-parallel ranks step in lockstep: they all-reduce every step, drain a wave together, and then park. A request handed to one rank leaves the others asleep, so somebody has to tell them to start the next wave. Upstream that somebody is a separate DP coordinator process; the engines only need it because their Python frontend has no other way to reach every rank at once. We already hold a socket to each rank, so the connector plays the wake role itself over the sockets it has. Widen the protocol seam with a wave notification on the decoded batch and an `encode_start_wave` hook. vLLM surfaces `wave_complete`/`start_wave` control messages (previously dropped on the floor) and encodes `START_DP_WAVE`; TokenSpeed reports neither and returns `Ok(None)`, so nothing wakes. The connector keeps the same wave state the coordinator would: the current wave, and whether the group is stepping. It arms only when the handshake reports a data-parallel size above one, which is exactly the lockstep case — dense DP ranks reconfigure themselves to a size of one and run independently. Submit adds the request first and wakes after, so the group stays parked for the whole window; if the wake fails the request is taken back rather than left stranded in a sleeping engine. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds lockstep wave coordination for multi-rank ZMQ clients. Protocols encode start-wave commands and decode wave lifecycle events. The connector tracks group state, wakes peer ranks, handles failures, and re-arms drained waves. Tests cover lockstep and independent-rank behavior. ChangesLockstep wave coordination
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant RankClient
participant ZmqConnector
participant PeerRank
participant vLLMProtocol
RankClient->>ZmqConnector: Submit request
ZmqConnector->>vLLMProtocol: Encode start-wave command
vLLMProtocol->>PeerRank: Send wave and excluded rank
PeerRank-->>ZmqConnector: Return wave start or completion event
ZmqConnector-->>RankClient: Advance or re-arm wave state
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
crates/engine_zmq_client/src/connector.rs (3)
957-971: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit —
independent_ranks_are_never_wokendoes not prove the negative.The test reads one message from rank 0 and asserts it is the Add. A wake sent after that Add would still pass. The name promises that no wake occurs at all.
Assert the absence directly. A short
tokio::time::timeoutaround a secondrecvproves no further message arrives. Assertingclient.inner.wave.is_none()also pins the mode.Two paths from the PR objectives are also untested: a failed wake that deregisters and aborts the request, and the
WaveEvent::Starthandling inobserve_wave.As per coding guidelines: "Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality."
💚 Proposed strengthening of the assertion
// 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:?}"), } + // No wake follows: independent ranks carry no wave state at all. + assert!(client.inner.wave.is_none()); + let extra = tokio::time::timeout(Duration::from_millis(200), engines[0].recv()).await; + assert!(extra.is_err(), "independent rank received an unexpected message");🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/engine_zmq_client/src/connector.rs` around lines 957 - 971, Strengthen `independent_ranks_are_never_woken` by asserting `client.inner.wave.is_none()` and wrapping a second `engines[0].recv()` in a short `tokio::time::timeout`, expecting it to time out after the initial `Add`. Add focused tests covering failed wake deregistration/request abortion and `WaveEvent::Start` handling in `observe_wave`, then run the PR test analyzer to verify coverage.Source: Coding guidelines
174-191: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win🟡 Nit —
broadcast_start_wavestops at the first failed peer send.The loop returns on the first
send_to_engineerror. Ranks already sent a wake stay awake, and the remaining ranks stay parked.submitthen aborts the request, so the woken ranks step with no work until they drain and reportWaveComplete.Consider sending to every peer and returning an error only after the full pass. This keeps the wake attempt symmetric and records which rank failed.
♻️ Proposed change to complete the broadcast before failing
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?; + if let Err(error) = self + .send_to_engine( + &engine.engine_id, + frame.clone(), + payload.clone(), + Vec::new(), + ) + .await + { + warn!(%error, ?engine.engine_id, wave, "failed to wake a rank"); + first_error.get_or_insert(error); + } } - Ok(()) + match first_error { + Some(error) => Err(error), + None => Ok(()), + }Declare
let mut first_error = None;before the loop.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/engine_zmq_client/src/connector.rs` around lines 174 - 191, Update broadcast_start_wave to continue sending to every eligible engine instead of returning immediately when send_to_engine fails. Track the first send error while completing the loop, then return that error after all peers have been attempted; preserve the existing exclusion and successful no-error behavior.
281-283: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value🟡 Nit — A mixed
data_parallel_sizereport resolves silently.
anyenables wave coordination when at least one rank reports a size above 1. If ranks disagree, the client picks a mode without reporting the disagreement.anyis the safe direction, because enabling coordination on an independent group only adds harmless wakes.Log a warning when the reported sizes are not uniform, so a misconfigured handshake is visible.
As per coding guidelines: "Do not silently fall back to None or a default when configuration validation should fail loudly."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/engine_zmq_client/src/connector.rs` around lines 281 - 283, Update the lockstep determination near engines iteration to detect whether all engines report the same data_parallel_size, while preserving the existing any-based behavior for enabling coordination. When the reported sizes are mixed, log a warning through the connector’s existing logging mechanism with enough context to identify the disagreement.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/engine_zmq_client/src/connector.rs`:
- Around line 237-247: Update the WaveEvent::Start branch to resolve and retain
the updated state.current value before broadcasting, then pass that resolved
wave to broadcast_start_wave instead of the original wave argument. Keep running
state management and failure handling unchanged, ensuring peers start the same
wave recorded by the connector when state.current is already higher.
---
Nitpick comments:
In `@crates/engine_zmq_client/src/connector.rs`:
- Around line 957-971: Strengthen `independent_ranks_are_never_woken` by
asserting `client.inner.wave.is_none()` and wrapping a second
`engines[0].recv()` in a short `tokio::time::timeout`, expecting it to time out
after the initial `Add`. Add focused tests covering failed wake
deregistration/request abortion and `WaveEvent::Start` handling in
`observe_wave`, then run the PR test analyzer to verify coverage.
- Around line 174-191: Update broadcast_start_wave to continue sending to every
eligible engine instead of returning immediately when send_to_engine fails.
Track the first send error while completing the loop, then return that error
after all peers have been attempted; preserve the existing exclusion and
successful no-error behavior.
- Around line 281-283: Update the lockstep determination near engines iteration
to detect whether all engines report the same data_parallel_size, while
preserving the existing any-based behavior for enabling coordination. When the
reported sizes are mixed, log a warning through the connector’s existing logging
mechanism with enough context to identify the disagreement.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: d17655d9-e042-4e72-9de1-d614ae5ec656
📒 Files selected for processing (5)
crates/engine_zmq_client/src/connector.rscrates/engine_zmq_client/src/mock_engine.rscrates/engine_zmq_client/src/protocol/mod.rscrates/engine_zmq_client/src/protocol/tokenspeed/mod.rscrates/engine_zmq_client/src/protocol/vllm/mod.rs
| 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; | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find tests exercising the StartWave path and any staleness assertions.
rg -n -C 6 'DpControlMessage::StartWave|WaveEvent::Start' --type=rustRepository: smg-project/smg
Length of output: 153
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate connector.rs =="
fd -a 'connector\.rs$' . || true
echo
echo "== target file outline/section =="
if [ -f crates/engine_zmq_client/src/connector.rs ]; then
wc -l crates/engine_zmq_client/src/connector.rs
sed -n '180,270p' crates/engine_zmq_client/src/connector.rs | nl -ba -v180
fi
echo
echo "== nearby state/update symbols =="
rg -n "struct WaveState|enum WaveEvent|observe_wave|broadcast_start_wave|submit|wake_group|DpControlMessage::StartWave|WaveEvent::Start|state\.current" crates/engine_zmq_client/src/connector.rsRepository: smg-project/smg
Length of output: 362
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== target file section =="
sed -n '180,270p' crates/engine_zmq_client/src/connector.rs
echo
echo "== relevant symbol usages =="
rg -n "struct WaveState|enum WaveEvent|observe_wave|broadcast_start_wave|submit\(|wake_group|DpControlMessage::StartWave|WaveEvent::Start|state\.current" crates/engine_zmq_client/src/connector.rs
echo
echo "== behavioral invariant probe =="
python3 - <<'PY'
class State:
def __init__(self):
self.current = 0
self.running = False
def handle_start(state, wave):
state.current = max(state.current, wave)
state.running = True
broadcast = wave
return broadcast
state = State()
state.current = 5
state.running = False
broadcast = handle_start(state, 3)
print(f"state.current={state.current} running={state.running} broadcast={broadcast}")
print(f"state and broadcast disagree: {state.current != broadcast}")
PYRepository: smg-project/smg
Length of output: 5567
🔴 Important — Broadcast the resolved start wave in WaveEvent::Start.
Complete(wave) advances state.current only for wave >= state.current, but Start(wave) always sets state.current = state.current.max(wave) and running = true while sending broadcast_start_wave(wave). If state.current is already higher, the connector marks the lockstep group running and tells peers to start an old wave, so later submit calls skip wake_group. Broadcast the resolved state.current value instead.
🐛 Proposed fix
WaveEvent::Start(wave) => {
- {
+ let wave = {
let mut state = state.lock();
state.current = state.current.max(wave);
state.running = true;
- }
+ state.current
+ };
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;
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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; | |
| } | |
| } | |
| WaveEvent::Start(wave) => { | |
| let wave = { | |
| let mut state = state.lock(); | |
| state.current = state.current.max(wave); | |
| state.running = true; | |
| state.current | |
| }; | |
| 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; | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/engine_zmq_client/src/connector.rs` around lines 237 - 247, Update the
WaveEvent::Start branch to resolve and retain the updated state.current value
before broadcasting, then pass that resolved wave to broadcast_start_wave
instead of the original wave argument. Keep running state management and failure
handling unchanged, ensuring peers start the same wave recorded by the connector
when state.current is already higher.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean implementation — wave bookkeeping, lockstep detection, and the add-then-wake ordering are all correct. Concurrency around WaveState is sound (short critical sections, lock released before async work, proper error rollback). Both protocol implementations updated consistently. Good test coverage of the three key scenarios. LGTM.
The mock worker matches `EngineInbound` exhaustively, so adding the start-wave variant broke its build. It simulates a single independent engine, never a lockstep group, so it never parks and has no wave to start: log and ignore. Also name vLLM as the source of the word "wave" where the enum is defined, with the upstream identifiers to grep for. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
1 participant
Why
vLLM's MoE data-parallel ranks step in lockstep: they all-reduce every step, drain a wave together, and then park. A request handed to one rank leaves the others asleep, so somebody has to tell them to start the next wave.
Upstream that somebody is a separate DP coordinator process. The engines only need it because their Python frontend has no other way to reach every rank at once — and the engine's poller dispatches the coordinator socket and the client socket through the same request-type switch, so
START_DP_WAVEis accepted on the socket we already own. We hold a socket to each rank, so the connector plays the wake role itself. No coordinator process, no extra sockets.What
EngineBatchcarries an optionalWaveEvent(Complete/Start), andEngineProtocolgainsencode_start_wave. vLLM surfaces thewave_complete/start_wavecontrol messages it was previously dropping on the floor and encodesSTART_DP_WAVE; TokenSpeed reports neither and returnsOk(None), so nothing wakes.Two decisions worth reviewing
When the wake arms. Only when the handshake reports
data_parallel_size > 1. That is precisely the lockstep case: dense DP ranks callreconfigure_for_independent_dp_rank()and answer the handshake with a size of one, so they never park as a group and never get woken. The connector branches on what the engine reports about itself, not on the model.Order of add and wake. Add first, wake after. Waking first admits a window where the group re-drains and parks again before the add lands, stranding the request in a sleeping engine. Adding first keeps the group parked for the whole window. If the wake then fails, the request is taken back (deregistered + aborted) rather than left behind.
Tests
Three new connector tests over N mock ranks: a first submit wakes the paused group excluding the holder, a drained wave re-arms the wake, and independent ranks are never woken.