Skip to content

fix(utils): tolerate coalesced msgpack frames in zmq pull queues - #466

Open
Cyber-Marty wants to merge 1 commit into
FlashML-org:mainfrom
Cyber-Marty:fix/mp-unpacker-coalesced-frames
Open

Cyber-Marty wants to merge 1 commit into
FlashML-org:mainfrom
Cyber-Marty:fix/mp-unpacker-coalesced-frames

Conversation

@Cyber-Marty

Copy link
Copy Markdown

Summary

ZmqPullQueue.get(), ZmqSubQueue.get() and ZmqAsyncPullQueue.get() decoded each received frame with msgpack.unpackb(frame, raw=False). When a frame carries more than one packed object, unpackb raises ExtraData — and both crash sites in #452 (freetoken/tokenizer/server.py:160 tokenize worker, freetoken/scheduler/io.py:85 scheduler) die on exactly that helper (utils/mp.py:68).

The reporter's hypothesis matches the failure profile: a coalesced read (two packed objects arriving in one frame) is non-deterministic, not tied to input, and the same helper serves both the tokenizer and scheduler receive paths. 2 crashes in 81 sequential requests.

Fix

Route every received frame through one buffered msgpack.Unpacker per queue (_CoalescedUnpacker):

  • get() returns the first object from the buffer; a coalesced frame's remainder is held for the next call instead of crashing;
  • empty() accounts for buffered objects (a socket-only check would report "empty" with undelivered messages pending);
  • get_raw() / decode() stay unbuffered — the multi-rank broadcast path (_recv_msg_multi_rank0) pairs empty()/get_raw() with a rank-wide count, which a buffered remainder would desynchronize. Mixing them with buffered get() now raises a clear RuntimeError instead of silently skipping messages;
  • the scheduler itself is unchanged — this only fixes the decode side.

Tests

New tests/utils/test_mp_coalesced.py (4 passed, 1 skipped):

  • sync pull: a frame with two packed objects yields two get() results in order;
  • sub queue: same through PUB/SUB;
  • buffered-unpacker unit: feed once, take twice;
  • empty() accounting with a buffered remainder;
  • async variant skipped on Windows (zmq.asyncio needs a Selector loop; prod runs on Linux CI).

Red-verified: msgpack.unpackb on a coalesced frame raises ExtraData against the previous implementation.

Not in scope

get_raw()/decode() are deliberately left unbuffered with a mixing guard — the multi-rank broadcast path counts raw frames across ranks, and a buffered remainder would break that count. If coalescing is ever observed on that path, it needs a design pass (length-prefix framing) rather than a local patch.

Fixes #452

ZmqPullQueue/ZmqSubQueue/ZmqAsyncPullQueue decoded each received frame
with msgpack.unpackb, which raises ExtraData when a frame carries more
than one packed object. On Windows with offload MoE this surfaced in
production as two mid-batch worker deaths in 81 requests: the tokenizer
and scheduler processes crashed on the same helper
(utils/mp.py), taking the whole API server down.

Route every received frame through one buffered msgpack.Unpacker per
queue: the first object is returned by get() and any remainder is held
for subsequent calls instead of crashing. empty() accounts for the
buffer, and get_raw()/decode() stay unbuffered (the multi-rank
broadcast path pairs them with a rank-wide count that a buffered
remainder would desynchronize) -- mixing them with buffered get() now
raises instead of silently skipping messages.

Regression tests cover the coalesced frame on sync pull, sub queue,
decode buffer, and empty() accounting; the async variant is skipped on
Windows (zmq.asyncio needs a Selector loop; prod runs on Linux CI).

Fixes FlashML-org#452
gdevenyi added a commit to gdevenyi/FreeToken that referenced this pull request Sep 18, 2026
gdevenyi added a commit to gdevenyi/FreeToken that referenced this pull request Sep 19, 2026
JUNQINGV587 added a commit to JUNQINGV587/FreeToken that referenced this pull request Sep 23, 2026
The bare `async def` test from FlashML-org#466 needs pytest-asyncio plus `asyncio_mode = auto`;
this repo's dev extra declares only pytest, so it failed collection instead of
exercising the queue. Same assertions, driven through asyncio.run.
JUNQINGV587 added a commit to JUNQINGV587/FreeToken that referenced this pull request Sep 23, 2026
Adds the A/B numbers behind declining upstream FlashML-org#385's expert sharding (cold 262K prefill
3368 vs 2840 tok/s, six concurrent 8K prompts 3.68 vs 5.13 s median TTFT, functionally equal,
and the structural reason: EP partitions the expert set so a rank caches only what it can
compute, while sharding replicates the whole routed working set on every rank), and the sync
log for FlashML-org#85 and FlashML-org#466 with the survey result that 21 community picks were already carried.
gdevenyi added a commit to gdevenyi/FreeToken that referenced this pull request Sep 23, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Engine worker dies mid-batch: msgpack ExtraData in utils/mp.py IPC receive (2 crashes / 81 requests, Windows, offload MoE)

1 participant