vllm.v1.kv_offload.tiering.p2p.control ¶

class ControlConnection(ABC):
    """Bidirectional message channel to a single remote peer.

    Lifecycle:
        1. Created by ControlTransport (connect() or poll())
        2. Used for send/recv by sessions
        3. Marked dead on disconnect (mark_dead())
        4. Cleaned up with close()

    Once mark_dead() is called, alive becomes False and the owning
    session should stop using this connection. close() releases
    underlying resources (sockets, monitors).
    """

    def __init__(self, peer_id: str) -> None:
        self.peer_id = peer_id

    @property
    @abstractmethod
    def alive(self) -> bool:
        """True if the connection is usable. False after mark_dead/close."""
        ...

    @abstractmethod
    def send(self, msg: dict) -> None:
        """Enqueue a message for delivery to the peer.

        Must not block. Serialization happens internally.
        Raises on closed connection.
        """
        ...

    @abstractmethod
    def recv(self) -> Sequence[dict]:
        """Drain and return all messages received since the last call.

        Returns an empty sequence if no messages are pending.
        The returned sequence is read-only — callers must not mutate it.
        Messages are dicts deserialized from the wire format.
        """
        ...

    @abstractmethod
    def mark_dead(self) -> None:
        """Mark this connection as dead (peer disconnected).

        After this call, alive returns False. The session should
        stop using this connection and the transport will clean it up.
        Resources are not released here — close() must still run.
        """
        ...

    @abstractmethod
    def close(self) -> None:
        """Release all resources (sockets, monitors).

        Idempotent. After close(), alive returns False.
        """
        ...

class ControlTransport(ABC):
    """Manages peer connections and drives all control-plane I/O.

    Owns the listening socket and all active connections.
    The caller must invoke poll() periodically to process I/O.

    Lifecycle:
        1. Constructed with a local identity and listen address
        2. connect() to reach remote peers
        3. poll() to accept inbound peers and process messages
        4. close() to shut down
    """

    @abstractmethod
    def connect(self, peer_id: str) -> ControlConnection:
        """Create an outbound connection to a remote peer.

        Args:
            peer_id: Remote peer identity (format: "host:port").

        Returns:
            A new ControlConnection ready for send/recv.

        The connection's send queue is live immediately — messages
        sent before the remote peer's poll() will be buffered.

        A connection to peer_id that is registered but no longer alive is
        retired and replaced; a live one is a duplicate and must not be
        replaced.
        """
        ...

    @abstractmethod
    def poll(self) -> Sequence[ControlConnection]:
        """Process all pending I/O and return newly accepted connections.

        This is the main I/O driver. Each call:
          - Receives messages from all connected peers (buffered in
            each connection's recv() queue)
          - Accepts new inbound peers and creates connections for them
            (first message already in recv() buffer)
          - Detects disconnections and marks connections dead

        Returns:
            Newly accepted inbound connections (not previously returned).
            The returned sequence is read-only — callers must not mutate
            it. The caller is responsible for creating sessions for
            these connections.
        """
        ...

    @abstractmethod
    def close(self) -> None:
        """Shut down the transport and all connections.

        Closes the listening socket and all active connections.
        Idempotent.
        """
        ...

class ZmqConnection(ControlConnection):
    """Bidirectional message channel to a single remote peer."""

    def __init__(self, peer_id: str, sockets: _Sockets) -> None:
        super().__init__(peer_id)
        self._sockets = sockets
        # _dead: the peer is gone. _closed: the sockets have been released.
        # Distinct, so mark_dead() cannot turn close() into a no-op and leak
        # the DEALER and its monitor socket.
        self._dead = False
        self._closed = False
        self._inbox: list[dict] = []

    def send(self, msg: dict) -> None:
        """Send a msgpack-encoded message to this peer."""
        if not self.alive:
            raise RuntimeError(
                f"ZmqConnection: send on closed connection to {self.peer_id}"
            )
        data = msgspec.msgpack.encode(msg)
        self._sockets.dealer.send(data)

    def recv(self) -> Sequence[dict]:
        """Drain and return all buffered incoming messages."""
        if not self._inbox:
            return _EMPTY_INBOX
        msgs = self._inbox
        self._inbox = []
        return msgs

    @property
    def alive(self) -> bool:
        return not (self._dead or self._closed)

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        self._dead = True
        logger.info("ZmqConnection: closing connection to %s", self.peer_id)
        self._sockets.monitor.close()
        self._sockets.dealer.close()

    def enqueue(self, msg: dict) -> None:
        """Buffer an incoming message."""
        self._inbox.append(msg)

    def mark_dead(self) -> None:
        """Mark connection as disconnected.

        Only flips liveness: the sockets stay open until close() releases
        them, so the owner can still drain recv() before tearing down.
        """
        self._dead = True

    @property
    def monitor_socket(self) -> zmq.Socket:
        """Monitor socket for disconnect detection (used by ZmqTransport)."""
        return self._sockets.monitor

