| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ZMQ-based transport layer for P2P KV cache sharing. |
| 4 | |
| 5 | Provides ZmqConnection (per-peer messaging) and ZmqTransport (connection |
| 6 | management). Message-content agnostic. |
| 7 | """ |
| 8 | |
| 9 | from __future__ import annotations |
| 10 | |
| 11 | from collections.abc import Sequence |
| 12 | from dataclasses import dataclass |
| 13 | |
| 14 | import msgspec |
| 15 | import zmq |
| 16 | import zmq.utils.monitor |
| 17 | |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.v1.kv_offload.tiering.p2p.control.base import ( |
| 20 | ControlConnection, |
| 21 | ControlTransport, |
| 22 | ) |
| 23 | |
| 24 | logger = init_logger(__name__) |
| 25 | |
| 26 | _HEARTBEAT_IVL_MS = 2000 |
| 27 | _HEARTBEAT_TIMEOUT_MS = 10000 |
| 28 | _HEARTBEAT_TTL_MS = 10000 |
| 29 | |
| 30 | # Shared sentinels returned when there is nothing to report. |
| 31 | _EMPTY_INBOX: tuple[dict, ...] = () |
| 32 | _EMPTY_NEW_CONNECTIONS: tuple[ControlConnection, ...] = () |
| 33 | |
| 34 | |
| 35 | def _tcp_addr(host: str, port: int | str) -> str: |
| 36 | return f"tcp://{host}:{port}" |
| 37 | |
| 38 | |
| 39 | def _apply_heartbeat(sock: zmq.Socket) -> None: |
| 40 | sock.setsockopt(zmq.HEARTBEAT_IVL, _HEARTBEAT_IVL_MS) |
| 41 | sock.setsockopt(zmq.HEARTBEAT_TIMEOUT, _HEARTBEAT_TIMEOUT_MS) |
| 42 | sock.setsockopt(zmq.HEARTBEAT_TTL, _HEARTBEAT_TTL_MS) |
| 43 | |
| 44 | |
| 45 | @dataclass |
| 46 | class _Sockets: |
| 47 | dealer: zmq.Socket |
| 48 | monitor: zmq.Socket |
| 49 | |
| 50 | |
| 51 | class ZmqConnection(ControlConnection): |
| 52 | """Bidirectional message channel to a single remote peer.""" |
| 53 | |
| 54 | def __init__(self, peer_id: str, sockets: _Sockets) -> None: |
| 55 | super().__init__(peer_id) |
| 56 | self._sockets = sockets |
| 57 | # _dead: the peer is gone. _closed: the sockets have been released. |
| 58 | # Distinct, so mark_dead() cannot turn close() into a no-op and leak |
| 59 | # the DEALER and its monitor socket. |
| 60 | self._dead = False |
| 61 | self._closed = False |
| 62 | self._inbox: list[dict] = [] |
| 63 | |
| 64 | def send(self, msg: dict) -> None: |
| 65 | """Send a msgpack-encoded message to this peer.""" |
| 66 | if not self.alive: |
| 67 | raise RuntimeError( |
| 68 | f"ZmqConnection: send on closed connection to {self.peer_id}" |
| 69 | ) |
| 70 | data = msgspec.msgpack.encode(msg) |
| 71 | self._sockets.dealer.send(data) |
| 72 | |
| 73 | def recv(self) -> Sequence[dict]: |
| 74 | """Drain and return all buffered incoming messages.""" |
| 75 | if not self._inbox: |
| 76 | return _EMPTY_INBOX |
| 77 | msgs = self._inbox |
| 78 | self._inbox = [] |
| 79 | return msgs |
| 80 | |
| 81 | @property |
| 82 | def alive(self) -> bool: |
| 83 | return not (self._dead or self._closed) |
| 84 | |
| 85 | def close(self) -> None: |
| 86 | if self._closed: |
| 87 | return |
| 88 | self._closed = True |
| 89 | self._dead = True |
| 90 | logger.info("ZmqConnection: closing connection to %s", self.peer_id) |
| 91 | self._sockets.monitor.close() |
| 92 | self._sockets.dealer.close() |
| 93 | |
| 94 | def enqueue(self, msg: dict) -> None: |
| 95 | """Buffer an incoming message.""" |
| 96 | self._inbox.append(msg) |
| 97 | |
| 98 | def mark_dead(self) -> None: |
| 99 | """Mark connection as disconnected. |
| 100 | |
| 101 | Only flips liveness: the sockets stay open until close() releases |
| 102 | them, so the owner can still drain recv() before tearing down. |
| 103 | """ |
| 104 | self._dead = True |
| 105 | |
| 106 | @property |
| 107 | def monitor_socket(self) -> zmq.Socket: |
| 108 | """Monitor socket for disconnect detection (used by ZmqTransport).""" |
| 109 | return self._sockets.monitor |
| 110 | |
| 111 | |
| 112 | class ZmqTransport(ControlTransport): |
| 113 | """ZMQ implementation of ControlTransport. |
| 114 | |
| 115 | Manages a ROUTER socket for accepting connections and DEALER sockets |
| 116 | for outbound connections. Message-content agnostic. |
| 117 | """ |
| 118 | |
| 119 | def __init__(self, local_id: str, host: str, port: int) -> None: |
| 120 | self._local_id = local_id |
| 121 | self._closed = False |
| 122 | |
| 123 | self._connections: dict[str, ZmqConnection] = {} |
| 124 | self._pending_inbound: list[tuple[str, dict]] = [] |
| 125 | # Monotonic suffix for inproc monitor endpoints — see |
| 126 | # _open_connection() for why peer_id alone is not enough. |
| 127 | self._monitor_seq = 0 |
| 128 | |
| 129 | self._zmq_ctx = zmq.Context() |
| 130 | self._router: zmq.Socket = self._zmq_ctx.socket(zmq.ROUTER) |
| 131 | _apply_heartbeat(self._router) |
| 132 | bind_addr = _tcp_addr(host, port) |
| 133 | self._router.bind(bind_addr) |
| 134 | logger.info("ZmqTransport %s: ROUTER bound on %s", self._local_id, bind_addr) |
| 135 | |
| 136 | # ------------------------------------------------------------------ |
| 137 | # ZmqConnection lifecycle |
| 138 | # ------------------------------------------------------------------ |
| 139 | |
| 140 | def connect(self, peer_id: str) -> ZmqConnection: |
| 141 | """Open an outbound connection to a remote peer. |
| 142 | |
| 143 | A dead connection can still be registered: its owning session may mark |
| 144 | it dead after this tick's sweep already ran, and poll() only |
| 145 | unregisters it on the next pass. Retire such an entry instead of |
| 146 | asserting, so a reconnect landing in that window succeeds. A live |
| 147 | entry is still a genuine duplicate. |
| 148 | """ |
| 149 | existing = self._connections.get(peer_id) |
| 150 | if existing is not None: |
| 151 | assert not existing.alive, f"ZmqConnection to {peer_id} already exists" |
| 152 | logger.info( |
| 153 | "ZmqTransport %s: retiring dead connection to %s before reconnect", |
| 154 | self._local_id, |
| 155 | peer_id, |
| 156 | ) |
| 157 | self._connections.pop(peer_id).close() |
| 158 | logger.info( |
| 159 | "ZmqTransport %s: opening OUTBOUND connection to %s", |
| 160 | self._local_id, |
| 161 | peer_id, |
| 162 | ) |
| 163 | return self._open_connection(peer_id, direction="outbound") |
| 164 | |
| 165 | def poll(self) -> Sequence[ControlConnection]: |
| 166 | """Process all pending I/O. Returns newly accepted connections. |
| 167 | |
| 168 | - Receives messages (buffered in each connection's inbox) |
| 169 | - Creates connections for new inbound peers (connect msg in inbox) |
| 170 | - Checks monitors for disconnections |
| 171 | - Removes and closes dead connections |
| 172 | """ |
| 173 | # Retire connections a session killed since the last poll() before |
| 174 | # routing traffic: otherwise a reconnecting peer's first message is |
| 175 | # enqueued into the dead connection and discarded along with it. |
| 176 | # This closes only that between-polls window. A peer that dies while |
| 177 | # poll() is running is not seen until the next _check_monitors(), and |
| 178 | # one that dies silently not until the ZMQ heartbeat expires; both |
| 179 | # still take a message into a doomed connection and are out of scope. |
| 180 | self._sweep_dead_connections() |
| 181 | |
| 182 | self._recv_router() |
| 183 | self._check_monitors() |
| 184 | self._sweep_dead_connections() |
| 185 | |
| 186 | # Create connections for new inbound peers |
| 187 | new_connections: list[ControlConnection] | None = None |
| 188 | for sender_id, msg in self._pending_inbound: |
| 189 | conn = self._connections.get(sender_id) |
| 190 | if conn is None: |
| 191 | logger.info( |
| 192 | "ZmqTransport %s: accepting INBOUND connection from %s", |
| 193 | self._local_id, |
| 194 | sender_id, |
| 195 | ) |
| 196 | conn = self._open_connection(sender_id, direction="inbound") |
| 197 | if new_connections is None: |
| 198 | new_connections = [] |
| 199 | new_connections.append(conn) |
| 200 | conn.enqueue(msg) |
| 201 | self._pending_inbound.clear() |
| 202 | |
| 203 | return ( |
| 204 | new_connections if new_connections is not None else _EMPTY_NEW_CONNECTIONS |
| 205 | ) |
| 206 | |
| 207 | def close(self) -> None: |
| 208 | if self._closed: |
| 209 | return |
| 210 | self._closed = True |
| 211 | |
| 212 | for conn in self._connections.values(): |
| 213 | conn.close() |
| 214 | self._connections.clear() |
| 215 | |
| 216 | self._router.setsockopt(zmq.LINGER, 0) |
| 217 | self._router.close() |
| 218 | self._zmq_ctx.destroy(linger=0) |
| 219 | |
| 220 | # ------------------------------------------------------------------ |
| 221 | # Internal |
| 222 | # ------------------------------------------------------------------ |
| 223 | |
| 224 | def _open_connection( |
| 225 | self, peer_id: str, direction: str = "outbound" |
| 226 | ) -> ZmqConnection: |
| 227 | """Create a DEALER socket + monitor and register the connection.""" |
| 228 | host, port_str = peer_id.rsplit(":", 1) |
| 229 | dealer_addr = _tcp_addr(host, port_str) |
| 230 | |
| 231 | logger.debug( |
| 232 | "ZmqTransport %s: creating DEALER for %s peer %s -> %s", |
| 233 | self._local_id, |
| 234 | direction, |
| 235 | peer_id, |
| 236 | dealer_addr, |
| 237 | ) |
| 238 | |
| 239 | dealer = self._zmq_ctx.socket(zmq.DEALER) |
| 240 | _apply_heartbeat(dealer) |
| 241 | dealer.identity = self._local_id.encode() |
| 242 | |
| 243 | # Unique per connection, not per peer: libzmq releases an inproc |
| 244 | # endpoint on its reaper thread after the DEALER's close() has already |
| 245 | # returned, so reusing the peer-derived address on a reconnect races |
| 246 | # that teardown and fails with EADDRINUSE. |
| 247 | safe_id = peer_id.replace(":", "-").replace("/", "-") |
| 248 | monitor_addr = f"inproc://p2p-monitor-{safe_id}-{self._monitor_seq}" |
| 249 | self._monitor_seq += 1 |
| 250 | dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED) |
| 251 | |
| 252 | monitor_sock = self._zmq_ctx.socket(zmq.PAIR) |
| 253 | monitor_sock.connect(monitor_addr) |
| 254 | |
| 255 | dealer.connect(dealer_addr) |
| 256 | |
| 257 | sockets = _Sockets(dealer=dealer, monitor=monitor_sock) |
| 258 | conn = ZmqConnection(peer_id, sockets) |
| 259 | self._connections[peer_id] = conn |
| 260 | logger.info( |
| 261 | "ZmqTransport %s: %s connection established to %s (active connections: %d)", |
| 262 | self._local_id, |
| 263 | direction, |
| 264 | peer_id, |
| 265 | len(self._connections), |
| 266 | ) |
| 267 | return conn |
| 268 | |
| 269 | def _sweep_dead_connections(self) -> None: |
| 270 | """Unregister and release every connection that is no longer alive.""" |
| 271 | for pid in [p for p, c in self._connections.items() if not c.alive]: |
| 272 | self._connections.pop(pid).close() |
| 273 | |
| 274 | def _recv_router(self) -> None: |
| 275 | """Non-blocking: receive all pending messages from ROUTER.""" |
| 276 | while True: |
| 277 | try: |
| 278 | frames = self._router.recv_multipart(zmq.NOBLOCK) |
| 279 | except zmq.Again: |
| 280 | break |
| 281 | except zmq.ZMQError as exc: |
| 282 | logger.warning("ZmqTransport %s: recv error: %s", self._local_id, exc) |
| 283 | break |
| 284 | |
| 285 | if len(frames) != 2: |
| 286 | logger.warning( |
| 287 | "ZmqTransport %s: dropping message with %d frames (expected 2)", |
| 288 | self._local_id, |
| 289 | len(frames), |
| 290 | ) |
| 291 | continue |
| 292 | |
| 293 | identity, data = frames |
| 294 | sender_id = identity.decode() |
| 295 | |
| 296 | logger.debug( |
| 297 | "ZmqTransport %s: ROUTER recv from %s (%d bytes)", |
| 298 | self._local_id, |
| 299 | sender_id, |
| 300 | len(data), |
| 301 | ) |
| 302 | |
| 303 | try: |
| 304 | msg = msgspec.msgpack.decode(data) |
| 305 | except Exception as exc: |
| 306 | logger.warning( |
| 307 | "ZmqTransport %s: failed to decode message from %s: %s", |
| 308 | self._local_id, |
| 309 | sender_id, |
| 310 | exc, |
| 311 | ) |
| 312 | continue |
| 313 | |
| 314 | conn = self._connections.get(sender_id) |
| 315 | |
| 316 | if conn is not None: |
| 317 | conn.enqueue(msg) |
| 318 | else: |
| 319 | self._pending_inbound.append((sender_id, msg)) |
| 320 | |
| 321 | def _check_monitors(self) -> None: |
| 322 | """Non-blocking: check all monitor sockets for disconnection.""" |
| 323 | for conn in self._connections.values(): |
| 324 | if not conn.alive: |
| 325 | continue |
| 326 | try: |
| 327 | event = zmq.utils.monitor.recv_monitor_message( |
| 328 | conn.monitor_socket, zmq.NOBLOCK |
| 329 | ) |
| 330 | except zmq.Again: |
| 331 | continue |
| 332 | except zmq.ZMQError as exc: |
| 333 | logger.warning( |
| 334 | "ZmqTransport %s: monitor error for peer %s: %s", |
| 335 | self._local_id, |
| 336 | conn.peer_id, |
| 337 | exc, |
| 338 | ) |
| 339 | continue |
| 340 | |
| 341 | if event["event"] == zmq.EVENT_DISCONNECTED: |
| 342 | logger.debug( |
| 343 | "ZmqTransport %s: peer %s disconnected", |
| 344 | self._local_id, |
| 345 | conn.peer_id, |
| 346 | ) |
| 347 | conn.mark_dead() |