feat(zmq): vLLM DP>1 as grouped engine workers with least-loaded selection - #2087
Merged
Conversation
slin1237
requested review from
CatherineSue,
gongwei-130 and
key4ng
as code owners
August 11, 2026 04:18
github-actions
Bot
added
python-bindings
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change adds optional grouped ZMQ engine configuration across Python and Rust interfaces. Worker startup and connection handshakes propagate engine counts. vLLM launchers support grouped engines. Unpinned ZMQ requests use load and in-flight state for engine selection. ChangesGrouped ZMQ engine support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Router
participant WorkerBuilder
participant ZmqEngineClient
participant Engine
Router->>WorkerBuilder: configure zmq_engine_count
WorkerBuilder->>ZmqEngineClient: connect with engine_count
ZmqEngineClient->>Engine: establish grouped handshake
Router->>ZmqEngineClient: submit unpinned request
ZmqEngineClient->>ZmqEngineClient: score and reserve an engine
ZmqEngineClient->>Engine: send request
Engine-->>ZmqEngineClient: report completion
Possibly related issues
Possibly related PRs
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 |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Solid PR — the grouped-worker model is well designed and the config plumbing is complete across all layers (CLI, types, builder, Python bindings, serve launcher, job queue, worker builder, connector).
Summary: 0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing
The one nit is about _backend_arg_int silently swallowing invalid values — a warning log would help operators catch misconfigured --data-parallel-size flags. Everything else checks out: the least-loaded engine selection mirrors vLLM's frontend DP client correctly, _filter_backend_args properly handles --key=value forms, the TokenSpeed rejection is well-placed, and the test coverage (connector load balancing, builder shape, serve flag ownership) exercises the key behaviors.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
crates/engine_zmq_client/src/connector.rs (1)
768-818: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit — Add coverage for both load sort keys and lockstep routing.
This test only proves that an engine without scheduler statistics scores zero. Add cases where both engines report load to verify
num_waitinghas priority overnum_running, and where equalnum_waitingusesnum_running. Add an unpinned lockstep-group case that verifies the selected rank receivesAddand peer ranks receive the required wake message.As per coding guidelines, “Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality,” and “Account for dual-dispatch complexity introduced by PD disaggregation in addition to regular routing.”
🤖 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 768 - 818, Expand unpinned routing coverage around unpinned_requests_prefer_the_least_loaded_engine: add scenarios where both engines report scheduler statistics, verifying num_waiting is compared before num_running and num_running breaks ties. Add an unpinned lockstep-group scenario that asserts the selected rank receives Add while peer ranks receive the required wake message, covering both regular and PD dual-dispatch routing behavior.Source: Coding guidelines
bindings/python/src/smg/serve.py (1)
899-903: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add a test for the
_build_router_argsvLLM DP>1 gating.No test in this diff exercises
router_args.zmq_engine_count = engine_dpbeing set only whenengine_dp > 1andself.backend == "vllm". A test constructingServeOrchestrator("vllm", args, ["--data-parallel-size", "2"])and assertingorch._build_router_args().zmq_engine_count == 2, plus a negative case forbackend != "vllm", would close this gap.🤖 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 `@bindings/python/src/smg/serve.py` around lines 899 - 903, Add tests covering _build_router_args: verify a ServeOrchestrator using backend "vllm" with --data-parallel-size 2 sets zmq_engine_count to 2, and verify a non-vLLM backend does not apply this override.model_gateway/src/workflow/job_queue.rs (1)
575-582: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add a focused unit test for the grouped-ZMQ-worker gating in
InitializeWorkersFromConfig.
worker/builder.rstests coverzmq_engine_groupat the builder level, but this call site's own gating —zmq_engine_count > 1combined withConnectionMode::from_url(&spec.url) == Some(ConnectionMode::Zmq)— has no direct test in this diff. A test with a mixedworker_urlslist (oneipc://, onehttp://) confirming only the ZMQ entry getsdp_sizeset would catch a future regression here.🤖 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 `@model_gateway/src/workflow/job_queue.rs` around lines 575 - 582, Add a focused unit test for InitializeWorkersFromConfig covering a mixed worker_urls list with one ipc:// (ZMQ) entry and one http:// entry, using zmq_engine_count greater than 1; assert that only the ZMQ worker receives dp_size and the HTTP worker remains unset, preserving the existing gating logic.model_gateway/src/workflow/steps/local/create_worker.rs (1)
136-154: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win🟡 Nit: Add coverage for the TokenSpeed grouped-ZMQ rejection path.
The existing tests at lines 588-612 cover
validate_zmq_dp(dp-aware expansion), not this new check, which rejectsdp_size > 1+ ZMQ + TokenSpeed at registration. A test constructing aWorkerSpecwithdp_size = Some(2),ConnectionMode::Zmq, andRuntimeType::TokenSpeed, then assertingexecute()returns the expectedStepFailed, would directly confirm this new registration-time guard.🤖 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 `@model_gateway/src/workflow/steps/local/create_worker.rs` around lines 136 - 154, Extend the worker creation tests to cover the registration guard in the ZMQ/TokenSpeed path: construct a WorkerSpec with dp_size set to Some(2), ConnectionMode::Zmq, and RuntimeType::TokenSpeed, then assert execute() returns WorkflowError::StepFailed with the expected rejection message. Keep the existing validate_zmq_dp coverage unchanged.
🤖 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 `@bindings/python/src/lib.rs`:
- Line 959: Move zmq_engine_count to the append-only end position in all
affected definitions: in bindings/python/src/lib.rs, relocate the signature
entry and the matching new(...) parameter after worker_ports_annotation; in
bindings/python/src/smg/router_args.py, relocate the dataclass field past the
“Append new fields here to preserve positional callers” marker near the final
fields. Preserve the existing type and default.
In `@crates/engine_zmq_client/src/connector.rs`:
- Around line 150-160: The select_engine logic must account for requests
dispatched since the last scheduler update instead of relying solely on stale
load values. Update the load selection state around the self.load lookup to
reserve or increment local pending dispatches for the selected engine, and
reconcile those reservations when scheduler statistics arrive; alternatively
ensure equally cold engines are rotated so repeated unpinned selections do not
target the same engine.
---
Nitpick comments:
In `@bindings/python/src/smg/serve.py`:
- Around line 899-903: Add tests covering _build_router_args: verify a
ServeOrchestrator using backend "vllm" with --data-parallel-size 2 sets
zmq_engine_count to 2, and verify a non-vLLM backend does not apply this
override.
In `@crates/engine_zmq_client/src/connector.rs`:
- Around line 768-818: Expand unpinned routing coverage around
unpinned_requests_prefer_the_least_loaded_engine: add scenarios where both
engines report scheduler statistics, verifying num_waiting is compared before
num_running and num_running breaks ties. Add an unpinned lockstep-group scenario
that asserts the selected rank receives Add while peer ranks receive the
required wake message, covering both regular and PD dual-dispatch routing
behavior.
In `@model_gateway/src/workflow/job_queue.rs`:
- Around line 575-582: Add a focused unit test for InitializeWorkersFromConfig
covering a mixed worker_urls list with one ipc:// (ZMQ) entry and one http://
entry, using zmq_engine_count greater than 1; assert that only the ZMQ worker
receives dp_size and the HTTP worker remains unset, preserving the existing
gating logic.
In `@model_gateway/src/workflow/steps/local/create_worker.rs`:
- Around line 136-154: Extend the worker creation tests to cover the
registration guard in the ZMQ/TokenSpeed path: construct a WorkerSpec with
dp_size set to Some(2), ConnectionMode::Zmq, and RuntimeType::TokenSpeed, then
assert execute() returns WorkflowError::StepFailed with the expected rejection
message. Keep the existing validate_zmq_dp coverage unchanged.
🪄 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: 7dee87d9-130b-4970-85e2-bea1f97bf51d
📒 Files selected for processing (13)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pybindings/python/src/smg/serve.pybindings/python/tests/test_serve.pycrates/engine_zmq_client/src/connector.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/grpc/zmq_client.rsmodel_gateway/src/worker/builder.rsmodel_gateway/src/worker/worker.rsmodel_gateway/src/workflow/job_queue.rsmodel_gateway/src/workflow/steps/local/create_worker.rs
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
.github/workflows/pr-test-rust.yml (1)
1118-1144: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win🔴 Important — Gate
finishone2e-2gpu-chat-zmq-dp.Line 1118 does not include the new job in
finish.needs. The failure predicate also does not check it. The aggregate CI result can pass when the grouped ZMQ E2E job fails.Proposed fix
- needs: [..., e2e-2gpu-pd, ...] + needs: [..., e2e-2gpu-pd, e2e-2gpu-chat-zmq-dp, ...] + "${{ needs.e2e-2gpu-chat-zmq-dp.result }}" == "failure" || \🤖 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 @.github/workflows/pr-test-rust.yml around lines 1118 - 1144, Update the finish job’s needs list and its “Check CI result” failure predicate to include e2e-2gpu-chat-zmq-dp, ensuring the aggregate CI result waits for and fails when that job fails.
🤖 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 `@e2e_test/infra/worker.py`:
- Around line 625-628: Reject grouped ZMQ configurations for non-vLLM runtimes
before allocation or launch: in e2e_test/infra/worker.py lines 625-628, validate
get_zmq_engine_count() > 1 against engine and raise ValueError before GPU
sizing; in e2e_test/infra/gateway.py lines 243-247, reject counts above one
unless the selected backend is vllm. Add a regression test covering TokenSpeed
ZMQ with E2E_ZMQ_ENGINE_COUNT=2 and ensure validation fails loudly.
---
Outside diff comments:
In @.github/workflows/pr-test-rust.yml:
- Around line 1118-1144: Update the finish job’s needs list and its “Check CI
result” failure predicate to include e2e-2gpu-chat-zmq-dp, ensuring the
aggregate CI result waits for and fails when that job fails.
🪄 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: ecefb358-e8df-4396-9d01-bce3c93b8b49
📒 Files selected for processing (5)
.github/workflows/e2e-gpu-job.yml.github/workflows/pr-test-rust.ymle2e_test/infra/constants.pye2e_test/infra/gateway.pye2e_test/infra/worker.py
4 tasks
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 203-205: Keep reported EngineLoad values immutable by removing the
optimistic num_waiting update from scheduler state in the registration path.
After all fallible local registration and encoding work succeeds, atomically
select an engine and create a request-ID reservation before sending Add; roll it
back on duplicate, encode, send, or wake failures. Remove each reservation
exactly once on terminal output and both abort paths, including when batch.load
is absent. Add tests covering failed submission rollback and completion without
scheduler statistics, then run the silent-failure-hunter on changed files and
verify the tests.
🪄 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: 6835a2ef-4790-4e59-936a-eef4be9784ad
📒 Files selected for processing (1)
crates/engine_zmq_client/src/connector.rs
| if matches!(error, Error::EngineCoreDead | Error::Transport(_)) { | ||
| warn!(%error, "engine transport failed; failing all in-flight requests"); | ||
| inner.registry.lock().fail_all(Arc::new(error)); | ||
| inner.inflight.lock().clear(); |
There was a problem hiding this comment.
🟡 Nit: The inflight map is cleared here on transport/EngineCoreDead failure, but the two other death-returns in this same function — the silence-timeout watchdog (line ~553, fail_all + return) and the normal output-stream-ended shutdown (line ~563, fail_all) — don't clear it. The stale counts are harmless in practice because fail_all closes the registry so no new submit can land, but matching the cleanup across all three paths would be more consistent.
Member
Author
There was a problem hiding this comment.
Done — all three death returns now clear the map (connector.rs:569, :593, :611). You are right that fail_all makes the stale counts unreachable, but the consistency is worth having: the map is the DP scheduler's only picture of occupancy, and leaving it populated on two of three paths is the kind of asymmetry that becomes a real bug the moment someone adds a recovery path that reopens the registry.
A ZMQ worker can now await a group of DP engines on one socket set: dp_size on the worker spec (with no rank) becomes the engine count the handshake awaits, plumbed spec -> builder -> connect_zmq_backend -> connect_for_worker -> ZmqEngineClient::connect, replacing the hardcoded engine_count=1. Grouped-vs-expanded is deliberate: dp-aware expansion creates one rank-pinned worker per rank, but each ZMQ worker owns its own socket bind, so N expanded workers would fight over the same ipc paths (and unlink each other's live sockets). validate_zmq_dp keeps rejecting the expansion and now says why, pointing at the grouped form. TokenSpeed groups are rejected at registration - its wire has no DP-rank routing yet. The connector's unpinned-request path previously sent everything to engines.first(), which would silently funnel a whole group to rank 0. select_engine now picks the least-loaded engine from the piggybacked per-rank stats (queue depth, then in-flight batch), mirroring vLLM's own frontend DP client; unreported engines score zero so cold ranks fill first, and a single-engine group degenerates to today's behavior. SMG-pinned data_parallel_rank remains authoritative. Launcher (grouped headless launch), a connector selection test on the mock-engine harness, and the 2-GPU e2e lane follow in this PR before review. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
RouterConfig.zmq_engine_count stamps dp_size onto each --worker-urls ZMQ worker at startup registration (mirroring how --backend pins their runtime through startup_worker_runtime_type): the worker becomes a grouped ZMQ worker whose handshake awaits that many DP engines on one socket set. Guarded to ZMQ URLs only - dp_size on an HTTP/gRPC startup worker would misread as dp-awareness. CLI, Python bindings, and the serve launcher's grouped headless vLLM launch ride on this next. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Completes the zmq_engine_count plumbing surface: - gateway CLI: --zmq-engine-count feeds the router knob directly. - Python bindings: zmq_engine_count rides the Router constructor and RouterArgs (generic field mapping picks it up; --zmq-engine-count is also exposed on the args parser). - serve launcher: an engine-level --data-parallel-size N after now launches a grouped headless vLLM (N engines on one socket set, size and size-local both N) instead of being filtered down to 1, and the router args are stamped with the matching zmq_engine_count so the handshake awaits all N engines. The smg-level --data-parallel-size keeps meaning worker replicas; the two compose as replicas of groups. 146 serve tests pass, including new coverage for the grouped launch (space and equals flag forms, launcher flag ownership) and the single-engine default. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Two engines, one reports (waiting=5, running=3) via a pinned warm-up request, the other never reports and scores zero: an unpinned request must land on the idle engine, not engines.first(). Also corrects the module header, which still described the pre-balancing design. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
E2E_ZMQ_ENGINE_COUNT runs a ZMQ lane with grouped workers, riding the lanes #2060 landed: the vLLM worker builder appends the engine-level --data-parallel-size (flowing through the same smg serve launcher as production), the gateway gains --zmq-engine-count so its handshake awaits every engine, and start_workers sizes the worker's GPU slice as tp x engine count. e2e-2gpu-chat-zmq-dp runs the chat suite with dp=2 vLLM groups on the 2-GPU runner, exercising the grouped handshake, the connector's least-loaded selection, and the wave protocol live. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Replace the snapshot-min selection with the algorithm vLLM's own frontend DP client uses (DPLBAsyncMPClient.get_core_engine_for_request, single-client form): - score = max(own in-flight count, reported waiting + running): the exact local floor survives stale snapshots, so a burst between load reports cannot dogpile one rank; - KV-pressure penalty: waiting scaled by 6 x max(0, kv_usage - 0.5) - a queue on a KV-bound engine drains slowly; - optimistic bump of the chosen rank's waiting count until the next real report; - rotating scan start so all-zero ties (cold start) don't always resolve to the first engine. In-flight counts increment at submit and decrement on unique finished ids per batch (terminal outputs and the out-of-band finished list can name the same request), on abort, and on the wake-failure rollback; the map clears with fail_all. New mock-harness test: a cold two-engine burst must spread one request to each engine (the old code sent both to engines.first()); the least-loaded and pinned-rank tests hold. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
…-at-end
Addresses the dp-lane CI failure and the review comments on this PR:
- e2e-2gpu-chat-zmq-dp exited 5 with every test deselected: the tier
filter (hooks.py) requires each test's gpu marker to EQUAL the lane
tier, and chat tests default to gpu=1. The lane now declares
gpu_tier=1 (the tier names the test set and model downloads; the
second GPU serves the second engine of the group).
- selection reservations could leak: the optimistic num_waiting bump
mutated REPORTED load before registration/encode/send/wake could
fail, and nothing rolled it back - an idle vLLM engine emits no
reports, so a single failed submit penalized it permanently. Dropped
the report-cache mutation entirely (vLLM needs it only because
multiple client processes share its engines; this connector is the
group's sole client, so the in-flight floor covers bursts) and moved
the in-flight increment into select_engine under the scoring lock;
every failed submit path now releases via unreserve(), including the
public abort() API, and all three dispatcher death paths clear the
map. New regression test: a duplicate-id submit must not leave a
reservation behind (counts stay (1,0), next request routes to the
idle engine).
- E2E_ZMQ_ENGINE_COUNT>1 with a non-vLLM engine now fails fast in
start_workers instead of reserving GPUs for engines that never
launch.
- zmq_engine_count moved to the documented append point in the pyo3
signature/struct/constructor and the RouterArgs dataclass ('Append
new fields here to preserve positional callers').
- _backend_arg_int logs a warning when an unparseable value falls back
to the default instead of silently discarding operator intent.
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 @.github/workflows/pr-test-rust.yml:
- Line 613: Change the PR-controlled pytest job using the 1-GPU/2-GPU runner
manifests so it runs on an isolated runner without credentials, or remove
HF_TOKEN injection from those PR jobs; do not only remove secrets: inherit,
since the token is supplied by the runner environment.
🪄 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: bf7fdfde-5d78-4fa6-a8d7-b15d185763b4
📒 Files selected for processing (6)
.github/workflows/pr-test-rust.ymlbindings/python/src/lib.rsbindings/python/src/smg/router_args.pybindings/python/src/smg/serve.pycrates/engine_zmq_client/src/connector.rse2e_test/infra/worker.py
🚧 Files skipped from review as they are similar to previous changes (5)
- bindings/python/src/smg/router_args.py
- bindings/python/src/lib.rs
- e2e_test/infra/worker.py
- bindings/python/src/smg/serve.py
- crates/engine_zmq_client/src/connector.rs
The dp=2 lane's first live run caught a stall the mock tests could not: intermittent 'no engine output for 300s with in-flight requests' followed by the watchdog killing the group. Root cause is a race in the wake gate. wake_group skipped the wake whenever the tracked WaveState said the group was running - but that state is only as fresh as the last processed wave_complete. Between the ranks agreeing to park (their all-reduce) and this client's dispatcher processing their completion, a submit saw stale running=true, sent only the Add, and skipped the wake. Without a DP coordinator a parked vLLM engine does NOT self-wake on a received Add (core.py gates that branch on has_coordinator), so the request sat on a parked rank until the 300s silence watchdog failed the group. Sequential traffic - one request in flight, exactly the e2e suite's shape - hits the window constantly. Wake unconditionally instead: vLLM's START_DP_WAVE handler is idempotent (new_wave >= current_wave while stepping is a no-op, stale waves are ignored), so the gate bought one saved message per submit at the price of a stall window. The running flag remains as observed state for the drain bookkeeping. Also: codespell fix (unparsable). The wave tests now pin the unconditional-wake contract, including the racing-submit case. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
## Description ### Problem The dp lane hung on gpt-oss with `no engine output for 300s`, and the engine's own log shows why: `EngineCore_DP0` died on `TimeoutError: RPC call to sample_tokens timed out` after five minutes of `No available shared memory broadcast block found in 60 seconds`. Rank 0 was blocked in the model forward waiting for rank 1 to join the all-reduce; rank 1 was parked and never woke. The wake was dropped on the floor. `START_DP_WAVE` is ignored by an engine whose `current_wave` already exceeds the wave it names (`core.py:_handle_client_request`), and the ranks bump `current_wave` themselves as they park (`core.py:run_busy_loop`). So between the group draining a wave and this client processing the resulting `wave_complete`, the client's clock reads one wave behind the engines, and the next submit's wake is stale — silently discarded. The holder still gets its Add, starts stepping, and blocks forever on peers that were never told to wake. Waking unconditionally (d71f4af) did not help: the message was sent, just ignored. ### Solution Treat the wave number as a logical clock this client owns, and bump it on every wake rather than replaying the last number the engines reported. That is provably never stale: a rank sits at the last number we broadcast plus at most one self-increment (a second increment needs another wave, which needs another wake), so `last + 1` always clears it. Reported waves are still folded in, upward only, so a rank found running ahead (kept its `current_wave` across a gateway restart) also gets cleared. The wake now goes to every rank, with no excluded index. Excluding the request holder left it on a lower `current_wave` than its peers, which breaks the argument above for the price of one saved message. `WaveState` collapses to the bare clock: its `running` flag existed for the wake gate removed in d71f4af. The clock widens to u64 so the monotonic bump cannot wrap a long-lived gateway into the same lost-wake hang. ## Changes - `connector.rs`: `wake_group` bumps the clock and broadcasts to all ranks; `observe_wave` folds reported waves upward only; `WaveState` is now a `Mutex<u64>`. - `protocol/mod.rs`, `vllm/mod.rs`, `tokenspeed/mod.rs`: `encode_start_wave(wave: u64)` drops the exclusion parameter and sends a sentinel index that matches no rank. - `vllm/output.rs`, `mock_engine.rs`: wave fields widened to u64. ## Test Plan - `cargo test -p engine-zmq-client` — 72 passing. `every_wake_names_a_wave_the_ranks_will_accept` replaces `a_drained_wave_re_arms_the_wake`: it asserts each wake strictly outranks the previous one without waiting for any `wave_complete`, and that a rank reporting a wave ahead of the clock is cleared by the next wake. - `cargo clippy --workspace --all-targets -- -D warnings`. - `e2e-2gpu-chat-zmq-dp (vllm)` is the live exercise — this is the lane that hung. ## Checklist - [x] Tests added/updated - [x] Lint and formatting pass - [x] No breaking changes to public APIs Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Comment on lines
+528
to
533
| let mut finished: HashSet<String> = batch | ||
| .finished_request_ids | ||
| .iter() | ||
| .cloned() | ||
| .collect(); | ||
| let mut registry = inner.registry.lock(); |
There was a problem hiding this comment.
🟡 Nit: This /// block runs directly into the existing doc comment above (lines 522-527, which describes the silence-death timeout), so rustdoc merges them into one comment for ENGINE_SEND_TIMEOUT. ENGINE_SILENCE_DEATH_TIMEOUT on line 535 loses its doc comment entirely.
A blank line between the two blocks (or reordering so each constant sits directly below its own doc) would fix the attribution.
Member
Author
There was a problem hiding this comment.
Good catch — rustdoc was attaching that text to ENGINE_SILENCE_DEATH_TIMEOUT, so run_dispatcher documented nothing. Split in 354491fb: the constant keeps the silence-death rationale and the dispatcher gets its own doc.
## Description ### Problem Review catch on the reservation design: a request has several racing ways to finish, and each one decremented the rank's in-flight counter. `Client::abort` released a slot; dropping that same request's unfinished stream sent an auto-abort, and the dispatcher released a slot again; the engine's own terminal output or `finished_requests` report released a third. One request could retire three slots, leaving the rank looking emptier than it is and attracting traffic it cannot serve until the next load report corrects the score. `saturating_sub` kept the counter off the floor, which is what hid this: the count was wrong, never negative. ### Solution Hold in-flight slots as a set of request ids per rank rather than a count. Releasing is then keyed by identity, so the second and third retirement of the same request find nothing to remove and the other requests on that rank keep their slots. The scoring floor reads the set size, so selection behaves as before. Registration now happens before selection. Slots are keyed by id, so a reservation taken before the duplicate-id gate could collide with the live request's own slot and release it on rollback; registration is already the admission gate, so admitting first makes the collision unreachable and drops one rollback path. ## Changes - `connector.rs`: `inflight` is `HashMap<u32, HashSet<String>>`; `unreserve`/`inflight_remove` become `release`/`release_one` (both idempotent); `select_engine` takes the request id and records it; `submit` registers before selecting. ## Test Plan - `cargo test -p engine-zmq-client` — 73 passing. `retiring_one_request_twice_frees_only_its_own_slot` submits two requests to one rank, aborts one and drops its stream (two retirements, one request), and asserts the next unpinned request still routes to the idle rank. Verified it fails when `release` retires an arbitrary slot the way the counter did. - `cargo clippy -p engine-zmq-client --all-targets -- -D warnings`. ## Checklist - [x] Tests added/updated - [x] Lint and formatting pass - [x] No breaking changes to public APIs Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The dispatcher's description had been written above `ENGINE_SILENCE_DEATH_TIMEOUT`, so rustdoc attached it to the constant and `run_dispatcher` documented nothing. Each now sits above its own item. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Problem
DP>1 over the ZMQ direct backend is rejected end-to-end:
validate_zmq_dpfails registration, the client hardcodesengine_count=1, and the serve launcher filters the engine-level--data-parallel-sizedown to 1. The connector's wave protocol (C1, #2078) landed with nothing able to use it. Additionally, the connector sent every unpinned request toengines.first()— any multi-engine group would have funneled all traffic to rank 0.Solution
A ZMQ DP group is one worker, one socket set, N engines (
dp_size: Non the worker spec with no rank — the existing seam, no new field). dp-aware expansion stays rejected: each ZMQ worker owns its socket bind, so N expanded workers would fight over the same ipc paths and unlink each other's live sockets;validate_zmq_dpnow says so and points at the grouped form. Unpinned requests balance inside the group on the piggybacked per-rank scheduler stats — least(num_waiting, num_running)first, mirroring vLLM's own frontend DP client — so the group stays opaque to gateway policies. An SMG-stampeddata_parallel_rankremains authoritative, and TokenSpeed groups are rejected at registration (its wire has no DP-rank routing yet; upstream fix follows).Changes
crates/engine_zmq_client/src/connector.rs:select_enginepicks the least-loaded engine for unpinned requests (unreported engines score zero, so cold ranks fill first; single-engine groups degenerate to prior behavior); module header updated.model_gateway:WorkerMetadata::zmq_engine_count()feeds both ZMQ connect paths, replacing the hardcoded 1;BasicWorkerBuilder::zmq_engine_group;create_workercarriesconfig.dp_sizeonto grouped ZMQ workers and rejects TokenSpeed groups;RouterConfig.zmq_engine_countstamps startup--worker-urlsZMQ workers (thestartup_worker_runtime_typepattern, ZMQ-URLs-only);--zmq-engine-countCLI flag.bindings/python:zmq_engine_counton the Router constructor andRouterArgs(+--zmq-engine-count); serve launcher launches a grouped headless vLLM when an engine-level--data-parallel-size Nis passed after--and stamps the router args to await N engines. The smg-level--data-parallel-sizekeeps meaning worker replicas; the two compose as replicas of groups.Test Plan
Live dp=2: the new
e2e-2gpu-chat-zmq-dp (vllm)lane on this PR runs the chat suite against a real two-engine group on the 2-GPU runner — grouped handshake, least-loaded selection, and wave protocol end-to-end (rides the ZMQ lane scaffolding test(zmq): direct-backend e2e tests and CI lanes #2060 merged).Connector (mock-engine harness): two engines, engine 0 reports
(waiting=5, running=3)via a pinned warm-up, engine 1 never reports — an unpinned request must land on engine 1. Fails on the oldengines.first()behavior. 70 crate tests green.Builder: grouped spec shape (
dp_sizeset, no rank, URL untouched) and the single-engine default.serve: grouped launch coverage (space and equals flag forms, launcher flag ownership, single-engine default) — 146 tests green.
Live
dp=2validation comes with C4's 2-GPU e2e lane, which extends the ZMQ lane scaffolding landing in test(zmq): direct-backend e2e tests and CI lanes #2060 and follows that merge as a small separate PR.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspasses