feat(zmq): piggyback a scheduler-load snapshot on slim output batches - #1079
feat(zmq): piggyback a scheduler-load snapshot on slim output batches#1079slin1237 wants to merge 1 commit into
Conversation
An external frontend balancing a DP group needs per-rank load, and the msgpack wire has no channel for it: the pickle-mode GetLoad poll's reply is a control message this transport drops by design, so today a frontend routes on its own in-flight counts alone — blind to queue depth and KV pressure that another client's traffic creates. Append a load snapshot to BatchTokenIDOutSlim, sampled at send time: num_running, num_waiting, kv_active_pages, kv_total_pages — the same sources as the per-iteration Prometheus snapshot (resident request states, scheduler waiting_size, active/usable KV pages), with the KV ratio carried as its exact integer numerator/denominator. Appended tail fields defaulting to 0 keep the wire append-only compatible in both directions; kv_total_pages == 0 tells the frontend "no snapshot". The sampler late-binds onto MsgpackSendSocket after the scheduler exists (the socket is created during the startup handshake, before it), mirroring how engine_index is injected, and is sampled fresh per batch. Known limit, deliberate: outputs only flow when tokens flow, so an idle rank's last snapshot goes stale (and predates the final commit). The frontend routed every request, so its own empty in-flight set for a rank is ground truth for quiescence; correcting for staleness is the frontend's job, not a heartbeat's. Tests: snapshot round-trip; 10-element (pre-load era) decode keeps zero defaults; the send socket samples load_fn fresh per batch; the wire- order pin extended to the 14-element form. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 534e3aa9e1
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| num_running: int = 0 | ||
| num_waiting: int = 0 | ||
| kv_active_pages: int = 0 | ||
| kv_total_pages: int = 0 |
There was a problem hiding this comment.
In msgpack mode this struct is encoded as an array_like positional tuple, and MsgpackSendSocket now sends the new 14-field shape for every BatchTokenIDOutSlim, including when the load values are just defaults. That only lets a new decoder accept older 10-field senders; any SMG/frontend still decoding the previous 10-slot tuple will reject or misread every output batch during a mixed-version rollout, so generation streams break rather than staying compatible. Please gate the extra tail on a wire version/handshake or otherwise preserve the old arity until the peer has advertised support.
Useful? React with 👍 / 👎.
| self.send_to_tokenizer.load_fn = lambda: ( | ||
| len(self.output_processor.rid_to_state), | ||
| self.scheduler.waiting_size(), | ||
| self.scheduler.active_kv_pages(), |
There was a problem hiding this comment.
For terminal batches, send_pyobj() is called from stream_output() before the FinishEvent/AbortEvent returned by post_process_forward_op() is applied with advance_forward(), so scheduler.active_kv_pages() still includes pages for requests that are being finished in that same batch. Because the msgpack path only publishes load snapshots on output batches, a rank that just became idle can leave the frontend with a stale nonzero KV ratio until some later output, which can make least-cache routing avoid capacity that was already freed; send a post-advance snapshot or subtract the pending terminal changes from this sample.
Useful? React with 👍 / 👎.
1 participant
Problem
An external frontend balancing a DP group over the msgpack ZMQ wire has no per-rank load signal at all. The pickle-mode mechanism does not carry over:
AsyncLLM.watch_load_threadpollsGetLoadReqInput, but itsGetLoadReqOutputreply is a control messageMsgpackSendSocket.send_pyobjdrops by design (zmq_msgpack.py). The frontend is left routing on its own in-flight counts alone — blind to queue depth and KV pressure.Solution
Piggyback a scheduler-load snapshot on every
BatchTokenIDOutSlim, sampled at send time:num_runninglen(output_processor.rid_to_state)num_waitingscheduler.waiting_size()kv_active_pagesscheduler.active_kv_pages()kv_total_pagesscheduler_cache_geometry.num_usable_pagesThese are the same sources as the per-iteration Prometheus snapshot (
_record_scheduler_iteration_metrics), with the KV ratio carried as its exact integer numerator/denominator rather than a float. No extra scheduler wakeups and no control-channel round trip — four integers ride a send that already happens.Wire compatibility: appended tail fields defaulting to 0, the same append-only discipline as
engine_index(#1046). Older decoders skip them; older senders leave them 0, andkv_total_pages == 0is the frontend's explicit "no snapshot" signal, so it cannot mistake an old sender for an empty scheduler.Injection: the sampler late-binds onto
MsgpackSendSocketafter the scheduler exists (the socket is created during the startup handshake, before it), mirroring howengine_indexis injected, and is sampled fresh per batch.Known limit, deliberate: outputs only flow when tokens flow, so an idle rank's last snapshot goes stale (and is sampled before its own final finish commits). The frontend routed every request, so its empty in-flight set for a rank is ground truth for quiescence — correcting staleness is frontend-side by design, not a heartbeat. A future shared-memory load mailbox (scheduler publishes a snapshot every iteration, frontend reads on demand) can supersede this for the same-host case; the piggyback remains the transport-agnostic floor.
Tests
test/runtime/test_zmq_msgpack.py(all pass):test_slim_out_piggybacks_the_load_snapshot— snapshot round-trips.test_slim_out_load_defaults_for_older_senders— a 10-element (engine_index-era) array decodes with the zero "no snapshot" defaults.test_send_socket_samples_load_fn_per_batch— the socket samplesload_fnfresh on every batch.test_slim_out_is_tagged_positional_tuple— the wire-order pin extended to the 14-element form with the tail at positions 10-13.