vllm.v1.kv_offload.tiering.p2p.session ¶

class LoadResult(NamedTuple):
    """Result from a session poll, client side."""

    job_id: int
    kv_request_id: str
    success: bool

class P2PSession:
    """Bidirectional session — coordinator over ClientRole + ServerRole.

    Lifecycle:
      - Constructor with conn=None  ⇒ pending. Accepts add_stored_blocks
        but cannot send (used by the prefiller to buffer blocks before
        the decoder connects).
      - Constructor with conn != None ⇒ connected. Sends our own ConnectMsg
        immediately; the peer's ConnectMsg arrives in poll() and is
        dispatched to _on_connect (which calls transport.add_remote_peer
        and replies with ConnectAckMsg). Outgoing sends are queued until
        ConnectAckMsg confirms our metadata reached the peer.
      - attach_connection(conn) on a pending session ⇒ same as above,
        starting from pending.
    """

    def __init__(
        self,
        peer_id: str,
        local_id: str,
        transport: DataTransport,
        local_block_len: int,
        local_hash_seed: str,
        conn: ControlConnection | None = None,
    ) -> None:
        self.peer_id = peer_id
        self._local_id = local_id
        self._transport = transport
        self._local_block_len = local_block_len
        self._local_hash_seed = local_hash_seed
        self._conn: ControlConnection | None = None

        self._send_ready = False  # True after the peer acked our ConnectMsg
        # Msgs waiting to be sent on connection establishment
        self._queued: list[dict] = []

        # Consecutive non-protocol dispatch errors. Reset on success.
        self._dispatch_error_count: int = 0

        # kv_request_ids whose FetchMsg arrived during the current poll
        # tick. Drained and returned in the next poll() result so the
        # manager can bind kv_request_id → session and replay any
        # submit_store batches parked before the binding existed.
        self._new_fetch_ids: list[str] = []

        self._client = ClientRole(peer_id=peer_id, send=self._send)
        self._server = ServerRole(
            peer_id=peer_id,
            transport=transport,
            send=self._send,
        )

        if conn is not None:
            self.attach_connection(conn)

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def alive(self) -> bool:
        # Pending sessions (awaiting connection) are alive — only a
        # closed real connection counts as dead.
        return self._conn is None or self._conn.alive

    @property
    def connected(self) -> bool:
        return self._conn is not None

    @property
    def ready(self) -> bool:
        """True after the peer acked our ConnectMsg (we may send freely)."""
        return self._send_ready

    @property
    def has_pending_work(self) -> bool:
        """True while inbound loads or outbound transfers are outstanding."""
        return self._client.has_active_loads or self._server.has_inflight_transfers

    # ------------------------------------------------------------------
    # Connection lifecycle
    # ------------------------------------------------------------------

    def attach_connection(self, conn: ControlConnection) -> None:
        """Attach a connection to a pending session and announce ourselves.

        Symmetric: every side advertises its NIXL metadata on connect, so
        whichever peer receives a session first can register the other.
        """
        if self._conn is not None:
            raise ValueError(f"P2PSession {self.peer_id}: already connected")
        self._conn = conn
        self._send_connect()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def request_blocks(
        self,
        job_id: JobId,
        kv_request_id: str,
        keys: Sequence[OffloadKey],
        block_ids: Sequence[int],
    ) -> None:
        """Send fetch to the peer."""
        self._client.request_blocks(
            job_id, kv_request_id, keys, block_ids, send_ready=self._send_ready
        )

    def add_stored_blocks(
        self,
        kv_request_id: str,
        keys: Sequence[OffloadKey],
        block_ids: Sequence[int],
        job_id: JobId,
    ) -> None:
        """New blocks stored locally — match against pending fetch demand."""
        self._server.add_stored_blocks(kv_request_id, keys, block_ids, job_id)

    def finish_request(self, kv_request_id: str) -> None:
        """Called when the request is finishing locally.

        Finishes the client role (aborts any inbound load and drops any
        pending symmetric-P2P lookup state) and finalizes any outbound
        serving (server role) for this id. Roles that aren't active for
        this id are silent no-ops.
        """
        self._client.finish(kv_request_id)
        self._server.finish(kv_request_id)

    def register_lookup(self, kv_request_id: str, key: bytes) -> bool | None:
        """Register or resolve one (kv_request_id, key) probe.

        Called from the manager's lookup() for symmetric-P2P consumers
        (``remote_kv_source`` sub-dict in kv_transfer_params). See
        ``ClientRole.register_lookup`` for the state-machine contract.
        """
        return self._client.register_lookup(kv_request_id, key)

    def flush_pending_lookups(self) -> None:
        """Flush any aggregated symmetric-P2P lookups for this peer.

        Called once per scheduler step from the manager's
        ``on_schedule_end()``. Send-gating is handled inside the
        client's ``_send`` callback (queues until ConnectAckMsg).
        """
        self._client.flush_pending_lookups()

    def serve_external_requests(self, parent: ParentManager) -> None:
        """Resolve inbound peer lookups against the tiering manager.

        Delegates to the server role; the ``parent`` handle is valid
        only for the duration of this call.
        """
        self._server.serve_external_requests(parent)

    def poll(self) -> SessionPollResult:
        """Process incoming messages, drive transfers, apply timeouts."""
        if self._conn is None:
            # Pending session — store-job timeouts still apply so buffered
            # jobs that never get picked up are surfaced as failures.
            return SessionPollResult(
                loads=[],
                stores=self._server.collect_idle_timeouts(),
                new_fetch_ids=[],
            )

        for msg in self._conn.recv():
            self._on_message(msg)

        loads = self._client.collect_results()
        stores = self._server.collect_results()
        self._server.drain_pending_aborts()

        new_fetch_ids = self._new_fetch_ids
        self._new_fetch_ids = []
        return SessionPollResult(
            loads=loads, stores=stores, new_fetch_ids=new_fetch_ids
        )

    def close(self) -> SessionCloseResult:
        """Shut down.

        failed_jobs: client load job_ids to fail.
        failed_req_ids: client kv_request_ids to fail — in-flight loads plus
            requests with an unresolved symmetric-P2P probe toward the
            now-dead peer. The manager fails these so the consumer's lookup()
            falls back to local prefill instead of deferring forever on an
            answer that can never arrive.
        failed_stores: server store job_ids to fail.
        failed_serves: synthetic lookup ctxs still owing
            ``parent.on_request_finished`` (the manager flushes these on
            its next ``serve_external_requests``).
        """
        client_result = self._client.close()
        failed_stores, failed_serves = self._server.close()

        if self._conn is not None:
            with contextlib.suppress(Exception):
                self._conn.send({TYPE_KEY: DisconnectMsg.TYPE})
            self._conn.close()
            self._conn = None

        return SessionCloseResult(
            failed_jobs=client_result.failed_jobs,
            failed_req_ids=client_result.failed_req_ids,
            failed_stores=failed_stores,
            failed_serves=failed_serves,
        )

    # ------------------------------------------------------------------
    # Message dispatch
    # ------------------------------------------------------------------

    def _on_message(self, msg: dict) -> None:
        msg_type = msg.get(TYPE_KEY) if isinstance(msg, dict) else msg
        try:
            self._dispatch_message(msg)
        except ValueError as exc:
            # Protocol contract violation from the peer — *Msg.validate()
            # and handler-level checks raise ValueError. Retrying won't
            # help and may corrupt session state, so disconnect now.
            self._protocol_error(f"malformed {msg_type!r}: {exc}")
            return
        except Exception as exc:
            # Anything else is most likely an internal bug rather than a
            # peer fault. Log loudly with a traceback so it doesn't
            # disappear, but don't kill the session on a single hiccup.
            # Disconnect only if errors keep arriving — that pattern is
            # consistent with a peer wedging us into a broken state.
            self._dispatch_error_count += 1
            logger.exception(
                "P2PSession %s: error handling message %r (count=%d): %s",
                self.peer_id,
                msg_type,
                self._dispatch_error_count,
                exc,
            )
            if self._dispatch_error_count >= _MAX_CONSECUTIVE_DISPATCH_ERRORS:
                self._protocol_error(
                    f"too many consecutive dispatch errors "
                    f"({self._dispatch_error_count})"
                )
            return
        self._dispatch_error_count = 0

    def _protocol_error(self, reason: str) -> None:
        """Log a protocol violation and disconnect.

        Best-effort sends ``DisconnectMsg`` so the peer learns why we're
        going away, then marks the connection dead. The manager reaps
        the session on the next poll via ``alive``.
        """
        logger.error(
            "P2PSession %s: protocol error: %s — disconnecting",
            self.peer_id,
            reason,
        )
        if self._conn is not None:
            with contextlib.suppress(Exception):
                self._conn.send({TYPE_KEY: DisconnectMsg.TYPE})
            self._conn.mark_dead()

    def _dispatch_message(self, msg: dict) -> None:
        # Drop messages buffered before disconnect: a poll batch can
        # contain msg-after-DisconnectMsg, and dispatching them would
        # mutate state on a dead session.
        if self._conn is not None and not self._conn.alive:
            return
        msg_type = msg.get(TYPE_KEY) if isinstance(msg, dict) else None
        if msg_type == ConnectMsg.TYPE:
            self._on_connect(msg)
        elif msg_type == ConnectAckMsg.TYPE:
            ConnectAckMsg.validate(msg)
            self._on_connect_ack()
        elif msg_type == FetchMsg.TYPE:
            FetchMsg.validate(msg)
            kv_request_id = msg[FetchMsg.KV_REQUEST_ID]
            keys = [
                OffloadKey(bh if isinstance(bh, bytes) else bytes(bh))
                for bh in msg[FetchMsg.KEYS]
            ]
            block_indexes = msg[FetchMsg.BLOCK_INDEXES]
            round_seq = msg[FetchMsg.ROUND_SEQ]
            # Run the server-role state machine inline as today —
            # add_fetch_demand records demand against any blocks we've
            # already seen in `available`. Report the kv_request_id so
            # the manager (after poll() returns) can replay any parked
            # submit_store batches; their add_stored_blocks calls hit
            # the demand recorded here and submit transfers immediately.
            self._server.on_fetch(kv_request_id, keys, block_indexes, round_seq)
            self._new_fetch_ids.append(kv_request_id)
        elif msg_type == AbortFetchMsg.TYPE:
            AbortFetchMsg.validate(msg)
            self._server.on_abort_fetch(
                msg[AbortFetchMsg.KV_REQUEST_ID],
                msg[AbortFetchMsg.ROUND_SEQ],
            )
        elif msg_type == TransferDoneMsg.TYPE:
            TransferDoneMsg.validate(msg)
            self._client.on_transfer_done(
                msg[TransferDoneMsg.KV_REQUEST_ID],
                msg[TransferDoneMsg.SUCCESS],
                msg[TransferDoneMsg.ROUND_SEQ],
            )
        elif msg_type == AbortAckMsg.TYPE:
            AbortAckMsg.validate(msg)
            self._client.on_abort_ack(
                msg[AbortAckMsg.KV_REQUEST_ID],
                msg[AbortAckMsg.ROUND_SEQ],
            )
        elif msg_type == LookupMsg.TYPE:
            LookupMsg.validate(msg)
            kv_request_id = msg[LookupMsg.KV_REQUEST_ID]
            keys = [
                OffloadKey(bh if isinstance(bh, bytes) else bytes(bh))
                for bh in msg[LookupMsg.KEYS]
            ]
            self._server.on_lookup(kv_request_id, keys, msg[LookupMsg.ROUND_SEQ])
        elif msg_type == LookupRespMsg.TYPE:
            LookupRespMsg.validate(msg)
            kv_request_id = msg[LookupRespMsg.KV_REQUEST_ID]
            keys = [
                OffloadKey(bh if isinstance(bh, bytes) else bytes(bh))
                for bh in msg[LookupRespMsg.KEYS]
            ]
            hits = msg[LookupRespMsg.HITS]
            self._client.on_lookup_resp(kv_request_id, keys, hits)
        elif msg_type == DisconnectMsg.TYPE:
            if self._conn is not None:
                self._conn.mark_dead()
        else:
            logger.warning(
                "P2PSession %s: unknown message type %r", self.peer_id, msg_type
            )

    # ------------------------------------------------------------------
    # Handshake
    # ------------------------------------------------------------------

    def _on_connect(self, msg: dict) -> None:
        # Validation failures here mean an incompatible or malicious peer.
        # Mark the connection dead so the manager reaps the session;
        # don't call add_remote_peer or send connect_ack.
        if self._send_ready:
            # We've already received connect_ack, so the handshake is
            # complete. A second connect from the peer is a protocol
            # violation — re-registering would corrupt transport state.
            self._protocol_error("duplicate connect after handshake")
            return
        try:
            ConnectMsg.validate(msg)
            if msg[ConnectMsg.BLOCK_LEN] != self._local_block_len:
                raise ValueError(
                    f"block_len mismatch from {self.peer_id}: "
                    f"remote={msg[ConnectMsg.BLOCK_LEN]}, "
                    f"local={self._local_block_len}"
                )
            remote_fp = msg.get(ConnectMsg.CONFIG_FINGERPRINT, "")
            local_fp = self._transport.config_fingerprint
            if local_fp and remote_fp and remote_fp != local_fp:
                raise ValueError(
                    f"config fingerprint mismatch from {self.peer_id}: "
                    f"remote={remote_fp!r}, local={local_fp!r}"
                )
            if msg[ConnectMsg.HASH_SEED] != self._local_hash_seed:
                raise ValueError(
                    f"PYTHONHASHSEED mismatch from {self.peer_id}: "
                    f"remote={msg[ConnectMsg.HASH_SEED]!r}, "
                    f"local={self._local_hash_seed!r}"
                )
            self._transport.add_remote_peer(
                self.peer_id,
                agent_metadata=msg[ConnectMsg.AGENT_METADATA],
                base_addr=msg[ConnectMsg.BASE_ADDR],
                num_blocks=msg[ConnectMsg.NUM_BLOCKS],
                block_len=msg[ConnectMsg.BLOCK_LEN],
            )
        except ValueError as exc:
            logger.error("P2PSession %s: rejecting peer connect: %s", self.peer_id, exc)
            if self._conn is not None:
                self._conn.mark_dead()
            return

        if self._conn is not None:
            self._conn.send(
                {
                    TYPE_KEY: ConnectAckMsg.TYPE,
                    ConnectAckMsg.PEER_ID: self._local_id,
                }
            )

    def _on_connect_ack(self) -> None:
        if self._queued:
            logger.debug(
                "P2PSession %s: connect_ack received, flushing %d queued msg(s)",
                self.peer_id,
                len(self._queued),
            )
        self._send_ready = True
        for queued in self._queued:
            self._do_send(queued)
        self._queued.clear()

    # ------------------------------------------------------------------
    # Send helpers
    # ------------------------------------------------------------------

    def _send_connect(self) -> None:
        """Send our ConnectMsg announcing local NIXL metadata."""
        assert self._conn is not None
        self._conn.send(
            {
                TYPE_KEY: ConnectMsg.TYPE,
                ConnectMsg.PEER_ID: self._local_id,
                ConnectMsg.AGENT_METADATA: self._transport.get_agent_metadata(),
                ConnectMsg.BASE_ADDR: self._transport.base_addr,
                ConnectMsg.NUM_BLOCKS: self._transport.num_blocks,
                ConnectMsg.BLOCK_LEN: self._transport.block_len,
                ConnectMsg.CONFIG_FINGERPRINT: self._transport.config_fingerprint,
                ConnectMsg.HASH_SEED: self._local_hash_seed,
            }
        )

    def _send(self, msg: dict) -> None:
        if self._conn is None or not self._send_ready:
            logger.debug(
                "P2PSession %s: queueing %s (ready=%s queue_depth=%d)",
                self.peer_id,
                msg.get(TYPE_KEY),
                self._send_ready,
                len(self._queued) + 1,
            )
            self._queued.append(msg)
            return
        self._do_send(msg)

    def _do_send(self, msg: dict) -> None:
        if self._conn is None:
            return
        try:
            self._conn.send(msg)
            logger.debug("P2PSession %s: sent %s", self.peer_id, msg.get(TYPE_KEY))
        except Exception:
            # A send failure means the connection is broken. Swallowing it
            # silently strands every in-flight lookup/load toward this peer:
            # the session stays alive, is never reaped, and the consumer's
            # lookup() keeps returning RETRY until the HTTP client times out.
            # Mark the connection dead so the manager reaps the session on
            # its next poll and surfaces the stranded work as failures.
            logger.warning(
                "P2PSession %s: send of %s failed — marking connection dead",
                self.peer_id,
                msg.get(TYPE_KEY),
            )
            self._conn.mark_dead()

