vllm.v1.kv_offload.tiering.p2p.data ¶

class DataTransport(ABC):
    """Abstract data-plane transport for RDMA-style block transfers.

    Owns the local KV block memory region and manages transfers to/from
    registered remote peers.

    Construction:
        view: A 2D memoryview (num_blocks × block_len bytes) over the
              local KV cache block storage.
        config_fields: Dict of model config values used to compute the
                       compatibility fingerprint. None → empty fingerprint
                       (compatible with any peer).
    """

    def __init__(self, view: memoryview, config_fields: dict | None = None) -> None:
        assert view.shape is not None
        self._view = view
        self._base_addr = ctypes.addressof(ctypes.c_char.from_buffer(view))
        self._num_blocks = view.shape[0]
        self._block_len = view.shape[1]
        self._config_fingerprint = self._compute_fingerprint(config_fields)

    @property
    def base_addr(self) -> int:
        """Base address of the local block memory region."""
        return self._base_addr

    @property
    def num_blocks(self) -> int:
        """Number of blocks in the local region."""
        return self._num_blocks

    @property
    def block_len(self) -> int:
        """Size of each block in bytes."""
        return self._block_len

    @property
    def config_fingerprint(self) -> str:
        """Content-hash of the model configuration (hex string).

        Peers must have matching fingerprints to exchange blocks.
        Empty string means no fingerprint (always compatible).
        """
        return self._config_fingerprint

    @staticmethod
    def _compute_fingerprint(config_fields: dict | None) -> str:
        if not config_fields:
            return ""
        canonical = json.dumps(config_fields, sort_keys=True, separators=(",", ":"))
        return hashlib.sha256(canonical.encode()).hexdigest()[:16]

    @abstractmethod
    def get_agent_metadata(self) -> bytes:
        """Return opaque metadata needed by remote peers to connect.

        The returned bytes are sent during the control-plane handshake
        and passed to the remote peer's add_remote_peer().
        """
        ...

    @abstractmethod
    def add_remote_peer(
        self,
        peer_id: str,
        agent_metadata: bytes,
        base_addr: int,
        num_blocks: int,
        block_len: int,
    ) -> None:
        """Register a remote peer for block transfers.

        Must be called before write_blocks() to this peer.

        Args:
            peer_id: Unique identifier for the remote peer.
            agent_metadata: Opaque bytes from the peer's get_agent_metadata().
            base_addr: Base address of the peer's block memory region.
            num_blocks: Number of blocks in the peer's region.
            block_len: Size of each block (must match local block_len).
        """
        ...

    @abstractmethod
    def remove_remote_peer(self, peer_id: str) -> None:
        """Unregister a remote peer and release associated resources.

        Inflight transfers to this peer should be cancelled first.
        """
        ...

    @abstractmethod
    def write_blocks(
        self,
        peer_id: str,
        local_idxs: list[int],
        remote_idxs: list[int],
    ) -> int | None:
        """Submit a WRITE transfer: local blocks → remote peer's blocks.

        Args:
            peer_id: Target peer (must be registered via add_remote_peer).
            local_idxs: Indexes of local blocks to read from.
            remote_idxs: Indexes of remote blocks to write to.
                         Must be same length as local_idxs.

        Returns:
            A unique transfer_id (int) to track this transfer, or
            None if the peer is not registered or submission failed.
        """
        ...

    @abstractmethod
    def poll(self) -> PollResult:
        """Poll all inflight transfers for completion.

        Returns:
            PollResult with lists of completed and failed transfer_ids.
            Completed/failed transfers are removed from the inflight set.

        Must be called periodically to drive progress checking.
        """
        ...

    @abstractmethod
    def cancel(
        self,
        transfer_ids: Iterable[int],
        mode: CancelMode = "immediate",
    ) -> list[int]:
        """Cancel inflight transfers by their IDs.

        Best-effort: transfers that already completed are ignored.

        Args:
            transfer_ids: IDs to cancel. Unknown IDs are ignored.
            mode:
                "immediate" (default): pop and release each handle and
                    return []. Matches the legacy fire-and-forget
                    behavior — the caller does not wait for the
                    underlying transfer to drain.
                "wait": attempt to release each handle. If the release
                    cannot complete because the transfer is still
                    PROC/PEND, the entry stays in the inflight set and
                    its id is included in the returned list. The
                    caller is expected to keep calling poll() until
                    every returned id surfaces in done/failed.

        Returns:
            For mode="wait", the subset of *transfer_ids* still
            tracked as inflight after the cancel attempt. For
            mode="immediate", always [].
        """
        ...

    @abstractmethod
    def close(self) -> None:
        """Release all resources (registrations, handles, memory).

        Cancels any remaining inflight transfers. Idempotent.
        After close(), no other methods may be called.
        """
        ...

