| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | import random |
| 5 | |
| 6 | import msgspec |
| 7 | import msgspec.msgpack |
| 8 | import pytest |
| 9 | import zmq |
| 10 | |
| 11 | from vllm.config.kv_events import KVEventsConfig |
| 12 | from vllm.distributed.kv_events import EventPublisherFactory |
| 13 | |
| 14 | from .test_events import SampleBatch |
| 15 | |
| 16 | DP_RANK = 0 |
| 17 | |
| 18 | |
| 19 | @pytest.fixture |
| 20 | def random_port(): |
| 21 | """Generate a random port number for testing.""" |
| 22 | return random.randint(10000, 59900) |
| 23 | |
| 24 | |
| 25 | @pytest.fixture |
| 26 | def publisher_config(random_port, request): |
| 27 | """Create a publisher config with inproc transport.""" |
| 28 | how = request.param if hasattr(request, "param") else "inproc" |
| 29 | |
| 30 | if how == "inproc": |
| 31 | endpoint = f"inproc://test-{random_port}" |
| 32 | replay_endpoint = endpoint + "-replay" |
| 33 | else: |
| 34 | endpoint = f"tcp://*:{random_port}" |
| 35 | replay_endpoint = f"tcp://*:{random_port + 100}" |
| 36 | |
| 37 | return KVEventsConfig( |
| 38 | enable_kv_cache_events=True, |
| 39 | publisher="zmq", |
| 40 | endpoint=endpoint, |
| 41 | replay_endpoint=replay_endpoint, |
| 42 | buffer_steps=100, |
| 43 | hwm=1000, |
| 44 | topic="test", |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | @pytest.fixture |
| 49 | def publisher(publisher_config): |
| 50 | """Create and return a publisher instance.""" |
| 51 | pub = EventPublisherFactory.create(publisher_config, DP_RANK) |
| 52 | yield pub |
| 53 | pub.shutdown() |
| 54 | |
| 55 | |
| 56 | @pytest.fixture |
| 57 | def subscriber(publisher_config): |
| 58 | """Create and return a subscriber for testing.""" |
| 59 | endpoint = publisher_config.endpoint |
| 60 | replay_endpoint = publisher_config.replay_endpoint |
| 61 | |
| 62 | if endpoint.startswith("tcp://*"): |
| 63 | endpoint = endpoint.replace("*", "127.0.0.1") |
| 64 | if replay_endpoint and replay_endpoint.startswith("tcp://*"): |
| 65 | replay_endpoint = replay_endpoint.replace("*", "127.0.0.1") |
| 66 | |
| 67 | sub = MockSubscriber( |
| 68 | [endpoint], |
| 69 | [replay_endpoint] if replay_endpoint else None, |
| 70 | publisher_config.topic, |
| 71 | ) |
| 72 | yield sub |
| 73 | sub.close() |
| 74 | |
| 75 | |
| 76 | class MockSubscriber: |
| 77 | """Helper class to receive and verify published events.""" |
| 78 | |
| 79 | def __init__( |
| 80 | self, |
| 81 | pub_endpoints: str | list[str], |
| 82 | replay_endpoints: str | list[str] | None = None, |
| 83 | topic: str = "", |
| 84 | decode_type=SampleBatch, |
| 85 | ): |
| 86 | self.ctx = zmq.Context.instance() |
| 87 | |
| 88 | # Convert single endpoint to list for consistency |
| 89 | if isinstance(pub_endpoints, str): |
| 90 | pub_endpoints = [pub_endpoints] |
| 91 | if isinstance(replay_endpoints, str): |
| 92 | replay_endpoints = [replay_endpoints] |
| 93 | |
| 94 | ipv6 = any("[" in ep for ep in pub_endpoints + (replay_endpoints or [])) |
| 95 | |
| 96 | # Set up subscriber socket - connect to all endpoints |
| 97 | self.sub = self.ctx.socket(zmq.SUB) |
| 98 | self.sub.setsockopt(zmq.IPV6, ipv6) |
| 99 | self.sub.setsockopt(zmq.SUBSCRIBE, topic.encode("utf-8")) |
| 100 | for endpoint in pub_endpoints: |
| 101 | self.sub.connect(endpoint) |
| 102 | |
| 103 | # Set up replay sockets if provided. |
| 104 | # DEALER allows receiving multiple replies per request. |
| 105 | self.replay_sockets = [] |
| 106 | if replay_endpoints: |
| 107 | for replay_endpoint in replay_endpoints: |
| 108 | replay = self.ctx.socket(zmq.DEALER) |
| 109 | replay.setsockopt(zmq.IPV6, ipv6) |
| 110 | replay.connect(replay_endpoint) |
| 111 | self.replay_sockets.append(replay) |
| 112 | |
| 113 | self.topic = topic |
| 114 | self.topic_bytes = topic.encode("utf-8") |
| 115 | self.received_msgs: list[tuple[int, SampleBatch]] = [] |
| 116 | self.last_seq = -1 |
| 117 | self.decoder = msgspec.msgpack.Decoder(type=decode_type) |
| 118 | |
| 119 | def receive_one(self, timeout=1000) -> tuple[int, SampleBatch] | None: |
| 120 | """Receive a single message with timeout.""" |
| 121 | if not self.sub.poll(timeout): |
| 122 | return None |
| 123 | |
| 124 | topic_bytes, seq_bytes, payload = self.sub.recv_multipart() |
| 125 | assert topic_bytes == self.topic_bytes |
| 126 | |
| 127 | seq = int.from_bytes(seq_bytes, "big") |
| 128 | data = self.decoder.decode(payload) |
| 129 | self.last_seq = seq |
| 130 | self.received_msgs.append((seq, data)) |
| 131 | return seq, data |
| 132 | |
| 133 | def request_replay(self, start_seq: int, socket_idx: int = 0) -> None: |
| 134 | """Request replay of messages starting from start_seq.""" |
| 135 | if not self.replay_sockets: |
| 136 | raise ValueError("Replay sockets not initialized") |
| 137 | if socket_idx >= len(self.replay_sockets): |
| 138 | raise ValueError(f"Invalid socket index {socket_idx}") |
| 139 | |
| 140 | self.replay_sockets[socket_idx].send_multipart( |
| 141 | [b"", start_seq.to_bytes(8, "big")] |
| 142 | ) |
| 143 | |
| 144 | def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]: |
| 145 | """Receive replayed messages from a specific replay socket.""" |
| 146 | if not self.replay_sockets: |
| 147 | raise ValueError("Replay sockets not initialized") |
| 148 | if socket_idx >= len(self.replay_sockets): |
| 149 | raise ValueError(f"Invalid socket index {socket_idx}") |
| 150 | |
| 151 | replay_socket = self.replay_sockets[socket_idx] |
| 152 | replayed: list[tuple[int, SampleBatch]] = [] |
| 153 | while True: |
| 154 | try: |
| 155 | if not replay_socket.poll(1000): |
| 156 | break |
| 157 | |
| 158 | # DEALER receives [empty_delim, topic, seq, payload] |
| 159 | frames = replay_socket.recv_multipart() |
| 160 | if frames and frames[0] == b"": |
| 161 | frames = frames[1:] |
| 162 | if len(frames) != 3 or not frames[-1]: |
| 163 | # End of replay marker |
| 164 | break |
| 165 | |
| 166 | topic, seq_bytes, payload = frames |
| 167 | assert topic == self.topic_bytes |
| 168 | seq = int.from_bytes(seq_bytes, "big") |
| 169 | data = self.decoder.decode(payload) |
| 170 | replayed.append((seq, data)) |
| 171 | except zmq.ZMQError as _: |
| 172 | break |
| 173 | |
| 174 | return replayed |
| 175 | |
| 176 | def close(self): |
| 177 | """Clean up resources.""" |
| 178 | self.sub.close() |
| 179 | for replay in self.replay_sockets: |
| 180 | replay.close() |
| 181 | |
| 182 | |
| 183 | @pytest.fixture |
| 184 | def enable_ray_v2_backend(): |
| 185 | """Set env vars for the Ray V2 executor backend and shut down Ray |
| 186 | between tests.""" |
| 187 | import ray |
| 188 | |
| 189 | saved = { |
| 190 | "VLLM_USE_RAY_V2_EXECUTOR_BACKEND": os.environ.get( |
| 191 | "VLLM_USE_RAY_V2_EXECUTOR_BACKEND" |
| 192 | ), |
| 193 | "VLLM_ENABLE_V1_MULTIPROCESSING": os.environ.get( |
| 194 | "VLLM_ENABLE_V1_MULTIPROCESSING" |
| 195 | ), |
| 196 | } |
| 197 | os.environ["VLLM_USE_RAY_V2_EXECUTOR_BACKEND"] = "1" |
| 198 | os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" |
| 199 | if ray.is_initialized(): |
| 200 | ray.shutdown() |
| 201 | try: |
| 202 | yield |
| 203 | finally: |
| 204 | if ray.is_initialized(): |
| 205 | ray.shutdown() |
| 206 | os.environ.update({k: v for k, v in saved.items() if v is not None}) |
| 207 | for key in (k for k, v in saved.items() if v is None): |
| 208 | os.environ.pop(key, None) |