vllm.distributed.ec_transfer.ec_connector.cpu.session ¶

class ConsumerSession:
    """Long-lived per-producer session on the consumer side.

    Owns the ZMQ DEALER connection and the NIXL agent registration for one
    producer peer. The connector calls poll(messages, now) once per engine
    step with the raw bytes collected from this peer's DEALER socket.

    poll() decodes XferAcks, dispatches them to the right ConsumerXfer,
    advances all active xfer state machines, and drains quarantined xfers.
    Results accumulate in internal sets; call take_results() to collect them.
    """

    def __init__(
        self,
        addr: PeerAddr,
        zmq_conn: ZmqClientConnection,
        transport: ZmqClientTransport,
        data: DataTransport,
        compat_hash: str,
    ) -> None:
        self._addr = addr
        self._zmq = zmq_conn
        self._transport = transport
        self._data = data
        self._compat_hash = compat_hash
        self._session_id = str(
            uuid.uuid4()
        )  # stable identity for (session_id:mm_hash) keys

        self._nixl_agent_name: str | None = None
        self._nixl_metadata_bytes: bytes | None = None

        self._xfers: dict[str, ConsumerXfer] = {}  # mm_hash → xfer
        self._quarantined: list[ConsumerXfer] = []  # DMA pending, to poll
        self._newly_quarantined: set[str] = set()  # quarantined since last take_results

        self._completed: set[str] = set()
        self._tombstoned: set[str] = set()
        self._retryable: set[str] = set()
        self._settled: list[tuple[str, list[int]]] = []

        self._encoder = msgspec.msgpack.Encoder()
        self._decoder = msgspec.msgpack.Decoder(XferAck)

    @property
    def addr(self) -> PeerAddr:
        return self._addr

    def start_xfer(
        self, mm_hash: str, block_indices: list[int], deadline: float
    ) -> None:
        """Create a ConsumerXfer and send the XferReq over the connection."""
        xfer = ConsumerXfer(
            mm_hash=mm_hash,
            block_indices=block_indices,
            addr=self._addr,
            deadline=deadline,
            data=self._data,
            consumer_session_id=self._session_id,
        )
        self._zmq.send(
            self._encoder.encode(
                XferReq(
                    mm_hash=mm_hash,
                    compatibility_hash=self._compat_hash,
                    session_id=self._session_id,
                )
            )
        )
        self._xfers[mm_hash] = xfer
        logger.debug(
            "EC Connector: XferReq sent mm_hash=%s to %s:%d",
            mm_hash,
            self._addr[0],
            self._addr[1],
        )

    def poll(self, messages: list[bytes], now: float) -> None:
        """Decode inbound XferAcks, dispatch to xfers, and advance all state machines.

        messages: raw bytes from ZmqClientTransport.poll() for this peer.
        Results accumulate until take_results() is called.
        """
        for raw in messages:
            try:
                ack = self._decoder.decode(raw)
            except (msgspec.DecodeError, msgspec.ValidationError):
                logger.warning(
                    "EC Connector: dropped malformed XferAck from %s:%d",
                    self._addr[0],
                    self._addr[1],
                )
                continue
            self._dispatch_ack(ack)

        for mm_hash, xfer in list(self._xfers.items()):
            state = xfer.poll(now)
            if state == XferState.DONE:
                self._completed.add(mm_hash)
                del self._xfers[mm_hash]
            elif state == XferState.ACK_TIMEOUT:
                # The scheduler bounds retries with a per-request deadline.
                logger.warning(
                    "EC consumer: no XferAck for mm_hash=%s from %s:%d in time; "
                    "will re-request",
                    mm_hash,
                    self._addr[0],
                    self._addr[1],
                )
                self._retryable.add(mm_hash)
                del self._xfers[mm_hash]
            elif state == XferState.READ_FAILED:
                self._tombstoned.add(mm_hash)
                del self._xfers[mm_hash]
            elif state == XferState.QUARANTINED:
                # DMA may still be writing — caller must NOT free blocks.
                self._quarantined.append(xfer)
                self._newly_quarantined.add(mm_hash)
                del self._xfers[mm_hash]

        self._drain_quarantine(now)

    def _dispatch_ack(self, ack: XferAck) -> None:
        xfer = self._xfers.get(ack.mm_hash)
        if xfer is None or xfer.transfer_handle is not None:
            return  # unknown mm_hash or duplicate ack

        if ack.status != XferStatus.OK:
            logger.log(
                logging.INFO if ack.status in EXPECTED_NACKS else logging.WARNING,
                "EC Connector: NACK %s from %s:%d for mm_hash=%s",
                ack.status.name,
                self._addr[0],
                self._addr[1],
                ack.mm_hash,
            )
            del self._xfers[ack.mm_hash]
            # The producer announced this encoding and its save is still in
            # flight, so another attempt on a later step can be granted. Every
            # other NACK is terminal for this mm_hash.
            if ack.status in RETRYABLE_NACKS:
                self._retryable.add(ack.mm_hash)
            else:
                self._tombstoned.add(ack.mm_hash)
            return

        try:
            agent_name = self._ensure_registered(ack.agent_metadata, ack.mem_descriptor)
        except Exception:
            logger.exception(
                "EC Connector: failed to register peer %s:%d for mm_hash=%s",
                self._addr[0],
                self._addr[1],
                ack.mm_hash,
            )
            del self._xfers[ack.mm_hash]
            self._tombstoned.add(ack.mm_hash)
            return

        if not xfer.handle_ack(ack, agent_name):
            del self._xfers[ack.mm_hash]
            self._tombstoned.add(ack.mm_hash)

    def _ensure_registered(self, metadata: bytes, mem_descriptor: bytes) -> str:
        """Register or re-register the peer in DataTransport. Returns agent_name."""
        if self._nixl_agent_name is not None and self._nixl_metadata_bytes == metadata:
            return self._nixl_agent_name
        if self._nixl_agent_name is not None:
            logger.info(
                "EC Connector: producer %s:%d restarted; re-registering",
                self._addr[0],
                self._addr[1],
            )
            self._data.remove_remote_peer(self._nixl_agent_name)
        agent_name = self._data.add_remote_peer(metadata, mem_descriptor)
        self._nixl_agent_name = agent_name
        self._nixl_metadata_bytes = metadata
        logger.debug(
            "EC Connector: registered peer %s:%d agent=%s",
            self._addr[0],
            self._addr[1],
            agent_name,
        )
        return agent_name

    def _drain_quarantine(self, now: float) -> None:
        still_pending: list[ConsumerXfer] = []
        for xfer in list(self._quarantined):
            state = xfer.poll(now)
            if state == XferState.SETTLED:
                self._settled.append((xfer.mm_hash, xfer.block_indices))
            else:
                still_pending.append(xfer)
        self._quarantined = still_pending

    def on_peer_down(self) -> None:
        """WAITING_ACK xfers → cancel (retry allowed); READING xfers → quarantine."""
        for mm_hash, xfer in list(self._xfers.items()):
            if xfer.transfer_handle is None:
                xfer.cancel()
                self._retryable.add(mm_hash)
            else:
                self._quarantined.append(xfer)
                self._newly_quarantined.add(mm_hash)
        self._xfers.clear()

    def take_results(self) -> ConsumerSessionResults:
        """Return and clear all accumulated results since the last call."""
        results = ConsumerSessionResults(
            completed=self._completed,
            tombstoned=self._tombstoned,
            quarantined=self._newly_quarantined,
            retryable=self._retryable,
            settled=self._settled,
        )
        self._completed = set()
        self._tombstoned = set()
        self._newly_quarantined = set()
        self._retryable = set()
        self._settled = []
        return results

    def has_in_flight(self) -> bool:
        return bool(self._xfers) or bool(self._quarantined)

    def close(self) -> None:
        """Release all transfer handles, deregister NIXL peer, close connection.

        Connection teardown is routed through the transport so the pooled
        connection is dropped in lock-step with the session — a session and its
        DEALER connection have a 1:1 lifecycle, so a later connect() to the same
        peer builds a fresh socket instead of handing back this dead one.
        """
        for xfer in self._xfers.values():
            xfer.release()
        self._xfers.clear()
        for xfer in self._quarantined:
            xfer.release()
        self._quarantined.clear()
        if self._nixl_agent_name is not None:
            self._data.remove_remote_peer(self._nixl_agent_name)
        self._transport.remove(self._addr)

