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

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 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]