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.
        """
        ...

    @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.
        """
        ...

    @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
        self._closed = False
        self._inbox: list[dict] = []

    def send(self, msg: dict) -> None:
        """Send a msgpack-encoded message to this peer."""
        if self._closed:
            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._closed

    def close(self) -> None:
        if self._closed:
            return
        self._closed = 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."""
        self._closed = 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]] = []

        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."""
        assert peer_id not in self._connections, (
            f"ZmqConnection to {peer_id} already exists"
        )
        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
        """
        self._recv_router()
        self._check_monitors()

        # 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()

        # Remove dead connections
        for pid in [p for p, c in self._connections.items() if not c.alive]:
            self._connections.pop(pid).close()

        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()

        safe_id = peer_id.replace(":", "-").replace("/", "-")
        monitor_addr = f"inproc://p2p-monitor-{safe_id}"
        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 _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()