@dataclass
class ConsumerSessionResults:
    """Results from ConsumerSession.take_results()."""

    completed: set[str]  # promote to local cache
    tombstoned: set[str]  # free blocks + prevent retry this step
    quarantined: set[str]  # timed-out, DMA still running: prevent retry
    # but DO NOT free blocks — session holds them
    retryable: set[str]  # free blocks, allow another attempt (not ready, peer down)
    settled: list[tuple[str, list[int]]]  # (mm_hash, block_indices) quarantine cleared

@dataclass
class ConsumerXfer:
    """State machine for one NIXL READ on the consumer side.

    transfer_handle is the in-flight NIXL handle returned by post_read() —
    used for check_xfer_state() / release_xfer_handle(). The dlist handle
    for the remote peer lives inside DataTransport, not here.
    """

    mm_hash: str
    block_indices: list[int]
    addr: PeerAddr
    deadline: float
    data: DataTransport
    consumer_session_id: str  # session identity; combined with mm_hash for notif_msg

    transfer_handle: Any | None = field(default=None)
    _quarantined: bool = field(default=False, init=False)

    def handle_ack(self, ack: XferAck, agent_name: str) -> bool:
        """Issue the NIXL READ on an OK XferAck. Returns False on any NACK."""
        if ack.status != XferStatus.OK:
            return False
        # notif_msg encodes both session_id and mm_hash so the producer can
        # do a direct (session_id, mm_hash) → ProducerXfer lookup on completion.
        notif_msg = f"{self.consumer_session_id}:{self.mm_hash}".encode()
        self.transfer_handle = self.data.post_read(
            self.block_indices,
            agent_name,
            ack.src_block_indices,
            notif_msg=notif_msg,
        )
        self.deadline = time.monotonic() + _CONSUMER_READ_TIMEOUT_S
        logger.debug(
            "EC consumer: NIXL READ posted mm_hash=%s agent=%s local_blocks=%d "
            "remote_blocks=%d",
            self.mm_hash,
            agent_name,
            len(self.block_indices),
            len(ack.src_block_indices),
        )
        return True

    def poll(self, now: float) -> XferState:
        if self._quarantined:
            return self._poll_quarantined(now)
        if self.transfer_handle is None:
            return (
                XferState.ACK_TIMEOUT if now > self.deadline else XferState.WAITING_ACK
            )
        return self._poll_reading(now)

    def _poll_reading(self, now: float) -> XferState:
        assert self.transfer_handle is not None
        try:
            state = self.data.check_xfer_state(self.transfer_handle)
        except Exception:
            logger.exception(
                "EC Connector: check_xfer_state failed for mm_hash=%s", self.mm_hash
            )
            self.data.release_xfer_handle(self.transfer_handle)
            self.transfer_handle = None
            return XferState.READ_FAILED
        if state == "DONE":
            self.data.release_xfer_handle(self.transfer_handle)
            self.transfer_handle = None
            logger.debug("EC consumer: NIXL READ completed mm_hash=%s", self.mm_hash)
            return XferState.DONE
        if state == "PROC":
            if now > self.deadline:
                self._quarantined = True
                self.deadline = now + _CONSUMER_QUARANTINE_TIMEOUT_S
                logger.warning(
                    "EC Connector: READ for mm_hash=%s from %s:%d timed out; "
                    "quarantining",
                    self.mm_hash,
                    self.addr[0],
                    self.addr[1],
                )
                return XferState.QUARANTINED
            return XferState.READING
        logger.warning(
            "EC Connector: READ for mm_hash=%s unexpected NIXL state %r",
            self.mm_hash,
            state,
        )
        self.data.release_xfer_handle(self.transfer_handle)
        self.transfer_handle = None
        return XferState.READ_FAILED

    def _poll_quarantined(self, now: float) -> XferState:
        if self.transfer_handle is None:
            return XferState.SETTLED
        # If quarantine has been too long, assume NIXL is stuck and evict blocks.
        if now > self.deadline:
            self.data.release_xfer_handle(self.transfer_handle)
            self.transfer_handle = None
            logger.warning(
                "EC Connector: quarantined mm_hash=%s from %s:%d timed out; "
                "evicting blocks",
                self.mm_hash,
                self.addr[0],
                self.addr[1],
            )
            return XferState.READ_FAILED
        try:
            state = self.data.check_xfer_state(self.transfer_handle)
        except Exception:
            logger.exception(
                "EC Connector: check_xfer_state failed for quarantined mm_hash=%s",
                self.mm_hash,
            )
            state = None
        if state == "PROC":
            return XferState.QUARANTINED
        self.data.release_xfer_handle(self.transfer_handle)
        self.transfer_handle = None
        logger.debug(
            "EC Connector: quarantined mm_hash=%s settled (state=%s)",
            self.mm_hash,
            state,
        )
        return XferState.SETTLED

    def cancel(self) -> None:
        """Cancel a WAITING_ACK xfer (peer down). No tombstone — caller allows retry."""
        assert self.transfer_handle is None, "cancel() requires WAITING_ACK state"

    def release(self) -> None:
        """Force-release the transfer handle. Shutdown only."""
        if self.transfer_handle is not None:
            self.data.release_xfer_handle(self.transfer_handle)
            self.transfer_handle = None