class NixlTransport(DataTransport):
    """Manages a NIXL agent, memory registration, and block transfers.

    Wraps the NIXL C library behind a Python interface so the rest of the
    P2P tier code never touches NIXL types directly. Tracks inflight
    handles internally and returns completed/failed tags on poll.
    """

    def __init__(
        self,
        local_id: str,
        view: memoryview,
        config_fields: dict | None = None,
        backends: list[str] | None = None,
        num_threads: int = 4,
    ) -> None:
        super().__init__(view, config_fields=config_fields)
        self._local_id = local_id
        self._backends = list(backends) if backends else ["UCX"]
        self._num_threads = num_threads
        self._agent: Any = None
        self._reg: Any = None
        self._local_dlist: Any = None
        self._remote_dlists: dict[str, object] = {}
        self._peer_nixl_names: dict[str, str] = {}
        self._inflight: dict[int, object] = {}  # transfer_id → handle
        self._next_id = itertools.count()

        self._init(view)

    @property
    def available(self) -> bool:
        return self._agent is not None

    def _init(self, view: memoryview) -> None:
        if _NixlAgent is None:
            return

        non_ucx_backends = [b for b in self._backends if b != "UCX"]
        if non_ucx_backends:
            cfg = _NixlAgentConfig(backends=self._backends, capture_telemetry=True)
            logger.info(
                "NixlTransport %s: NIXL backends=%s",
                self._local_id,
                self._backends,
            )
        else:
            cfg = _NixlAgentConfig(
                num_threads=self._num_threads, capture_telemetry=True
            )
            logger.info(
                "NixlTransport %s: NIXL backends=[UCX] num_threads=%d",
                self._local_id,
                self._num_threads,
            )
        self._agent = _NixlAgent(self._local_id, cfg)

        total_size = self._num_blocks * self._block_len
        reg_descs = [(self._base_addr, total_size, 0, "")]
        self._reg = self._agent.register_memory(reg_descs, mem_type="DRAM")

        block_tuples = [
            (self._base_addr + i * self._block_len, self._block_len, 0)
            for i in range(self._num_blocks)
        ]
        xfer_dlist = self._agent.get_xfer_descs(block_tuples, mem_type="DRAM")
        self._local_dlist = self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", xfer_dlist)
        logger.info(
            "NixlTransport %s: registered %d blocks", self._local_id, self._num_blocks
        )

    def get_agent_metadata(self) -> bytes:
        assert self._agent is not None
        return self._agent.get_agent_metadata()

    # ------------------------------------------------------------------
    # Peer management
    # ------------------------------------------------------------------

    def add_remote_peer(
        self,
        peer_id: str,
        agent_metadata: bytes,
        base_addr: int,
        num_blocks: int,
        block_len: int,
    ) -> None:
        nixl_name = self._agent.add_remote_agent(agent_metadata)
        block_descs = [
            (base_addr + i * block_len, block_len, 0) for i in range(num_blocks)
        ]
        xfer_dlist = self._agent.get_xfer_descs(block_descs, mem_type="DRAM")
        remote_dlist = self._agent.prep_xfer_dlist(nixl_name, xfer_dlist)
        self._peer_nixl_names[peer_id] = nixl_name
        self._remote_dlists[peer_id] = remote_dlist

    def remove_remote_peer(self, peer_id: str) -> None:
        nixl_name = self._peer_nixl_names.pop(peer_id, None)
        dlist = self._remote_dlists.pop(peer_id, None)
        if self._agent is not None:
            if dlist is not None:
                self._agent.release_dlist_handle(dlist)
            if nixl_name:
                self._agent.remove_remote_agent(nixl_name)

    # ------------------------------------------------------------------
    # Transfer submission and polling
    # ------------------------------------------------------------------

    def write_blocks(
        self,
        peer_id: str,
        local_idxs: list[int],
        remote_idxs: list[int],
    ) -> int | None:
        """Submit a WRITE transfer to *peer_id*.

        Returns a transfer ID, or None if the peer is not registered.
        The ID is returned via poll() when the transfer completes or fails.
        """
        remote_dlist = self._remote_dlists.get(peer_id)
        if remote_dlist is None:
            logger.warning(
                "NixlTransport %s: write_blocks NO REMOTE DLIST for peer=%s "
                "(known peers=%s)",
                self._local_id,
                peer_id,
                list(self._remote_dlists.keys()),
            )
            return None
        logger.debug(
            "NixlTransport %s: write_blocks NIXL.transfer peer=%s blocks=%d",
            self._local_id,
            peer_id,
            len(local_idxs),
        )
        handle = self._agent.make_prepped_xfer(
            "WRITE",
            self._local_dlist,
            local_idxs,
            remote_dlist,
            remote_idxs,
        )
        self._agent.transfer(handle)
        transfer_id = next(self._next_id)
        self._inflight[transfer_id] = handle
        return transfer_id

    def poll(self) -> PollResult:
        """Poll all inflight transfers.

        Returns PollResult(done=..., failed=...) with transfer IDs.
        Completed handles are released automatically.
        """
        if not self._inflight:
            return _EMPTY_POLL_RESULT

        done_ids: list[int] | None = None
        failed_ids: list[int] | None = None

        for transfer_id, handle in self._inflight.items():
            try:
                state = self._agent.check_xfer_state(handle)
            except Exception as exc:
                logger.warning(
                    "NixlTransport %s: check_xfer_state failed for transfer_id=%d: %s",
                    self._local_id,
                    transfer_id,
                    exc,
                )
                continue
            if state == "DONE":
                if done_ids is None:
                    done_ids = []
                done_ids.append(transfer_id)
            elif state not in ("PROC", "PEND"):
                if failed_ids is None:
                    failed_ids = []
                failed_ids.append(transfer_id)

        if done_ids is None and failed_ids is None:
            return _EMPTY_POLL_RESULT

        handles_to_release = []
        for tid in done_ids or ():
            handles_to_release.append(self._inflight.pop(tid))
        for tid in failed_ids or ():
            handles_to_release.append(self._inflight.pop(tid))
        self._release_handles(handles_to_release)

        return PollResult(
            done=done_ids if done_ids is not None else _EMPTY_POLL_RESULT.done,
            failed=failed_ids if failed_ids is not None else _EMPTY_POLL_RESULT.failed,
        )

    def cancel(
        self,
        transfer_ids: Iterable[int],
        mode: CancelMode = "immediate",
    ) -> list[int]:
        """Cancel inflight transfers by their IDs.

        See ``DataTransport.cancel`` for the contract. In "wait" mode,
        transfers whose ``release_xfer_handle`` raises (NIXL could not
        complete the abort because the backend is still draining) stay
        in ``self._inflight`` so a later ``poll()`` will observe them.
        """
        if mode == "immediate":
            handles = [
                self._inflight.pop(tid) for tid in transfer_ids if tid in self._inflight
            ]
            self._release_handles(handles)
            return []

        still_inflight: list[int] = []
        for tid in transfer_ids:
            handle = self._inflight.get(tid)
            if handle is None:
                continue
            try:
                self._agent.release_xfer_handle(handle)
            except Exception as exc:
                logger.debug(
                    "NixlTransport %s: cancel pending for transfer_id=%d: %s",
                    self._local_id,
                    tid,
                    exc,
                )
                still_inflight.append(tid)
                continue
            del self._inflight[tid]
        return still_inflight

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def close(self) -> None:
        if self._agent is None:
            return
        self._release_handles(list(self._inflight.values()))
        self._inflight.clear()
        for peer_id in list(self._remote_dlists):
            self.remove_remote_peer(peer_id)
        if self._local_dlist is not None:
            self._agent.release_dlist_handle(self._local_dlist)
            self._local_dlist = None
        if self._reg is not None:
            self._agent.deregister_memory(self._reg)
            self._reg = None
        self._agent = None

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

    def _release_handles(self, handles: list[object]) -> None:
        if self._agent is None:
            return
        for handle in handles:
            try:
                self._agent.release_xfer_handle(handle)
            except Exception as exc:
                logger.warning(
                    "NixlTransport %s: release_xfer_handle failed: %s",
                    self._local_id,
                    exc,
                )

class PollResult(NamedTuple):
    """Result of polling inflight transfers.

    Attributes:
        done: Transfer IDs that completed successfully.
        failed: Transfer IDs that failed (error, timeout, etc.).
    """

    done: Sequence[int]
    failed: Sequence[int]