fix(scheduler): wait for rank subscribers before first broadcast - #495
Open
Cyber-Marty wants to merge 1 commit into
Open
Cyber-Marty wants to merge 1 commit into
Cyber-Marty wants to merge 1 commit into
Conversation
JUNQINGV587
added a commit
to JUNQINGV587/FreeToken
that referenced
this pull request
Sep 23, 2026
…broadcast The TP control channel broadcasts the first message right after startup, but a SUB's subscription only reaches the publisher asynchronously; anything broadcast before it lands is silently dropped and both ranks block forever. Switches PUB to XPUB and waits for size-1 subscription events. This fork runs TP=2 in production, so the race is live.
This branch has not been deployed
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.
Fixes #364.
Root cause
A ZeroMQ slow-joiner race with a deadlock twist. In the TP>1 path, the SUB side's
connect+SUBSCRIBEcomplete asynchronously, and nothing in the current setup waits for the registration to land on the publisher. If rank 0 picks up the first request and broadcasts before rank 1's subscription is registered on the PUB socket, ZMQ silently drops the message (PUB has no queue for unregistered peers). Rank 0 then waits for rank 1's reply, rank 1 waits for the next broadcast, and both ranks block forever while/v1/modelskeeps answering.Fix
Switch the rank-0 broadcast socket from PUB to XPUB, which sends identically to PUB but surfaces inbound subscription events (first byte
0x01). The scheduler's broadcast path now callswait_for_subscribers(tp_info.size - 1)before entering the receive loop, with a 30s timeout so a rank that died before connecting fails fast with a clear error instead of hanging the publisher silently.Tests
tests/utils/test_mp_pubsub.py, 4 cases over real loopback TCP sockets:wait_for_subscribers, payload decoded equalwait_for_subscribers(0)returns immediately (TP=1 behavior unchanged)TimeoutErrorwithout subscribersput_rawis identical (XPUB does not change the wire format)Red/green verified: all 4 fail on unpatched
main(the new API does not exist), all 4 pass after the fix.Tested on
Windows 11, CPU only, no multi-GPU box - pytest over real loopback TCP sockets (
PYTHONPATH=python .venv/Scripts/python.exe -m pytest tests/utils/test_mp_pubsub.py). I could not run the full TP=2 serve path end to end; the race window itself is covered by the socket-level tests rather than an end-to-end run.