| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ZMQ control-plane implementations for ECCPUConnector. |
| 4 | |
| 5 | ZmqClientConnection — one DEALER socket to a single producer peer. |
| 6 | ZmqClientTransport — pool of ZmqClientConnection objects (consumer side). |
| 7 | ZmqServerTransport — ROUTER socket (producer side). |
| 8 | """ |
| 9 | |
| 10 | import contextlib |
| 11 | from dataclasses import dataclass, field |
| 12 | |
| 13 | import zmq |
| 14 | from zmq.utils.monitor import recv_monitor_message |
| 15 | |
| 16 | from vllm.distributed.ec_transfer.ec_connector.cpu.control.base import ( |
| 17 | ControlConnection, |
| 18 | ) |
| 19 | from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.utils.network_utils import make_zmq_path, make_zmq_socket |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | _HEARTBEAT_IVL_MS = 2000 |
| 26 | _HEARTBEAT_TIMEOUT_MS = 4000 |
| 27 | _HEARTBEAT_TTL_MS = 8000 |
| 28 | |
| 29 | |
| 30 | # ── Client side ─────────────────────────────────────────────────────────────── |
| 31 | |
| 32 | |
| 33 | @dataclass |
| 34 | class ZmqClientConnection(ControlConnection): |
| 35 | """One DEALER socket to one producer peer. |
| 36 | |
| 37 | Owns the socket and its optional monitor socket. NIXL state is NOT |
| 38 | stored here — it belongs in ECSession above this layer. |
| 39 | """ |
| 40 | |
| 41 | dealer: zmq.Socket |
| 42 | monitor: zmq.Socket | None |
| 43 | _alive: bool = field(default=True, init=False) |
| 44 | |
| 45 | @property |
| 46 | def alive(self) -> bool: |
| 47 | return self._alive |
| 48 | |
| 49 | def send(self, msg: bytes) -> None: |
| 50 | self.dealer.send_multipart([b"", msg]) |
| 51 | |
| 52 | def recv(self) -> list[bytes]: |
| 53 | msgs: list[bytes] = [] |
| 54 | while True: |
| 55 | try: |
| 56 | frames = self.dealer.recv_multipart(flags=zmq.NOBLOCK) |
| 57 | except zmq.Again: |
| 58 | break |
| 59 | except Exception: |
| 60 | logger.exception("ec: DEALER recv failed") |
| 61 | break |
| 62 | if len(frames) != 2 or frames[0] != b"": |
| 63 | logger.warning( |
| 64 | "ec: dropped malformed frame envelope (%d frames)", len(frames) |
| 65 | ) |
| 66 | continue |
| 67 | msgs.append(frames[1]) |
| 68 | return msgs |
| 69 | |
| 70 | def mark_dead(self) -> None: |
| 71 | self._alive = False |
| 72 | |
| 73 | def close(self) -> None: |
| 74 | if self.monitor is not None: |
| 75 | with contextlib.suppress(Exception): |
| 76 | self.dealer.disable_monitor() |
| 77 | with contextlib.suppress(Exception): |
| 78 | self.monitor.close(linger=0) |
| 79 | with contextlib.suppress(Exception): |
| 80 | self.dealer.close(linger=0) |
| 81 | self._alive = False |
| 82 | |
| 83 | |
| 84 | class ZmqClientTransport: |
| 85 | """Pool of ZmqClientConnection objects — one per producer peer. |
| 86 | |
| 87 | Creates DEALER sockets lazily on first contact. Does not touch NIXL; |
| 88 | NIXL registration is handled by ECCPUConsumer once an XferAck arrives. |
| 89 | """ |
| 90 | |
| 91 | def __init__(self) -> None: |
| 92 | self._ctx = zmq.Context() |
| 93 | self._connections: dict[PeerAddr, ZmqClientConnection] = {} |
| 94 | |
| 95 | def connect(self, addr: PeerAddr) -> ZmqClientConnection: |
| 96 | """Return the existing connection for addr, or create a new one.""" |
| 97 | existing = self._connections.get(addr) |
| 98 | if existing is not None: |
| 99 | return existing |
| 100 | |
| 101 | host, port = addr |
| 102 | dealer = make_zmq_socket( |
| 103 | ctx=self._ctx, |
| 104 | path=make_zmq_path(scheme="tcp", host=host, port=port), |
| 105 | socket_type=zmq.DEALER, |
| 106 | bind=False, |
| 107 | ) |
| 108 | dealer.setsockopt(zmq.HEARTBEAT_IVL, _HEARTBEAT_IVL_MS) |
| 109 | dealer.setsockopt(zmq.HEARTBEAT_TIMEOUT, _HEARTBEAT_TIMEOUT_MS) |
| 110 | dealer.setsockopt(zmq.HEARTBEAT_TTL, _HEARTBEAT_TTL_MS) |
| 111 | |
| 112 | monitor_addr = f"inproc://ec-peer-mon-{host}-{port}" |
| 113 | dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED) |
| 114 | monitor = self._ctx.socket(zmq.PAIR) |
| 115 | monitor.connect(monitor_addr) |
| 116 | |
| 117 | conn = ZmqClientConnection(dealer=dealer, monitor=monitor) |
| 118 | self._connections[addr] = conn |
| 119 | return conn |
| 120 | |
| 121 | def poll(self) -> dict[PeerAddr, list[bytes]]: |
| 122 | """Non-blocking drain of messages from every peer's DEALER.""" |
| 123 | msgs: dict[PeerAddr, list[bytes]] = {} |
| 124 | for addr, conn in list(self._connections.items()): |
| 125 | conn_msgs = conn.recv() |
| 126 | msgs[addr] = conn_msgs |
| 127 | return msgs |
| 128 | |
| 129 | def poll_dead(self) -> list[PeerAddr]: |
| 130 | """Non-blocking sweep of peer monitors; return addresses that died.""" |
| 131 | dead: list[PeerAddr] = [] |
| 132 | for addr, conn in list(self._connections.items()): |
| 133 | if conn.monitor is None: |
| 134 | continue |
| 135 | try: |
| 136 | while True: |
| 137 | evt = recv_monitor_message(conn.monitor, flags=zmq.NOBLOCK) |
| 138 | if evt["event"] == zmq.EVENT_DISCONNECTED: |
| 139 | conn.mark_dead() |
| 140 | dead.append(addr) |
| 141 | break |
| 142 | except zmq.Again: |
| 143 | pass |
| 144 | except Exception: |
| 145 | logger.warning("ec: monitor poll failed for %s", addr, exc_info=True) |
| 146 | return dead |
| 147 | |
| 148 | def remove(self, addr: PeerAddr) -> ZmqClientConnection | None: |
| 149 | """Remove and close the connection for addr, returning it (or None).""" |
| 150 | conn = self._connections.pop(addr, None) |
| 151 | if conn is not None: |
| 152 | conn.close() |
| 153 | return conn |
| 154 | |
| 155 | def close(self) -> None: |
| 156 | for conn in list(self._connections.values()): |
| 157 | conn.close() |
| 158 | self._connections.clear() |
| 159 | try: |
| 160 | self._ctx.destroy(linger=0) |
| 161 | except Exception: |
| 162 | logger.debug("ec: zmq context destroy failed", exc_info=True) |
| 163 | |
| 164 | |
| 165 | # ── Server side ─────────────────────────────────────────────────────────────── |
| 166 | |
| 167 | |
| 168 | class ZmqServerTransport: |
| 169 | """ROUTER socket for the producer side. Byte-level only, driven by |
| 170 | ProducerSession. |
| 171 | |
| 172 | poll(timeout_ms) waits up to timeout_ms for messages, then drains all that |
| 173 | are available; timeout_ms=0 drains whatever has already arrived. send() |
| 174 | routes a reply back to a specific peer by ZMQ identity. |
| 175 | """ |
| 176 | |
| 177 | def __init__(self, host: str, port: int) -> None: |
| 178 | self._ctx = zmq.Context() |
| 179 | self._router = make_zmq_socket( |
| 180 | ctx=self._ctx, |
| 181 | path=make_zmq_path(scheme="tcp", host=host, port=port), |
| 182 | socket_type=zmq.ROUTER, |
| 183 | bind=True, |
| 184 | ) |
| 185 | self._poller = zmq.Poller() |
| 186 | self._poller.register(self._router, zmq.POLLIN) |
| 187 | |
| 188 | def poll(self, timeout_ms: int = 0) -> list[tuple[bytes, bytes]]: |
| 189 | """Wait up to timeout_ms for incoming messages, then drain all pending. |
| 190 | |
| 191 | Returns a list of (identity, payload) pairs. |
| 192 | """ |
| 193 | if timeout_ms > 0: |
| 194 | try: |
| 195 | events = dict(self._poller.poll(timeout=timeout_ms)) |
| 196 | except zmq.ContextTerminated: |
| 197 | return [] |
| 198 | except Exception: |
| 199 | logger.exception("ec: router poller failed") |
| 200 | return [] |
| 201 | if self._router not in events: |
| 202 | return [] |
| 203 | |
| 204 | msgs: list[tuple[bytes, bytes]] = [] |
| 205 | while True: |
| 206 | try: |
| 207 | identity, _, payload = self._router.recv_multipart(flags=zmq.NOBLOCK) |
| 208 | msgs.append((identity, payload)) |
| 209 | except zmq.Again: |
| 210 | break |
| 211 | except zmq.ContextTerminated: |
| 212 | break |
| 213 | except Exception: |
| 214 | logger.exception("ec: router recv failed") |
| 215 | break |
| 216 | return msgs |
| 217 | |
| 218 | def send(self, identity: bytes, payload: bytes) -> None: |
| 219 | """Send a reply payload to the peer identified by ZMQ identity.""" |
| 220 | try: |
| 221 | self._router.send_multipart([identity, b"", payload]) |
| 222 | except Exception: |
| 223 | logger.exception("ec: router send failed") |
| 224 | |
| 225 | def close(self) -> None: |
| 226 | try: |
| 227 | self._router.close(linger=0) |
| 228 | except Exception: |
| 229 | logger.debug("ec: router close failed", exc_info=True) |
| 230 | try: |
| 231 | self._ctx.destroy(linger=0) |
| 232 | except Exception: |
| 233 | logger.debug("ec: zmq context destroy failed", exc_info=True) |