fix: include topic frame in KV events replay response (#45177)

Signed-off-by: RishabhSaini <rishabhsaini01@gmail.com>

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

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))

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):

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(