fix: include topic frame in KV events replay response (#45177)
Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>
Rishabh Saini committed
Jul 7, 2026 at 14:48 UTC
2f3f441f84bd5b35ec8aa9fcfffb540f107da8a7
4 files changed
+50
-27
examples/features/kv_events/kv_events_subscriber.py
+1
-1
@@ -99,7 +99,7 @@ def main():
99
replay.send((last_seq + 1).to_bytes(8, "big"))
100
101
while poller.poll(timeout=200):
102
- seq_bytes, replay_payload = replay.recv_multipart()
102
+ _, seq_bytes, replay_payload = replay.recv_multipart()
103
if not replay_payload:
104
# End of replay marker is sent as an empty frame
105
# for the payload
tests/distributed/conftest.py
+12
-5
@@ -97,11 +97,12 @@ class MockSubscriber:
97
for endpoint in pub_endpoints:
98
self.sub.connect(endpoint)
99
100
- # Set up replay sockets if provided
100
+ # Set up replay sockets if provided.
101
+ # DEALER allows receiving multiple replies per request.
102
self.replay_sockets = []
103
if replay_endpoints:
104
for replay_endpoint in replay_endpoints:
104
- replay = self.ctx.socket(zmq.REQ)
105
+ replay = self.ctx.socket(zmq.DEALER)
106
replay.connect(replay_endpoint)
107
self.replay_sockets.append(replay)
108
@@ -132,7 +133,9 @@ class MockSubscriber:
133
if socket_idx >= len(self.replay_sockets):
134
raise ValueError(f"Invalid socket index {socket_idx}")
135
135
- self.replay_sockets[socket_idx].send(start_seq.to_bytes(8, "big"))
136
+ self.replay_sockets[socket_idx].send_multipart(
137
+ [b"", start_seq.to_bytes(8, "big")]
138
+ )
139
140
def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]:
141
"""Receive replayed messages from a specific replay socket"""
@@ -148,12 +151,16 @@ class MockSubscriber:
151
if not replay_socket.poll(1000):
152
break
153
154
+ # DEALER receives [empty_delim, topic, seq, payload]
155
frames = replay_socket.recv_multipart()
152
- if not frames or not frames[-1]:
156
+ if frames and frames[0] == b"":
157
+ frames = frames[1:]
158
+ if len(frames) != 3 or not frames[-1]:
159
# End of replay marker
160
break
161
156
- seq_bytes, payload = frames
162
+ topic, seq_bytes, payload = frames
163
+ assert topic == self.topic_bytes
164
seq = int.from_bytes(seq_bytes, "big")
165
data = self.decoder.decode(payload)
166
replayed.append((seq, data))
tests/distributed/test_events.py
+34
-15
@@ -80,20 +80,38 @@ def test_replay_mechanism(publisher, subscriber):
80
batch = create_test_events(1)
81
publisher.publish(batch)
82
83
- time.sleep(0.5) # Need publisher to process above requests
84
- subscriber.request_replay(10)
83
+ # Drain live events to ensure publisher has buffered them.
84
+ for _ in range(19):
85
+ assert subscriber.receive_one(timeout=1000) is not None
86
86
- batch = create_test_events(1)
87
- publisher.publish(batch) # 20th message
87
+ subscriber.request_replay(10)
88
89
replayed = subscriber.receive_replay()
90
91
- assert len(replayed) > 0, "No replayed messages received"
92
- seqs = [seq for seq, _ in replayed]
93
- assert all(seq >= 10 for seq in seqs), "Replayed messages not in order"
94
- assert seqs == list(range(min(seqs), max(seqs) + 1)), (
95
- "Replayed messages not consecutive"
91
+ assert len(replayed) == 9, (
92
+ f"Expected 9 replayed messages (seq 10-18), got {len(replayed)}"
93
)
94
+ seqs = [seq for seq, _ in replayed]
95
+ assert seqs == list(range(10, 19)), "Replayed sequences should be 10-18"
96
+
97
+
98
+def test_replay_includes_topic(publisher, subscriber, publisher_config):
99
+ """Test that replay responses include the topic, matching PUB format"""
100
+ for _ in range(5):
101
+ publisher.publish(create_test_events(1))
102
+
103
+ # Drain live events to ensure publisher has processed them.
104
+ for _ in range(5):
105
+ assert subscriber.receive_one(timeout=1000) is not None
106
+
107
+ subscriber.request_replay(0)
108
+
109
+ # receive_replay unpacks (topic, seq, payload) and asserts
110
+ # topic == publisher topic for each message.
111
+ replayed = subscriber.receive_replay()
112
+ assert len(replayed) == 5, f"Expected 5 replayed messages, got {len(replayed)}"
113
+ seqs = [seq for seq, _ in replayed]
114
+ assert seqs == list(range(5)), "Replayed sequences should be 0-4"
115
116
117
def test_buffer_limit(publisher, subscriber, publisher_config):
@@ -108,15 +126,16 @@ def test_buffer_limit(publisher, subscriber, publisher_config):
126
time.sleep(0.5) # Need publisher to process above requests
127
subscriber.request_replay(0)
128
111
- batch = create_test_events(1)
112
- publisher.publish(batch)
113
-
129
replayed = subscriber.receive_replay()
130
116
- assert len(replayed) <= buffer_size, "Can't replay more than buffer size"
131
+ assert len(replayed) == buffer_size, (
132
+ f"Expected {buffer_size} replayed messages, got {len(replayed)}"
133
+ )
134
118
- oldest_seq = min(seq for seq, _ in replayed)
119
- assert oldest_seq >= 10, "The oldest sequence should be at least 10"
135
+ seqs = [seq for seq, _ in replayed]
136
+ assert seqs == list(range(10, buffer_size + 10)), (
137
+ "Should replay seq 11 through buffer_size+10"
138
+ )
139
140
141
def test_topic_filtering(publisher_config):
vllm/distributed/kv_events.py
+3
-6
@@ -459,15 +459,12 @@ class ZmqEventPublisher(EventPublisher):
459
460
for seq, buf in self._buffer:
461
if seq >= start_seq:
462
- # [identity, empty_delim, seq_bytes, payload]
463
- # (identity, empty_delim) are stripped off by the router
464
- # receiving payload is (seq_bytes, payload)
462
+ # Subscriber receives (topic, seq_bytes, payload)
463
self._replay.send_multipart(
466
- (client_id, b"", seq.to_bytes(8, "big"), buf)
464
+ (client_id, b"", self._topic_bytes, seq.to_bytes(8, "big"), buf)
465
)
466
# Send end of sequence marker
469
- # receiving payload is (-1, b""")
470
- self._replay.send_multipart((client_id, b"", self.END_SEQ, b""))
467
+ self._replay.send_multipart((client_id, b"", b"", self.END_SEQ, b""))
468
469
@staticmethod
470
def offset_endpoint_port(