class ProducerSession:
    """Long-lived producer session. Owns ZmqServerTransport.

    poll_step() runs once per engine step: drains the transport without
    blocking, decodes XferReqs, grants or NACKs each one, drains NIXL completion
    notifications, and sweeps expired xfers.
    """

    def __init__(
        self,
        transport: ZmqServerTransport,
        data: DataTransport,
        cache: EmbeddingCache,
        compat_hash: str,
    ) -> None:
        self._transport = transport
        self._data = data
        self._cache = cache  # EmbeddingCache — single block owner
        self._compat_hash = compat_hash

        # Key: "{consumer_session_id}:{mm_hash}" — matches the NIXL notif_msg exactly,
        # so completion notification lookup is a direct dict pop with no parsing.
        self._active_xfers: dict[str, ProducerXfer] = {}

        self._req_decoder = msgspec.msgpack.Decoder(XferReq)
        self._req_encoder = msgspec.msgpack.Encoder()

    def close(self) -> None:
        for xfer in self._active_xfers.values():
            self._cache.unpin(xfer.mm_hash)  # release any still-held pins
        self._active_xfers.clear()
        self._transport.close()

    def poll_step(self) -> None:
        """Drain the transport without blocking and process what arrived."""
        self.poll(self._transport.poll(timeout_ms=0))

    def poll(self, messages: list[tuple[bytes, bytes]]) -> None:
        """Decode and process inbound XferReqs, drain NIXL notifs, sweep timeouts.

        messages: (identity, payload) pairs from ZmqServerTransport.poll().
        """
        for identity, payload in messages:
            try:
                req = self._req_decoder.decode(payload)
            except (msgspec.DecodeError, msgspec.ValidationError):
                logger.warning("EC Connector: dropped malformed XferReq")
                continue
            ack = self._grant_or_nack(req)
            self._transport.send(identity, self._req_encoder.encode(ack))

        self._drain_notifs()
        self._sweep_timeouts()

    def _grant_or_nack(self, req: XferReq) -> XferAck:
        if req.connector_version != EC_CONNECTOR_VERSION:
            logger.warning(
                "EC Connector: incompatible version for mm_hash=%s", req.mm_hash
            )
            return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_VERSION)
        if req.compatibility_hash != self._compat_hash:
            logger.warning(
                "EC Connector: incompatible compat hash for mm_hash=%s", req.mm_hash
            )
            return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INCOMPAT)
        if not req.session_id:
            # Grants are keyed by (session_id, mm_hash), so an omitted id from
            # one peer aliases another's: the first completion would unpin
            # blocks the other peer is still reading.
            logger.warning(
                "EC Connector: XferReq without session_id for mm_hash=%s", req.mm_hash
            )
            return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INTERNAL)
        entry = self._cache.pin_if_ready(req.mm_hash)
        if entry is None or not entry.ready:
            status = (
                XferStatus.NACK_MISSING if entry is None else XferStatus.NACK_NOT_READY
            )
            logger.log(
                logging.INFO if status in EXPECTED_NACKS else logging.WARNING,
                "EC Connector: mm_hash=%s not serveable in local cache (%s); NACKing",
                req.mm_hash,
                status.name,
            )
            return XferAck(mm_hash=req.mm_hash, status=status)
        block_indices = list(entry.block_ids)
        key = f"{req.session_id}:{req.mm_hash}"
        if key in self._active_xfers:
            # This peer is re-requesting a grant it still holds, having timed
            # out waiting. Overwriting it drops the pin taken for it, which
            # neither a completion nor the sweep would ever release.
            self._cache.unpin(req.mm_hash)
        self._active_xfers[key] = ProducerXfer(
            mm_hash=req.mm_hash,
            block_indices=block_indices,
            deadline=time.monotonic() + PRODUCER_PIN_LEASE_S,
        )
        logger.debug(
            "EC producer: serving NIXL READ request mm_hash=%s key=%s blocks=%d",
            req.mm_hash,
            key,
            len(block_indices),
        )
        return XferAck(
            mm_hash=req.mm_hash,
            status=XferStatus.OK,
            session_id=req.session_id,  # echo consumer's session_id back
            src_block_indices=block_indices,
            agent_metadata=self._data.get_agent_metadata(),
            mem_descriptor=self._data.get_mem_descriptor(),
        )

    def _drain_notifs(self) -> None:
        try:
            notifs = self._data.get_new_notifs()
        except Exception:
            logger.exception("EC Connector: get_new_notifs failed")
            return
        for msgs in notifs.values():
            for msg in msgs:
                try:
                    key = msg.decode("utf-8")  # "{session_id}:{mm_hash}"
                except Exception:
                    logger.exception("EC Connector: Dropped malformed NIXL notif")
                    continue
                xfer = self._active_xfers.pop(key, None)
                if xfer is not None:
                    self._cache.unpin(xfer.mm_hash)
                    logger.debug(
                        "EC producer: NIXL READ completed mm_hash=%s key=%s",
                        xfer.mm_hash,
                        key,
                    )

    def _sweep_timeouts(self) -> None:
        for key, xfer in list(self._active_xfers.items()):
            if xfer.is_expired():
                logger.warning(
                    "EC producer: consumer did not complete NIXL READ "
                    "mm_hash=%s key=%s within the pin lease; releasing "
                    "pinned blocks",
                    xfer.mm_hash,
                    key,
                )
                self._cache.unpin(xfer.mm_hash)
                del self._active_xfers[key]

@dataclass
class ProducerXfer:
    """One pinned read grant on the producer side.

    Pure data — lifecycle operations (pin/unpin) are performed by ProducerSession.
    Keyed by (consumer_session_id, mm_hash) in ProducerSession._active_xfers.
    """

    mm_hash: str
    block_indices: list[int]
    deadline: float

    def is_expired(self) -> bool:
        return time.monotonic() > self.deadline

class XferState(enum.Enum):
    """Lifecycle states of a single ConsumerXfer."""

    WAITING_ACK = "waiting_ack"  # XferReq sent, awaiting XferAck
    READING = "reading"  # NIXL READ in flight
    DONE = "done"  # READ completed; promote to local cache
    ACK_TIMEOUT = "ack_timeout"  # XferAck never arrived; free blocks + retry
    READ_FAILED = "read_failed"  # Unexpected NIXL state; free blocks + tombstone
    QUARANTINED = "quarantined"  # Read timed out, NIXL unabortable; keep blocks
    SETTLED = "settled"  # Quarantined DMA done; safe to free blocks