diff --git a/python/freetoken/scheduler/io.py b/python/freetoken/scheduler/io.py
index 37557a9..ca67b01 100644
--- a/python/freetoken/scheduler/io.py
+++ b/python/freetoken/scheduler/io.py
@@ -1,5 +1,7 @@
from __future__ import annotations
+import os
+import time
from typing import TYPE_CHECKING, Final, List
import torch
@@ -64,6 +66,77 @@ class SchedulerIOMixin:
self.receive_msg = recv
self.send_result = send
+ if tp_info.size > 1:
+ self._handshake_rank_relay(tp_info)
+
+ # ------------------------------------------------------------------------------------------
+ # The rank-0 -> rank-N request relay is ZeroMQ PUB/SUB. A PUB socket DROPS every message for
+ # which no subscription is registered yet, and rank N's connect + SUBSCRIBE travel through
+ # ZeroMQ's I/O thread asynchronously -- nothing above waits for them to land. Without a
+ # handshake the first request published after readiness can be lost: rank 0 then blocks in
+ # the gloo broadcast of `_recv_msg_multi_rank0` waiting for a rank N that blocks in the SUB
+ # receive of `_recv_msg_multi_rank1`, no timeout, no error, while the frontend keeps
+ # answering GETs (llm-server #795, gates H4/2b: wedged at ready+1 s, served at ready+60 s).
+ #
+ # The handshake: rank 0 publishes a hello frame and every rank reports, through the gloo
+ # group (the same `broadcast` primitive the relay already uses), whether it has received one;
+ # rank 0 repeats until every subscriber has. Then rank 0 publishes ONE done frame and each
+ # subscriber drains hello frames until it sees it -- ZeroMQ preserves order on a connection
+ # whose subscription is registered, so the done frame is the last handshake frame that can
+ # arrive, and nothing of the handshake can be mistaken for a request later. Only after that
+ # does the caller reach `sync_all_ranks()` and the "Scheduler is ready" ack.
+ # ------------------------------------------------------------------------------------------
+ _RELAY_HELLO: Final = b"\x00freetoken-relay-hello"
+ _RELAY_DONE: Final = b"\x00freetoken-relay-done"
+ _RELAY_POLL_MS: Final = 50
+ _RELAY_MAX_HELLOS: Final = int(os.environ.get("FREETOKEN_RELAY_HANDSHAKE_MAX_HELLOS", "2400"))
+
+ def _handshake_rank_relay(self, tp_info) -> None:
+ t0 = time.monotonic()
+ size = tp_info.size
+ if tp_info.is_primary():
+ pub = self._send_into_ranks.socket
+ hellos = 0
+ while True:
+ hellos += 1
+ pub.send(self._RELAY_HELLO)
+ if self._relay_subscribers_seen(tp_info, seen=False) == size - 1:
+ break
+ if hellos >= self._RELAY_MAX_HELLOS:
+ raise RuntimeError(
+ f"TP relay handshake: no subscriber acknowledged after {hellos} hellos "
+ f"({time.monotonic() - t0:.1f} s); the rank-0 PUB never reached rank>=1's SUB"
+ )
+ pub.send(self._RELAY_DONE)
+ logger.info(
+ f"TP relay handshake: {size - 1} subscriber(s) joined after {hellos} hello(s) "
+ f"in {(time.monotonic() - t0) * 1000:.0f} ms"
+ )
+ else:
+ sub = self._recv_from_rank0.socket
+ seen = False
+ while True:
+ if not seen and sub.poll(timeout=self._RELAY_POLL_MS):
+ frame = sub.recv()
+ assert frame == self._RELAY_HELLO, f"unexpected relay frame during handshake: {frame!r}"
+ seen = True
+ if self._relay_subscribers_seen(tp_info, seen=seen) == size - 1:
+ break
+ while True:
+ frame = sub.recv()
+ if frame == self._RELAY_DONE:
+ break
+ assert frame == self._RELAY_HELLO, f"unexpected relay frame during handshake: {frame!r}"
+
+ def _relay_subscribers_seen(self, tp_info, seen: bool) -> int:
+ """One round of the handshake: every rank >= 1 broadcasts whether it has received a hello."""
+ total = 0
+ for root in range(1, tp_info.size):
+ flag = torch.tensor(int(seen) if tp_info.rank == root else 0)
+ self.tp_cpu_group.broadcast(flag, root=root).wait()
+ total += int(flag.item())
+ return total
+
def run_when_idle(self):
raise NotImplementedError("should be implemented")
diff --git a/tests/scheduler/test_io_relay_handshake.py b/tests/scheduler/test_io_relay_handshake.py
new file mode 100644
index 0000000..bad2676
--- /dev/null
+++ b/tests/scheduler/test_io_relay_handshake.py
@@ -0,0 +1,125 @@
+"""The rank-0 -> rank-N request relay must not lose the first request (llm-server #795, H4/2b).
+
+Two real processes, a real gloo group, real ZeroMQ ipc sockets -- the served shape of
+`SchedulerIOMixin.__init__` at TP=2, minus the engine. The failing case is deterministic when rank 1
+builds its SUB after rank 0 has already published: a PUB drops what nobody subscribes to.
+"""
+from __future__ import annotations
+
+import socket
+import time
+from types import SimpleNamespace
+
+import pytest
+import torch.multiprocessing as mp
+
+PAYLOAD = b"first-request-bytes"
+JOIN_TIMEOUT_S = 60
+
+
+def _config(tmp: str, rank: int, size: int) -> SimpleNamespace:
+ from freetoken.distributed.info import DistributedInfo
+
+ return SimpleNamespace(
+ tp_info=DistributedInfo(rank=rank, size=size),
+ offline_mode=False,
+ zmq_backend_addr=f"ipc://{tmp}/backend",
+ zmq_detokenizer_addr=f"ipc://{tmp}/detok",
+ backend_create_detokenizer_link=True,
+ zmq_scheduler_broadcast_addr=f"ipc://{tmp}/broadcast",
+ )
+
+
+def _close(io) -> None:
+ """The engine's shutdown does this through each queue's stop(); a context with an open socket
+ blocks its term() at interpreter exit, so the test closes explicitly and with linger 0."""
+ for name in ("_recv_from_tokenizer", "_send_into_tokenizer", "_send_into_ranks", "_recv_from_rank0"):
+ q = getattr(io, name, None)
+ if q is not None:
+ q.socket.close(linger=0)
+ q.context.term()
+
+
+def _rank_main(rank: int, size: int, tmp: str, port: int, delay: dict, handshake: bool, result):
+ import faulthandler
+ import sys
+
+ import torch.distributed as dist
+
+ from freetoken.scheduler.io import SchedulerIOMixin
+
+ faulthandler.dump_traceback_later(30, exit=True, file=sys.stderr)
+ dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=size)
+ group = dist.group.WORLD
+ io = None
+ try:
+ time.sleep(delay.get(rank, 0.0))
+ io = SchedulerIOMixin.__new__(SchedulerIOMixin)
+ if not handshake:
+ io._handshake_rank_relay = lambda tp_info: None # the pre-fix engine
+ SchedulerIOMixin.__init__(io, _config(tmp, rank, size), group)
+ if handshake:
+ io.sync_all_ranks() # launch.py: the barrier that precedes "Scheduler is ready"
+ if rank == 0:
+ io._send_into_ranks.put_raw(PAYLOAD) # the first request, relayed at ready+0 s
+ if not handshake:
+ io.sync_all_ranks()
+ result.put((rank, "sent"))
+ else:
+ sub = io._recv_from_rank0.socket
+ if not handshake:
+ io.sync_all_ranks()
+ got = sub.recv() if sub.poll(timeout=5000) else None
+ result.put((rank, got))
+ finally:
+ result.close()
+ result.join_thread()
+ dist.barrier()
+ dist.destroy_process_group()
+ if io is not None:
+ _close(io)
+
+
+def _run(tmp_path, delay: dict, handshake: bool):
+ ctx = mp.get_context("spawn")
+ result = ctx.Queue()
+ with socket.socket() as s:
+ s.bind(("127.0.0.1", 0))
+ port = s.getsockname()[1]
+ size = 2
+ procs = [
+ ctx.Process(target=_rank_main, args=(r, size, str(tmp_path), port, delay, handshake, result))
+ for r in range(size)
+ ]
+ for p in procs:
+ p.start()
+ for p in procs:
+ p.join(JOIN_TIMEOUT_S)
+ alive = [p.pid for p in procs if p.is_alive()]
+ for p in procs:
+ if p.is_alive():
+ p.kill()
+ assert not alive, f"ranks still running after {JOIN_TIMEOUT_S} s (a wedge): {alive}"
+ assert all(p.exitcode == 0 for p in procs), [p.exitcode for p in procs]
+ out = {}
+ while not result.empty():
+ r, v = result.get()
+ out[r] = v
+ return out
+
+
+@pytest.mark.parametrize("delay", [{}, {1: 1.0}, {0: 1.0}], ids=["together", "rank1-late", "rank0-late"])
+def test_first_request_reaches_rank1_with_handshake(tmp_path, delay):
+ out = _run(tmp_path, delay, handshake=True)
+ assert out[0] == "sent"
+ assert out[1] == PAYLOAD, out[1]
+
+
+def test_without_handshake_a_late_subscriber_loses_the_first_request(tmp_path):
+ """The mechanism the handshake exists for, made deterministic: rank 0 publishes before rank 1's
+ SUB exists (rank 1 is 1 s late), so the PUB has no subscription to match and drops the frame.
+ In the served engine the same loss happens on a subscription that is created but not yet
+ registered at the PUB (llm-server #795: 5 hellos / 203 ms on ipc with both ranks together)."""
+ out = _run(tmp_path, {1: 1.0}, handshake=False)
+ assert out[0] == "sent"
+ assert out[1] is None, "the pre-fix relay delivered the first request; the race did not reproduce"
Seen on
qwen4_expat--tp-size 2 --disable-pynccl(ROCm gfx1201, but nothing below is platform-specific), main@4b94bdc3 plus our local patches, pyzmq 27.2.0 / libzmq 4.3.5.What happens
SchedulerIOMixin.__init__(scheduler/io.py) builds the inter-rank relay at the end ofScheduler.__init__: rank 0 binds azmq.PUBonzmq_scheduler_broadcast_addr, each rank ≥ 1 connects azmq.SUBand subscribes to""(utils/mp.py,ZmqPubQueue/ZmqSubQueue).launch.pythen runssync_all_ranks()(a gloo barrier) and puts"Scheduler is ready". Nothing waits for the SUBSCRIBE to reach the PUB — connect and subscribe travel through ZeroMQ's I/O thread asynchronously, and a PUB socket drops every frame no registered subscription matches (the slow-joiner case in the zguide).For the first request,
_recv_msg_multi_rank0doesget_raw()→_send_into_ranks.put_raw(raw)→tp_cpu_group.broadcast(len(pending)).wait(). Rank ≥ 1 in_recv_msg_multi_rank1sits in_recv_from_rank0.get()before its matching broadcast. If the publish was dropped, rank 0 waits in the broadcast for a rank that waits in the SUB receive. No timeout, no handshake, no error. Every later request queues behind it in rank 0's PULL, and the frontend keeps servingGET /v1/modelsin 0.5 s, so the server looks healthy.Evidence
Prefill batchever. py-spy: rank 0 atio.py:102(the broadcast), rank 1 atio.py:113→mp.py:143(socket.recv()). Rank 0's listener shows exactly one accepted connection (rank 1's SUB) — the topology was right, the frame was gone.TP relay handshake: 1 subscriber(s) joined after 2 hello(s) in 54 ms— the first hello was dropped, the second reached the subscriber — and the first request, sent 0.1 s afterAPI server is ready to serve, was served (first prefill 35 s later, 63 tokens back).Why it matters
Any client that fires at readiness — a health probe followed by the first user request, which is what every model router does — can wedge a TP>1 server. On a long cold start that is the whole load lost, with no error line.
Fix (proposed; full diff inline at the bottom)
A handshake at IO init, in
SchedulerIOMixin.__init__, whentp_info.size > 1: rank 0 publishes a hello frame and every rank ≥ 1 reports through the gloo group (thebroadcastprimitive the relay already uses) whether it has received one; rank 0 repeats until all have. Rank 0 then publishes one done frame and each subscriber drains hello frames until it sees it — ZeroMQ preserves order on a connection whose subscription is registered, so the done frame is the last handshake frame that can arrive and nothing can be mistaken for a request later. Only then does the caller reachsync_all_ranks()and the ready ack. TP=1 never enters it. A cap on hello rounds turns "the SUB never subscribed" into aRuntimeErrorinstead of a hang.Alternatives considered: waiting on the SUB monitor for
EVENT_CONNECTEDis not sufficient on its own (the SUBSCRIBE frame follows the connect); replacing PUB/SUB with one PUSH/PULL per rank removes the subscription state entirely at the cost of a socket per rank on rank 0. Independently of the fix, a timeout on the gloo broadcast in_recv_msg_multi_rank*would surface a lost message as an error rather than a hang.Known limitations of the diff below
Stated here so a maintainer does not have to find them. They are in our build as-is, because our image
build asserts the post-patch tree hash and a change would de-link the gates we ran; we would happily fold
any of these in before it lands upstream.
ranks ≥ 1 wait unbounded — a
while Truearound the collective, then a blockingrecv()for the doneframe. If rank 0 raises, the other ranks hang, which is the failure mode the patch removes on the happy
path. A deadline on the rank ≥ 1 side (and a broadcast of rank 0's failure) would close it.
asserts sit on a non-test path (the unexpected-frame checks), so they vanish underpython -O.They should be a raise.
_relay_subscribers_seen(tp_info, seen=False)is called from rank 0 with a flag that is never read —only roots ≥ 1 broadcast. Harmless, but it reads as meaningful.
more predictable.
Reproduce
Start a TP=2 server and POST a completion the moment
API server is ready to serveappears; compare with a 60 s wait. Or runtests/scheduler/test_io_relay_handshake.pyfrom the diff below: the negative test (handshake=False, rank 1 one second late) loses the frame every time.The diff
Against
main@4b94bdc3. Two files:scheduler/io.pyand a new two-process test. TP=1 never enters thehandshake. Happy to open it as a PR instead if you would rather review it that way — and happy to fold in
the four limitations listed above first.
0011-tp-relay-handshake.patch(221 lines)