class ZmqTransport(ControlTransport):
    """ZMQ implementation of ControlTransport.

    Manages a ROUTER socket for accepting connections and DEALER sockets
    for outbound connections. Message-content agnostic.
    """

    def __init__(self, local_id: str, host: str, port: int) -> None:
        self._local_id = local_id
        self._closed = False

        self._connections: dict[str, ZmqConnection] = {}
        self._pending_inbound: list[tuple[str, dict]] = []
        # Monotonic suffix for inproc monitor endpoints — see
        # _open_connection() for why peer_id alone is not enough.
        self._monitor_seq = 0

        self._zmq_ctx = zmq.Context()
        self._router: zmq.Socket = self._zmq_ctx.socket(zmq.ROUTER)
        _apply_heartbeat(self._router)
        bind_addr = _tcp_addr(host, port)
        self._router.bind(bind_addr)
        logger.info("ZmqTransport %s: ROUTER bound on %s", self._local_id, bind_addr)

    # ------------------------------------------------------------------
    # ZmqConnection lifecycle
    # ------------------------------------------------------------------

    def connect(self, peer_id: str) -> ZmqConnection:
        """Open an outbound connection to a remote peer.

        A dead connection can still be registered: its owning session may mark
        it dead after this tick's sweep already ran, and poll() only
        unregisters it on the next pass. Retire such an entry instead of
        asserting, so a reconnect landing in that window succeeds. A live
        entry is still a genuine duplicate.
        """
        existing = self._connections.get(peer_id)
        if existing is not None:
            assert not existing.alive, f"ZmqConnection to {peer_id} already exists"
            logger.info(
                "ZmqTransport %s: retiring dead connection to %s before reconnect",
                self._local_id,
                peer_id,
            )
            self._connections.pop(peer_id).close()
        logger.info(
            "ZmqTransport %s: opening OUTBOUND connection to %s",
            self._local_id,
            peer_id,
        )
        return self._open_connection(peer_id, direction="outbound")

    def poll(self) -> Sequence[ControlConnection]:
        """Process all pending I/O. Returns newly accepted connections.

        - Receives messages (buffered in each connection's inbox)
        - Creates connections for new inbound peers (connect msg in inbox)
        - Checks monitors for disconnections
        - Removes and closes dead connections
        """
        # Retire connections a session killed since the last poll() before
        # routing traffic: otherwise a reconnecting peer's first message is
        # enqueued into the dead connection and discarded along with it.
        # This closes only that between-polls window. A peer that dies while
        # poll() is running is not seen until the next _check_monitors(), and
        # one that dies silently not until the ZMQ heartbeat expires; both
        # still take a message into a doomed connection and are out of scope.
        self._sweep_dead_connections()

        self._recv_router()
        self._check_monitors()
        self._sweep_dead_connections()

        # Create connections for new inbound peers
        new_connections: list[ControlConnection] | None = None
        for sender_id, msg in self._pending_inbound:
            conn = self._connections.get(sender_id)
            if conn is None:
                logger.info(
                    "ZmqTransport %s: accepting INBOUND connection from %s",
                    self._local_id,
                    sender_id,
                )
                conn = self._open_connection(sender_id, direction="inbound")
                if new_connections is None:
                    new_connections = []
                new_connections.append(conn)
            conn.enqueue(msg)
        self._pending_inbound.clear()

        return (
            new_connections if new_connections is not None else _EMPTY_NEW_CONNECTIONS
        )

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True

        for conn in self._connections.values():
            conn.close()
        self._connections.clear()

        self._router.setsockopt(zmq.LINGER, 0)
        self._router.close()
        self._zmq_ctx.destroy(linger=0)

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _open_connection(
        self, peer_id: str, direction: str = "outbound"
    ) -> ZmqConnection:
        """Create a DEALER socket + monitor and register the connection."""
        host, port_str = peer_id.rsplit(":", 1)
        dealer_addr = _tcp_addr(host, port_str)

        logger.debug(
            "ZmqTransport %s: creating DEALER for %s peer %s -> %s",
            self._local_id,
            direction,
            peer_id,
            dealer_addr,
        )

        dealer = self._zmq_ctx.socket(zmq.DEALER)
        _apply_heartbeat(dealer)
        dealer.identity = self._local_id.encode()

        # Unique per connection, not per peer: libzmq releases an inproc
        # endpoint on its reaper thread after the DEALER's close() has already
        # returned, so reusing the peer-derived address on a reconnect races
        # that teardown and fails with EADDRINUSE.
        safe_id = peer_id.replace(":", "-").replace("/", "-")
        monitor_addr = f"inproc://p2p-monitor-{safe_id}-{self._monitor_seq}"
        self._monitor_seq += 1
        dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED)

        monitor_sock = self._zmq_ctx.socket(zmq.PAIR)
        monitor_sock.connect(monitor_addr)

        dealer.connect(dealer_addr)

        sockets = _Sockets(dealer=dealer, monitor=monitor_sock)
        conn = ZmqConnection(peer_id, sockets)
        self._connections[peer_id] = conn
        logger.info(
            "ZmqTransport %s: %s connection established to %s (active connections: %d)",
            self._local_id,
            direction,
            peer_id,
            len(self._connections),
        )
        return conn

    def _sweep_dead_connections(self) -> None:
        """Unregister and release every connection that is no longer alive."""
        for pid in [p for p, c in self._connections.items() if not c.alive]:
            self._connections.pop(pid).close()

    def _recv_router(self) -> None:
        """Non-blocking: receive all pending messages from ROUTER."""
        while True:
            try:
                frames = self._router.recv_multipart(zmq.NOBLOCK)
            except zmq.Again:
                break
            except zmq.ZMQError as exc:
                logger.warning("ZmqTransport %s: recv error: %s", self._local_id, exc)
                break

            if len(frames) != 2:
                logger.warning(
                    "ZmqTransport %s: dropping message with %d frames (expected 2)",
                    self._local_id,
                    len(frames),
                )
                continue

            identity, data = frames
            sender_id = identity.decode()

            logger.debug(
                "ZmqTransport %s: ROUTER recv from %s (%d bytes)",
                self._local_id,
                sender_id,
                len(data),
            )

            try:
                msg = msgspec.msgpack.decode(data)
            except Exception as exc:
                logger.warning(
                    "ZmqTransport %s: failed to decode message from %s: %s",
                    self._local_id,
                    sender_id,
                    exc,
                )
                continue

            conn = self._connections.get(sender_id)

            if conn is not None:
                conn.enqueue(msg)
            else:
                self._pending_inbound.append((sender_id, msg))

    def _check_monitors(self) -> None:
        """Non-blocking: check all monitor sockets for disconnection."""
        for conn in self._connections.values():
            if not conn.alive:
                continue
            try:
                event = zmq.utils.monitor.recv_monitor_message(
                    conn.monitor_socket, zmq.NOBLOCK
                )
            except zmq.Again:
                continue
            except zmq.ZMQError as exc:
                logger.warning(
                    "ZmqTransport %s: monitor error for peer %s: %s",
                    self._local_id,
                    conn.peer_id,
                    exc,
                )
                continue

            if event["event"] == zmq.EVENT_DISCONNECTED:
                logger.debug(
                    "ZmqTransport %s: peer %s disconnected",
                    self._local_id,
                    conn.peer_id,
                )
                conn.mark_dead()