From 2f3f441f84bd5b35ec8aa9fcfffb540f107da8a7 Mon Sep 17 00:00:00 2001 From: Rishabh Saini Date: Tue, 7 Jul 2026 14:48:23 -0400 Subject: [PATCH] fix: include topic frame in KV events replay response (#45177) Signed-off-by: RishabhSaini --- .../kv_events/kv_events_subscriber.py | 2 +- tests/distributed/conftest.py | 17 +++++-- tests/distributed/test_events.py | 49 +++++++++++++------ vllm/distributed/kv_events.py | 9 ++-- 4 files changed, 50 insertions(+), 27 deletions(-) diff --git a/examples/features/kv_events/kv_events_subscriber.py b/examples/features/kv_events/kv_events_subscriber.py index b8561c73980..cfe131f000d 100644 --- a/examples/features/kv_events/kv_events_subscriber.py +++ b/examples/features/kv_events/kv_events_subscriber.py @@ -99,7 +99,7 @@ def main(): replay.send((last_seq + 1).to_bytes(8, "big")) while poller.poll(timeout=200): - seq_bytes, replay_payload = replay.recv_multipart() + _, seq_bytes, replay_payload = replay.recv_multipart() if not replay_payload: # End of replay marker is sent as an empty frame # for the payload diff --git a/tests/distributed/conftest.py b/tests/distributed/conftest.py index da661c5e13b..48df856f5e7 100644 --- a/tests/distributed/conftest.py +++ b/tests/distributed/conftest.py @@ -97,11 +97,12 @@ class MockSubscriber: for endpoint in pub_endpoints: self.sub.connect(endpoint) - # Set up replay sockets if provided + # Set up replay sockets if provided. + # DEALER allows receiving multiple replies per request. self.replay_sockets = [] if replay_endpoints: for replay_endpoint in replay_endpoints: - replay = self.ctx.socket(zmq.REQ) + replay = self.ctx.socket(zmq.DEALER) replay.connect(replay_endpoint) self.replay_sockets.append(replay) @@ -132,7 +133,9 @@ class MockSubscriber: if socket_idx >= len(self.replay_sockets): raise ValueError(f"Invalid socket index {socket_idx}") - self.replay_sockets[socket_idx].send(start_seq.to_bytes(8, "big")) + self.replay_sockets[socket_idx].send_multipart( + [b"", start_seq.to_bytes(8, "big")] + ) def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]: """Receive replayed messages from a specific replay socket""" @@ -148,12 +151,16 @@ class MockSubscriber: if not replay_socket.poll(1000): break + # DEALER receives [empty_delim, topic, seq, payload] frames = replay_socket.recv_multipart() - if not frames or not frames[-1]: + if frames and frames[0] == b"": + frames = frames[1:] + if len(frames) != 3 or not frames[-1]: # End of replay marker break - seq_bytes, payload = frames + topic, seq_bytes, payload = frames + assert topic == self.topic_bytes seq = int.from_bytes(seq_bytes, "big") data = self.decoder.decode(payload) replayed.append((seq, data)) diff --git a/tests/distributed/test_events.py b/tests/distributed/test_events.py index f17b7997c58..9b5601ad1d9 100644 --- a/tests/distributed/test_events.py +++ b/tests/distributed/test_events.py @@ -80,20 +80,38 @@ def test_replay_mechanism(publisher, subscriber): batch = create_test_events(1) publisher.publish(batch) - time.sleep(0.5) # Need publisher to process above requests - subscriber.request_replay(10) + # Drain live events to ensure publisher has buffered them. + for _ in range(19): + assert subscriber.receive_one(timeout=1000) is not None - batch = create_test_events(1) - publisher.publish(batch) # 20th message + subscriber.request_replay(10) replayed = subscriber.receive_replay() - assert len(replayed) > 0, "No replayed messages received" - seqs = [seq for seq, _ in replayed] - assert all(seq >= 10 for seq in seqs), "Replayed messages not in order" - assert seqs == list(range(min(seqs), max(seqs) + 1)), ( - "Replayed messages not consecutive" + assert len(replayed) == 9, ( + f"Expected 9 replayed messages (seq 10-18), got {len(replayed)}" ) + seqs = [seq for seq, _ in replayed] + assert seqs == list(range(10, 19)), "Replayed sequences should be 10-18" + + +def test_replay_includes_topic(publisher, subscriber, publisher_config): + """Test that replay responses include the topic, matching PUB format""" + for _ in range(5): + publisher.publish(create_test_events(1)) + + # Drain live events to ensure publisher has processed them. + for _ in range(5): + assert subscriber.receive_one(timeout=1000) is not None + + subscriber.request_replay(0) + + # receive_replay unpacks (topic, seq, payload) and asserts + # topic == publisher topic for each message. + replayed = subscriber.receive_replay() + assert len(replayed) == 5, f"Expected 5 replayed messages, got {len(replayed)}" + seqs = [seq for seq, _ in replayed] + assert seqs == list(range(5)), "Replayed sequences should be 0-4" def test_buffer_limit(publisher, subscriber, publisher_config): @@ -108,15 +126,16 @@ def test_buffer_limit(publisher, subscriber, publisher_config): time.sleep(0.5) # Need publisher to process above requests subscriber.request_replay(0) - batch = create_test_events(1) - publisher.publish(batch) - replayed = subscriber.receive_replay() - assert len(replayed) <= buffer_size, "Can't replay more than buffer size" + assert len(replayed) == buffer_size, ( + f"Expected {buffer_size} replayed messages, got {len(replayed)}" + ) - oldest_seq = min(seq for seq, _ in replayed) - assert oldest_seq >= 10, "The oldest sequence should be at least 10" + seqs = [seq for seq, _ in replayed] + assert seqs == list(range(10, buffer_size + 10)), ( + "Should replay seq 11 through buffer_size+10" + ) def test_topic_filtering(publisher_config): diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index 2f363f0bc06..a7d83bb378f 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -459,15 +459,12 @@ class ZmqEventPublisher(EventPublisher): for seq, buf in self._buffer: if seq >= start_seq: - # [identity, empty_delim, seq_bytes, payload] - # (identity, empty_delim) are stripped off by the router - # receiving payload is (seq_bytes, payload) + # Subscriber receives (topic, seq_bytes, payload) self._replay.send_multipart( - (client_id, b"", seq.to_bytes(8, "big"), buf) + (client_id, b"", self._topic_bytes, seq.to_bytes(8, "big"), buf) ) # Send end of sequence marker - # receiving payload is (-1, b""") - self._replay.send_multipart((client_id, b"", self.END_SEQ, b"")) + self._replay.send_multipart((client_id, b"", b"", self.END_SEQ, b"")) @staticmethod def offset_endpoint_port(