| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ZMQ control plane for Mooncake encoder-cache reservations and events. |
| 4 | |
| 5 | Tensor bytes never travel through this module. It coordinates destination |
| 6 | reservations, completion/cancellation, TP-shard discovery, and readiness |
| 7 | notifications while :mod:`transfer` owns the Mooncake data plane. |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | import threading |
| 13 | import time |
| 14 | from collections import deque |
| 15 | from collections.abc import Callable |
| 16 | from concurrent.futures import Executor, Future |
| 17 | from typing import Any |
| 18 | |
| 19 | import torch |
| 20 | import zmq |
| 21 | |
| 22 | from vllm.logger import init_logger |
| 23 | from vllm.utils.network_utils import make_zmq_path |
| 24 | |
| 25 | logger = init_logger(__name__) |
| 26 | |
| 27 | _MAX_PENDING_EVENTS = 4096 |
| 28 | _RESERVATION_REAP_INTERVAL_SECONDS = 1 |
| 29 | |
| 30 | |
| 31 | ControlRequest = dict[str, Any] |
| 32 | ControlResponse = dict[str, Any] |
| 33 | |
| 34 | |
| 35 | class ControlClient: |
| 36 | """Send control requests through reusable, thread-local REQ sockets.""" |
| 37 | |
| 38 | def __init__(self, timeout_ms: int) -> None: |
| 39 | self._context = zmq.Context() |
| 40 | self._timeout_ms = timeout_ms |
| 41 | self._local = threading.local() |
| 42 | self._topologies: dict[str, list[str]] = {} |
| 43 | self._closed = False |
| 44 | |
| 45 | def _sockets(self) -> dict[str, zmq.Socket]: |
| 46 | sockets = getattr(self._local, "sockets", None) |
| 47 | if sockets is None: |
| 48 | sockets = {} |
| 49 | self._local.sockets = sockets |
| 50 | return sockets |
| 51 | |
| 52 | def _discard(self, addr: str) -> None: |
| 53 | socket = self._sockets().pop(addr, None) |
| 54 | if socket is not None: |
| 55 | socket.close(linger=0) |
| 56 | |
| 57 | def _exchange(self, addr: str, payload: ControlRequest) -> ControlResponse: |
| 58 | sockets = self._sockets() |
| 59 | socket = sockets.get(addr) |
| 60 | if socket is None: |
| 61 | socket = self._context.socket(zmq.REQ) |
| 62 | socket.setsockopt(zmq.IPV6, 1) |
| 63 | socket.setsockopt(zmq.RCVTIMEO, self._timeout_ms) |
| 64 | socket.setsockopt(zmq.SNDTIMEO, self._timeout_ms) |
| 65 | socket.setsockopt(zmq.LINGER, 0) |
| 66 | socket.connect(addr) |
| 67 | sockets[addr] = socket |
| 68 | try: |
| 69 | socket.send_json(payload) |
| 70 | response = socket.recv_json() |
| 71 | except Exception: |
| 72 | self._discard(addr) |
| 73 | raise |
| 74 | assert isinstance(response, dict) |
| 75 | return response |
| 76 | |
| 77 | def request(self, addr: str, payload: ControlRequest) -> Any: |
| 78 | response = self._exchange(addr, payload) |
| 79 | if not response.get("ok"): |
| 80 | raise RuntimeError(response.get("error", "EC control request failed")) |
| 81 | return response.get("result") |
| 82 | |
| 83 | def discover_shards(self, base_addr: str) -> list[str] | None: |
| 84 | if base_addr in self._topologies: |
| 85 | return self._topologies[base_addr] |
| 86 | try: |
| 87 | reply = self.request(base_addr, {"op": "peers"}) |
| 88 | ports = reply.get("ports") if isinstance(reply, dict) else None |
| 89 | if not isinstance(ports, list) or not ports: |
| 90 | raise ValueError("invalid or empty peer list") |
| 91 | prefix = base_addr.rsplit(":", 1)[0] |
| 92 | shards = [f"{prefix}:{int(port)}" for port in ports] |
| 93 | except Exception: |
| 94 | logger.warning( |
| 95 | "EC Mooncake consumer at %s did not report its shards", |
| 96 | base_addr, |
| 97 | exc_info=True, |
| 98 | ) |
| 99 | return None |
| 100 | self._topologies[base_addr] = shards |
| 101 | return shards |
| 102 | |
| 103 | def close(self) -> None: |
| 104 | if self._closed: |
| 105 | return |
| 106 | self._closed = True |
| 107 | self._context.destroy(linger=0) |
| 108 | |
| 109 | |
| 110 | def make_cancel_request( |
| 111 | transfer_id: str, |
| 112 | reservation_id: str = "", |
| 113 | *, |
| 114 | abandon: bool = False, |
| 115 | refresh: bool = False, |
| 116 | ) -> ControlRequest: |
| 117 | request: ControlRequest = { |
| 118 | "op": "cancel", |
| 119 | "transfer_id": transfer_id, |
| 120 | "reservation_id": reservation_id, |
| 121 | } |
| 122 | if abandon: |
| 123 | request["abandon"] = True |
| 124 | if refresh: |
| 125 | request["refresh"] = True |
| 126 | return request |
| 127 | |
| 128 | |
| 129 | class EventInbox: |
| 130 | """Receive Consumer readiness events without blocking the Scheduler. |
| 131 | |
| 132 | Subscribing needs one blocking control request per shard, so it runs on |
| 133 | `executor` when one is given: an unreachable Consumer must cost the |
| 134 | Scheduler nothing, not one control timeout per drain. |
| 135 | """ |
| 136 | |
| 137 | def __init__(self, client: ControlClient, executor: Executor | None = None) -> None: |
| 138 | self._client = client |
| 139 | self._executor = executor |
| 140 | self._discovery: Future[list[str] | None] | None = None |
| 141 | self._context: zmq.Context | None = None |
| 142 | self._socket: zmq.Socket | None = None |
| 143 | self._closed = False |
| 144 | self.shard_count = 1 |
| 145 | |
| 146 | def _endpoints(self, base_addr: str) -> list[str] | None: |
| 147 | """Ask every shard where it publishes readiness events. |
| 148 | |
| 149 | Blocking: called on `self._executor` unless there is none. |
| 150 | """ |
| 151 | shards = self._client.discover_shards(base_addr) |
| 152 | if shards is None: |
| 153 | return None |
| 154 | endpoints = [] |
| 155 | for addr in shards: |
| 156 | try: |
| 157 | event_port = self._client.request(addr, {"op": "event_port"}) |
| 158 | address, _ = addr.rsplit(":", 1) |
| 159 | endpoints.append(f"{address}:{int(event_port)}") |
| 160 | except Exception: |
| 161 | logger.warning( |
| 162 | "EC Mooncake could not subscribe to the event channel of " |
| 163 | "consumer shard %s; retrying the complete topology later.", |
| 164 | addr, |
| 165 | ) |
| 166 | return None |
| 167 | return endpoints |
| 168 | |
| 169 | def _connect(self, base_addr: str) -> None: |
| 170 | if self._socket is not None or self._closed: |
| 171 | return |
| 172 | if self._executor is None: |
| 173 | self._install(self._endpoints(base_addr)) |
| 174 | return |
| 175 | if self._discovery is None: |
| 176 | self._discovery = self._executor.submit(self._endpoints, base_addr) |
| 177 | return |
| 178 | if not self._discovery.done(): |
| 179 | return |
| 180 | discovery, self._discovery = self._discovery, None |
| 181 | try: |
| 182 | self._install(discovery.result()) |
| 183 | except Exception: |
| 184 | logger.warning( |
| 185 | "EC Mooncake event-channel discovery for %s failed; retrying.", |
| 186 | base_addr, |
| 187 | exc_info=True, |
| 188 | ) |
| 189 | |
| 190 | def _install(self, endpoints: list[str] | None) -> None: |
| 191 | """Adopt a discovered topology. |
| 192 | |
| 193 | Runs on the caller's thread so the socket is only ever touched there. |
| 194 | """ |
| 195 | if not endpoints or self._closed: |
| 196 | return |
| 197 | context = zmq.Context() |
| 198 | socket = context.socket(zmq.PULL) |
| 199 | socket.setsockopt(zmq.IPV6, 1) |
| 200 | for endpoint in endpoints: |
| 201 | socket.connect(endpoint) |
| 202 | self._context = context |
| 203 | self._socket = socket |
| 204 | self.shard_count = len(endpoints) |
| 205 | |
| 206 | def drain(self, base_addr: str) -> list[dict[str, Any]]: |
| 207 | self._connect(base_addr) |
| 208 | if self._socket is None: |
| 209 | return [] |
| 210 | events = [] |
| 211 | while True: |
| 212 | try: |
| 213 | events.append(self._socket.recv_json(flags=zmq.DONTWAIT)) |
| 214 | except zmq.Again: |
| 215 | return events |
| 216 | except zmq.ZMQError: |
| 217 | logger.warning("EC Mooncake event channel failed", exc_info=True) |
| 218 | return events |
| 219 | except (ValueError, UnicodeError): |
| 220 | # The event channel is a plain PULL socket: an undecodable |
| 221 | # frame must cost one frame, not the engine. |
| 222 | logger.warning( |
| 223 | "Discarding an undecodable EC Mooncake event.", exc_info=True |
| 224 | ) |
| 225 | |
| 226 | def close(self) -> None: |
| 227 | if self._closed: |
| 228 | return |
| 229 | self._closed = True |
| 230 | if self._discovery is not None: |
| 231 | self._discovery.cancel() |
| 232 | self._discovery = None |
| 233 | if self._socket is not None: |
| 234 | self._socket.close(linger=0) |
| 235 | self._socket = None |
| 236 | if self._context is not None: |
| 237 | self._context.term() |
| 238 | |
| 239 | |
| 240 | class ConsumerControlServer: |
| 241 | """Expose Consumer reservations and readiness events over ZMQ. |
| 242 | |
| 243 | One server runs on every receiving TP rank. The REP channel handles |
| 244 | reservation operations, while a PUSH channel publishes newly ready items |
| 245 | to the Scheduler. |
| 246 | """ |
| 247 | |
| 248 | def __init__( |
| 249 | self, |
| 250 | host: str, |
| 251 | port: int, |
| 252 | reserve: Callable[[dict[str, Any]], dict[str, Any]], |
| 253 | status: Callable[[str], dict[str, Any] | None], |
| 254 | complete: Callable[[str, str], tuple[bool, bool]], |
| 255 | cancel: Callable[[str, str, bool, bool], bool], |
| 256 | reap: Callable[[], int], |
| 257 | peer_ports: list[int] | None = None, |
| 258 | device: torch.device | None = None, |
| 259 | drain_ready: Callable[[], list[str]] = lambda: [], |
| 260 | ) -> None: |
| 261 | self.host = host |
| 262 | self.port = port |
| 263 | self.peer_ports = peer_ports or [port] |
| 264 | self._device = device |
| 265 | self.event_port: int | None = None |
| 266 | self._reserve = reserve |
| 267 | self._status = status |
| 268 | self._complete = complete |
| 269 | self._cancel = cancel |
| 270 | self._reap = reap |
| 271 | self._drain_ready = drain_ready |
| 272 | self._stop = threading.Event() |
| 273 | self._started = threading.Event() |
| 274 | self._thread: threading.Thread | None = None |
| 275 | self._startup_error: Exception | None = None |
| 276 | |
| 277 | def start(self) -> None: |
| 278 | def loop() -> None: |
| 279 | if self._device is not None and self._device.type == "cuda": |
| 280 | torch.accelerator.set_device_index(self._device.index or 0) |
| 281 | context = zmq.Context() |
| 282 | socket = context.socket(zmq.REP) |
| 283 | event_socket = context.socket(zmq.PUSH) |
| 284 | socket.setsockopt(zmq.IPV6, 1) |
| 285 | event_socket.setsockopt(zmq.IPV6, 1) |
| 286 | pending_events: deque[dict[str, Any]] = deque() |
| 287 | |
| 288 | def queue_event(event: dict[str, Any]) -> None: |
| 289 | event["shard"] = self.port |
| 290 | if len(pending_events) >= _MAX_PENDING_EVENTS: |
| 291 | dropped = pending_events.popleft() |
| 292 | logger.warning( |
| 293 | "EC Mooncake event backlog full on port %d; dropping " |
| 294 | "readiness for transfer_id=%s", |
| 295 | self.port, |
| 296 | dropped.get("transfer_id"), |
| 297 | ) |
| 298 | pending_events.append(event) |
| 299 | |
| 300 | def queue_ready(transfer_id: str) -> None: |
| 301 | status = self._status(transfer_id) |
| 302 | if status is not None: |
| 303 | queue_event({"transfer_id": transfer_id, **status}) |
| 304 | |
| 305 | last_reap_at = time.monotonic() |
| 306 | socket.setsockopt(zmq.RCVTIMEO, 100) |
| 307 | try: |
| 308 | socket.bind(make_zmq_path("tcp", self.host, self.port)) |
| 309 | event_socket.bind(make_zmq_path("tcp", self.host, 0)) |
| 310 | self.event_port = int( |
| 311 | event_socket.getsockopt(zmq.LAST_ENDPOINT) |
| 312 | .decode() |
| 313 | .rsplit(":", 1)[1] |
| 314 | ) |
| 315 | except Exception as e: |
| 316 | self._startup_error = e |
| 317 | self._started.set() |
| 318 | socket.close(linger=0) |
| 319 | event_socket.close(linger=0) |
| 320 | context.term() |
| 321 | return |
| 322 | self._started.set() |
| 323 | try: |
| 324 | while not self._stop.is_set(): |
| 325 | while pending_events: |
| 326 | try: |
| 327 | event_socket.send_json( |
| 328 | pending_events[0], flags=zmq.DONTWAIT |
| 329 | ) |
| 330 | except zmq.Again: |
| 331 | break |
| 332 | pending_events.popleft() |
| 333 | now = time.monotonic() |
| 334 | if now - last_reap_at >= _RESERVATION_REAP_INTERVAL_SECONDS: |
| 335 | self._reap() |
| 336 | last_reap_at = now |
| 337 | try: |
| 338 | request = socket.recv_json() |
| 339 | except zmq.Again: |
| 340 | continue |
| 341 | except Exception: |
| 342 | # The frame arrived but did not decode. REP still owes a |
| 343 | # reply, so answer before returning to the loop. |
| 344 | logger.exception( |
| 345 | "EC Mooncake control channel on port %d received an " |
| 346 | "undecodable request", |
| 347 | self.port, |
| 348 | ) |
| 349 | socket.send_json( |
| 350 | {"ok": False, "error": "malformed control request"} |
| 351 | ) |
| 352 | continue |
| 353 | try: |
| 354 | op = request.get("op") |
| 355 | result: Any = None |
| 356 | if op in ("reserve", "reserve_batch"): |
| 357 | items = ( |
| 358 | request["items"] if op == "reserve_batch" else [request] |
| 359 | ) |
| 360 | results = [] |
| 361 | for item in items: |
| 362 | try: |
| 363 | reserved = self._reserve(item) |
| 364 | results.append({"ok": True, "result": reserved}) |
| 365 | if reserved.get("ready"): |
| 366 | queue_ready(str(item["transfer_id"])) |
| 367 | except Exception as exc: |
| 368 | results.append({"ok": False, "error": str(exc)}) |
| 369 | if op == "reserve_batch": |
| 370 | result = {"items": results} |
| 371 | elif not results[0]["ok"]: |
| 372 | raise RuntimeError(results[0]["error"]) |
| 373 | else: |
| 374 | result = results[0]["result"] |
| 375 | elif op == "status": |
| 376 | result = self._status(str(request["transfer_id"])) |
| 377 | elif op == "event_port": |
| 378 | result = self.event_port |
| 379 | elif op == "peers": |
| 380 | result = {"ports": self.peer_ports} |
| 381 | elif op in ("complete", "complete_batch"): |
| 382 | items = ( |
| 383 | request["items"] |
| 384 | if op == "complete_batch" |
| 385 | else [request] |
| 386 | ) |
| 387 | completions = [] |
| 388 | for item in items: |
| 389 | transfer_id = str(item["transfer_id"]) |
| 390 | accepted, became_ready = self._complete( |
| 391 | transfer_id, str(item["reservation_id"]) |
| 392 | ) |
| 393 | for ready_id in self._drain_ready(): |
| 394 | queue_ready(ready_id) |
| 395 | completions.append( |
| 396 | { |
| 397 | "completed": accepted, |
| 398 | "became_ready": became_ready, |
| 399 | } |
| 400 | ) |
| 401 | if not became_ready: |
| 402 | continue |
| 403 | queue_ready(transfer_id) |
| 404 | result = ( |
| 405 | {"items": completions} |
| 406 | if op == "complete_batch" |
| 407 | else completions[0] |
| 408 | ) |
| 409 | elif op == "cancel": |
| 410 | result = { |
| 411 | "cancelled": self._cancel( |
| 412 | str(request["transfer_id"]), |
| 413 | str(request.get("reservation_id", "")), |
| 414 | bool(request.get("abandon", False)), |
| 415 | bool(request.get("refresh", False)), |
| 416 | ) |
| 417 | } |
| 418 | else: |
| 419 | raise ValueError(f"unknown control op: {op!r}") |
| 420 | socket.send_json({"ok": True, "result": result}) |
| 421 | except Exception as e: |
| 422 | socket.send_json({"ok": False, "error": str(e)}) |
| 423 | except Exception: |
| 424 | # `finally` closes the sockets and the thread ends, so without |
| 425 | # this every later reserve against this shard would surface |
| 426 | # only as a control timeout with nothing to attribute it to. |
| 427 | logger.exception( |
| 428 | "EC Mooncake control channel on port %d stopped serving", |
| 429 | self.port, |
| 430 | ) |
| 431 | raise |
| 432 | finally: |
| 433 | socket.close(linger=0) |
| 434 | event_socket.close(linger=0) |
| 435 | context.term() |
| 436 | |
| 437 | self._thread = threading.Thread( |
| 438 | target=loop, name="ec-mooncake-control", daemon=True |
| 439 | ) |
| 440 | self._thread.start() |
| 441 | if not self._started.wait(timeout=5): |
| 442 | raise RuntimeError("EC Mooncake control channel failed to start") |
| 443 | if self._startup_error is not None: |
| 444 | raise RuntimeError("EC Mooncake control channel failed to bind") from ( |
| 445 | self._startup_error |
| 446 | ) |
| 447 | logger.info( |
| 448 | "EC Mooncake control channel listening on tcp://%s:%d (events tcp://%s:%d)", |
| 449 | self.host, |
| 450 | self.port, |
| 451 | self.host, |
| 452 | self.event_port, |
| 453 | ) |
| 454 | |
| 455 | def close(self) -> None: |
| 456 | if self._thread is None: |
| 457 | return |
| 458 | self._stop.set() |
| 459 | self._thread.join() |
| 460 | self._thread = None |