| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import threading |
| 4 | import time |
| 5 | |
| 6 | import msgspec |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.distributed.kv_events import ( |
| 10 | EventBatch, |
| 11 | EventPublisherFactory, |
| 12 | NullEventPublisher, |
| 13 | ) |
| 14 | |
| 15 | DP_RANK = 0 |
| 16 | |
| 17 | |
| 18 | class EventSample( |
| 19 | msgspec.Struct, |
| 20 | tag=True, # type: ignore |
| 21 | array_like=True, # type: ignore |
| 22 | ): |
| 23 | """Test event for publisher testing.""" |
| 24 | |
| 25 | id: int |
| 26 | value: str |
| 27 | |
| 28 | |
| 29 | class SampleBatch(EventBatch): |
| 30 | """Test event batch for publisher testing.""" |
| 31 | |
| 32 | events: list[EventSample] |
| 33 | |
| 34 | |
| 35 | def create_test_events(count: int) -> SampleBatch: |
| 36 | """Create a batch of test events.""" |
| 37 | events = [EventSample(id=i, value=f"test-{i}") for i in range(count)] |
| 38 | return SampleBatch(ts=time.time(), events=events) |
| 39 | |
| 40 | |
| 41 | def test_basic_publishing(publisher, subscriber): |
| 42 | """Test basic event publishing works.""" |
| 43 | test_batch = create_test_events(5) |
| 44 | publisher.publish(test_batch) |
| 45 | |
| 46 | result = subscriber.receive_one(timeout=1000) |
| 47 | assert result is not None, "No message received" |
| 48 | |
| 49 | seq, received = result |
| 50 | assert seq == 0, "Sequence number mismatch" |
| 51 | assert received.ts == pytest.approx(test_batch.ts, abs=0.1), "Timestamp mismatch" |
| 52 | assert len(received.events) == len(test_batch.events), "Number of events mismatch" |
| 53 | |
| 54 | for i, event in enumerate(received.events): |
| 55 | assert event.id == i, "Event id mismatch" |
| 56 | assert event.value == f"test-{i}", "Event value mismatch" |
| 57 | |
| 58 | |
| 59 | def test_multiple_events(publisher, subscriber): |
| 60 | """Test publishing and receiving multiple event batches.""" |
| 61 | for _ in range(10): |
| 62 | batch = create_test_events(2) |
| 63 | publisher.publish(batch) |
| 64 | |
| 65 | received = [] |
| 66 | for _ in range(10): |
| 67 | data = subscriber.receive_one(timeout=100) |
| 68 | if data: |
| 69 | received.append(data) |
| 70 | |
| 71 | assert len(received) == 10, "Number of messages mismatch" |
| 72 | seqs = [seq for seq, _ in received] |
| 73 | assert seqs == list(range(10)), "Sequence numbers mismatch" |
| 74 | |
| 75 | |
| 76 | def test_replay_mechanism(publisher, subscriber): |
| 77 | """Test the replay mechanism works correctly.""" |
| 78 | for _ in range(19): |
| 79 | batch = create_test_events(1) |
| 80 | publisher.publish(batch) |
| 81 | |
| 82 | # Drain live events to ensure publisher has buffered them. |
| 83 | for _ in range(19): |
| 84 | assert subscriber.receive_one(timeout=1000) is not None |
| 85 | |
| 86 | subscriber.request_replay(10) |
| 87 | |
| 88 | replayed = subscriber.receive_replay() |
| 89 | |
| 90 | assert len(replayed) == 9, ( |
| 91 | f"Expected 9 replayed messages (seq 10-18), got {len(replayed)}" |
| 92 | ) |
| 93 | seqs = [seq for seq, _ in replayed] |
| 94 | assert seqs == list(range(10, 19)), "Replayed sequences should be 10-18" |
| 95 | |
| 96 | |
| 97 | def test_replay_includes_topic(publisher, subscriber, publisher_config): |
| 98 | """Test that replay responses include the topic, matching PUB format.""" |
| 99 | for _ in range(5): |
| 100 | publisher.publish(create_test_events(1)) |
| 101 | |
| 102 | # Drain live events to ensure publisher has processed them. |
| 103 | for _ in range(5): |
| 104 | assert subscriber.receive_one(timeout=1000) is not None |
| 105 | |
| 106 | subscriber.request_replay(0) |
| 107 | |
| 108 | # receive_replay unpacks (topic, seq, payload) and asserts |
| 109 | # topic == publisher topic for each message. |
| 110 | replayed = subscriber.receive_replay() |
| 111 | assert len(replayed) == 5, f"Expected 5 replayed messages, got {len(replayed)}" |
| 112 | seqs = [seq for seq, _ in replayed] |
| 113 | assert seqs == list(range(5)), "Replayed sequences should be 0-4" |
| 114 | |
| 115 | |
| 116 | def test_buffer_limit(publisher, subscriber, publisher_config): |
| 117 | """Test buffer limit behavior.""" |
| 118 | buffer_size = publisher_config.buffer_steps |
| 119 | |
| 120 | # Publish more events than the buffer can hold |
| 121 | for i in range(buffer_size + 10): |
| 122 | batch = create_test_events(1) |
| 123 | publisher.publish(batch) |
| 124 | |
| 125 | time.sleep(0.5) # Need publisher to process above requests |
| 126 | subscriber.request_replay(0) |
| 127 | |
| 128 | replayed = subscriber.receive_replay() |
| 129 | |
| 130 | assert len(replayed) == buffer_size, ( |
| 131 | f"Expected {buffer_size} replayed messages, got {len(replayed)}" |
| 132 | ) |
| 133 | |
| 134 | seqs = [seq for seq, _ in replayed] |
| 135 | assert seqs == list(range(10, buffer_size + 10)), ( |
| 136 | "Should replay seq 11 through buffer_size+10" |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | def test_topic_filtering(publisher_config): |
| 141 | """Test that a subscriber only receives messages matching its topic filter.""" |
| 142 | publisher_config.replay_endpoint = None |
| 143 | |
| 144 | publisher_config.topic = "foo" |
| 145 | pub = EventPublisherFactory.create(publisher_config, DP_RANK) |
| 146 | |
| 147 | from .conftest import MockSubscriber |
| 148 | |
| 149 | sub_foo = MockSubscriber(publisher_config.endpoint, None, "foo") |
| 150 | sub_bar = MockSubscriber(publisher_config.endpoint, None, "bar") |
| 151 | |
| 152 | try: |
| 153 | time.sleep(0.1) |
| 154 | |
| 155 | for _ in range(3): |
| 156 | pub.publish(create_test_events(1)) |
| 157 | |
| 158 | foo_received = [sub_foo.receive_one(timeout=200) for _ in range(3)] |
| 159 | assert all(msg is not None for msg in foo_received), ( |
| 160 | "Subscriber with matching topic should receive messages" |
| 161 | ) |
| 162 | |
| 163 | bar_received = [sub_bar.receive_one(timeout=200) for _ in range(3)] |
| 164 | assert all(msg is None for msg in bar_received), ( |
| 165 | "Subscriber with non-matching topic should receive no messages" |
| 166 | ) |
| 167 | finally: |
| 168 | pub.shutdown() |
| 169 | sub_foo.close() |
| 170 | sub_bar.close() |
| 171 | |
| 172 | |
| 173 | def test_high_volume(publisher, subscriber): |
| 174 | """Test publishing and receiving a high volume of events.""" |
| 175 | num_batches = 10_000 |
| 176 | events_per_batch = 100 |
| 177 | |
| 178 | # Publish events in a separate thread to not block |
| 179 | def publish_events(): |
| 180 | for i in range(num_batches): |
| 181 | batch = create_test_events(events_per_batch) |
| 182 | publisher.publish(batch) |
| 183 | # Small delay to avoid overwhelming |
| 184 | if i % 100 == 0: |
| 185 | time.sleep(0.01) |
| 186 | |
| 187 | received: list[tuple[int, SampleBatch]] = [] |
| 188 | |
| 189 | publisher_thread = threading.Thread(target=publish_events) |
| 190 | publisher_thread.start() |
| 191 | |
| 192 | start_time = time.time() |
| 193 | while len(received) < num_batches: |
| 194 | if time.time() - start_time > 10: # Timeout after 10 seconds |
| 195 | break |
| 196 | |
| 197 | result = subscriber.receive_one(timeout=100) |
| 198 | if result: |
| 199 | received.append(result) |
| 200 | |
| 201 | publisher_thread.join() |
| 202 | |
| 203 | assert len(received) >= num_batches * 0.9, "We should have received most messages" |
| 204 | |
| 205 | seqs = [seq for seq, _ in received] |
| 206 | assert sorted(seqs) == seqs, "Sequence numbers should be in order" |
| 207 | |
| 208 | |
| 209 | def test_null_publisher(): |
| 210 | """Test that NullEventPublisher can be used without errors.""" |
| 211 | publisher = NullEventPublisher(DP_RANK) |
| 212 | |
| 213 | # This should not raise any errors |
| 214 | batch = create_test_events(5) |
| 215 | publisher.publish(batch) |
| 216 | publisher.shutdown() |
| 217 | |
| 218 | |
| 219 | def test_data_parallel_rank_tagging(publisher_config): |
| 220 | """Test that events are properly tagged with their data parallel rank.""" |
| 221 | publisher_config.topic = "foo" |
| 222 | pub_0 = EventPublisherFactory.create(publisher_config, DP_RANK) |
| 223 | pub_1 = EventPublisherFactory.create(publisher_config, DP_RANK + 1) |
| 224 | |
| 225 | # Hardcode the expected endpoints based on port offsetting behavior |
| 226 | # Both ranks get offsets according to _offset_endpoint_port function |
| 227 | base_endpoint = publisher_config.endpoint |
| 228 | if "tcp://" in base_endpoint: |
| 229 | # For TCP endpoints: tcp://localhost:5557 -> tcp://localhost:5557, tcp://localhost:5558 |
| 230 | expected_endpoint_0 = base_endpoint # rank 0 gets port + 0 = same port |
| 231 | expected_endpoint_1 = base_endpoint.replace( |
| 232 | ":5557", ":5558" |
| 233 | ) # rank 1 gets port + 1 |
| 234 | else: |
| 235 | # For inproc endpoints: inproc://test -> inproc://test_dp0, inproc://test_dp1 |
| 236 | expected_endpoint_0 = base_endpoint # rank 0 gets base |
| 237 | expected_endpoint_1 = base_endpoint + "_dp1" # rank 1 gets _dp1 |
| 238 | |
| 239 | from .conftest import MockSubscriber |
| 240 | |
| 241 | sub_0 = MockSubscriber(expected_endpoint_0, None, publisher_config.topic) |
| 242 | sub_1 = MockSubscriber(expected_endpoint_1, None, publisher_config.topic) |
| 243 | |
| 244 | try: |
| 245 | time.sleep(0.1) # Let publishers start up |
| 246 | |
| 247 | # Publish events from different ranks |
| 248 | batch_0 = create_test_events(2) |
| 249 | batch_1 = create_test_events(3) |
| 250 | |
| 251 | pub_0.publish(batch_0) |
| 252 | pub_1.publish(batch_1) |
| 253 | |
| 254 | # Receive events from rank 0 |
| 255 | result_0 = sub_0.receive_one(timeout=200) |
| 256 | assert result_0 is not None, "No message received from rank 0" |
| 257 | seq_0, received_0 = result_0 |
| 258 | |
| 259 | # Receive events from rank 1 |
| 260 | result_1 = sub_1.receive_one(timeout=200) |
| 261 | assert result_1 is not None, "No message received from rank 1" |
| 262 | seq_1, received_1 = result_1 |
| 263 | |
| 264 | # Verify DP rank tagging |
| 265 | assert received_0.data_parallel_rank == 0, ( |
| 266 | f"Expected DP rank 0, got {received_0.data_parallel_rank}" |
| 267 | ) |
| 268 | assert received_1.data_parallel_rank == 1, ( |
| 269 | f"Expected DP rank 1, got {received_1.data_parallel_rank}" |
| 270 | ) |
| 271 | |
| 272 | # Verify event content is correct |
| 273 | assert len(received_0.events) == 2, "Wrong number of events from rank 0" |
| 274 | assert len(received_1.events) == 3, "Wrong number of events from rank 1" |
| 275 | |
| 276 | finally: |
| 277 | pub_0.shutdown() |
| 278 | pub_1.shutdown() |
| 279 | sub_0.close() |
| 280 | sub_1.close() |
| 281 | |
| 282 | |
| 283 | def test_event_publisher_factory(random_port): |
| 284 | """Test event publisher factory creation behavior under different configurations.""" |
| 285 | from vllm.config.kv_events import KVEventsConfig |
| 286 | from vllm.distributed.kv_events import ZmqEventPublisher |
| 287 | |
| 288 | # test config is None |
| 289 | publisher = EventPublisherFactory.create(None, DP_RANK) |
| 290 | assert isinstance(publisher, NullEventPublisher) |
| 291 | publisher.shutdown() |
| 292 | |
| 293 | # test disable kv cache events |
| 294 | config = KVEventsConfig( |
| 295 | enable_kv_cache_events=False, |
| 296 | publisher="zmq", # Even if zmq is specified, should return NullEventPublisher |
| 297 | endpoint="tcp://localhost:5557", |
| 298 | ) |
| 299 | publisher = EventPublisherFactory.create(config, DP_RANK) |
| 300 | assert isinstance(publisher, NullEventPublisher) |
| 301 | publisher.shutdown() |
| 302 | |
| 303 | # test zmq publisher |
| 304 | config = KVEventsConfig( |
| 305 | enable_kv_cache_events=True, |
| 306 | publisher="zmq", |
| 307 | endpoint=f"tcp://*:{random_port}", |
| 308 | replay_endpoint=f"tcp://*:{random_port + 100}", |
| 309 | ) |
| 310 | publisher = EventPublisherFactory.create(config, DP_RANK + 1) |
| 311 | assert isinstance(publisher, ZmqEventPublisher) |
| 312 | resolved_config = publisher.get_publisher_config() |
| 313 | assert resolved_config.endpoint == f"tcp://*:{random_port + 1}" |
| 314 | assert resolved_config.replay_endpoint == f"tcp://*:{random_port + 101}" |
| 315 | assert config.endpoint == f"tcp://*:{random_port}" |
| 316 | publisher.shutdown() |
| 317 | |
| 318 | # test unknown publisher |
| 319 | with pytest.raises(ValueError, match="Input should be"): |
| 320 | KVEventsConfig( |
| 321 | enable_kv_cache_events=True, |
| 322 | publisher="unknown_publisher", |
| 323 | endpoint="tcp://localhost:5557", |
| 324 | ) |
| 325 | |
| 326 | # test publisher not specified |
| 327 | config = KVEventsConfig( |
| 328 | enable_kv_cache_events=True, |
| 329 | # publisher not specified, should default to "zmq" |
| 330 | endpoint="tcp://localhost:5557", |
| 331 | ) |
| 332 | publisher = EventPublisherFactory.create(config, DP_RANK) |
| 333 | assert isinstance(publisher, ZmqEventPublisher) |
| 334 | publisher.shutdown() |
| 335 | |
| 336 | |
| 337 | @pytest.mark.parametrize( |
| 338 | ("endpoint", "expected"), |
| 339 | [("tcp://*:5557", "tcp://*:5558"), ("tcp://*:0", "tcp://*:0")], |
| 340 | ) |
| 341 | def test_offset_endpoint_port_keeps_port_zero(endpoint, expected): |
| 342 | """Port 0 is not offset: every rank asks the OS for its own port.""" |
| 343 | from vllm.distributed.kv_events import ZmqEventPublisher |
| 344 | |
| 345 | assert ZmqEventPublisher.offset_endpoint_port(endpoint, 1) == expected |
| 346 | |
| 347 | |
| 348 | @pytest.mark.parametrize("node_ip", ["127.0.0.1", "::1"]) |
| 349 | def test_ephemeral_publishers_report_bound_endpoints(monkeypatch, node_ip): |
| 350 | """Two ranks binding tcp://*:0 get distinct OS-assigned ports on this |
| 351 | node's IP and report endpoints a subscriber can use for events and replay.""" |
| 352 | from vllm.config.kv_events import KVEventsConfig |
| 353 | from vllm.distributed import kv_events |
| 354 | from vllm.utils.network_utils import split_zmq_path |
| 355 | |
| 356 | from .conftest import MockSubscriber |
| 357 | |
| 358 | monkeypatch.setattr(kv_events, "get_ip", lambda: node_ip) |
| 359 | config = KVEventsConfig( |
| 360 | enable_kv_cache_events=True, |
| 361 | publisher="zmq", |
| 362 | endpoint="tcp://*:0", |
| 363 | replay_endpoint="tcp://*:0", |
| 364 | topic="kv", |
| 365 | ) |
| 366 | pubs = [EventPublisherFactory.create(config, rank) for rank in range(2)] |
| 367 | try: |
| 368 | resolved = [pub.get_publisher_config() for pub in pubs] |
| 369 | endpoints = [e for c in resolved for e in (c.endpoint, c.replay_endpoint)] |
| 370 | hosts_ports = {split_zmq_path(e)[1:] for e in endpoints} |
| 371 | assert {h for h, _ in hosts_ports} == {node_ip} |
| 372 | assert len(hosts_ports) == 4 and "0" not in {p for _, p in hosts_ports} |
| 373 | assert config.endpoint == "tcp://*:0" |
| 374 | |
| 375 | sub = MockSubscriber(endpoints[::2], endpoints[1::2], topic="kv") |
| 376 | try: |
| 377 | # PUB/SUB drops messages sent before the subscription lands, so |
| 378 | # keep publishing until both ranks are heard. |
| 379 | ranks: set[int] = set() |
| 380 | deadline = time.time() + 10 |
| 381 | while ranks != {0, 1} and time.time() < deadline: |
| 382 | for pub in pubs: |
| 383 | pub.publish(create_test_events(1)) |
| 384 | while (result := sub.receive_one(timeout=100)) is not None: |
| 385 | ranks.add(result[1].data_parallel_rank) |
| 386 | assert ranks == {0, 1} |
| 387 | sub.request_replay(0, socket_idx=1) |
| 388 | assert sub.receive_replay(socket_idx=1) |
| 389 | finally: |
| 390 | sub.close() |
| 391 | finally: |
| 392 | for pub in pubs: |
| 393 | pub.shutdown() |
| 394 | |
| 395 | |
| 396 | @pytest.mark.parametrize("node_ip", ["0.0.0.0", "::"]) |
| 397 | def test_ephemeral_bind_needs_a_node_ip(monkeypatch, node_ip): |
| 398 | """tcp://*:0 is advertised as this node's IP; with no usable IP the |
| 399 | publisher refuses instead of reporting a wildcard.""" |
| 400 | from vllm.config.kv_events import KVEventsConfig |
| 401 | from vllm.distributed import kv_events |
| 402 | |
| 403 | monkeypatch.setattr(kv_events, "get_ip", lambda: node_ip) |
| 404 | config = KVEventsConfig( |
| 405 | enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://*:0" |
| 406 | ) |
| 407 | with pytest.raises(ValueError, match="VLLM_HOST_IP"): |
| 408 | EventPublisherFactory.create(config, DP_RANK) |