class SessionCloseResult(NamedTuple):
    """Result of tearing down a P2PSession.

    `failed_jobs`/`failed_stores` are the in-flight jobs (client loads /
    server stores) the manager must mark failed. `failed_req_ids` is every
    client-side kv_request_id whose lookup() must fail (in-flight loads plus
    unresolved probes); `failed_serves` is the server-side lookup state the
    dead peer can no longer resolve.
    """

    failed_jobs: list[int]  # client load job_ids
    failed_req_ids: list[str]  # client kv_request_ids (loads + probes)
    failed_stores: list[int]  # server store job_ids
    failed_serves: list[ReqContext]  # server-side lookup ctxs needing release

class SessionPollResult(NamedTuple):
    """Result of one P2PSession.poll() tick.

    `loads`/`stores` are the same per-role results the manager has always
    consumed. `new_fetch_ids` reports kv_request_ids whose FetchMsg
    arrived this tick — the manager uses them to bind kv_request_id →
    session and replay any submit_store batches parked while no peer had
    asked yet. Reporting (rather than calling back into the manager
    mid-dispatch) keeps the dependency strictly top-down.
    """

    loads: list[LoadResult]
    stores: list[StoreResult]
    new_fetch_ids: list[str]

class StoreResult(NamedTuple):
    """Result from a session poll, server side."""

    job_id: int
    success: bool