| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Abstract base class for the P2P data-plane transport. |
| 4 | |
| 5 | The data plane handles RDMA (or similar) block transfers between peers. |
| 6 | It is independent of the control plane — the control plane establishes |
| 7 | who can talk to whom, then the data plane moves blocks at wire speed. |
| 8 | |
| 9 | Architecture |
| 10 | ------------ |
| 11 | |
| 12 | DataTransport (one per node) |
| 13 | ├── owns the local KV block memory region |
| 14 | ├── registers remote peers (add_remote_peer / remove_remote_peer) |
| 15 | ├── submits block writes (write_blocks → transfer_id) |
| 16 | ├── polls for completion (poll → done/failed IDs) |
| 17 | └── cancels inflight transfers (cancel) |
| 18 | |
| 19 | Memory model |
| 20 | ------------ |
| 21 | |
| 22 | The local node exposes a contiguous block region: |
| 23 | |
| 24 | base_addr ──► ┌─────────────┐ block 0 |
| 25 | ├─────────────┤ block 1 |
| 26 | ├─────────────┤ ... |
| 27 | └─────────────┘ block (num_blocks - 1) |
| 28 | |
| 29 | Each block is block_len bytes. |
| 30 | |
| 31 | Remote peers expose the same layout. write_blocks() copies local |
| 32 | blocks to a remote peer's block region by index: |
| 33 | |
| 34 | write_blocks("peer:1", local_idxs=[0, 3], remote_idxs=[5, 7]) |
| 35 | → writes local block 0 → remote block 5 |
| 36 | local block 3 → remote block 7 |
| 37 | |
| 38 | Transfer lifecycle |
| 39 | ------------------ |
| 40 | |
| 41 | 1. Register remote peer: add_remote_peer(peer_id, metadata, ...) |
| 42 | - Provides the remote memory layout so transfers can target it |
| 43 | 2. Submit: write_blocks(peer_id, local_idxs, remote_idxs) → int |
| 44 | - Returns a transfer_id (opaque int) for tracking |
| 45 | - Returns None if peer not registered or submission fails |
| 46 | 3. Poll: poll() → PollResult(done=[...], failed=[...]) |
| 47 | - Returns transfer_ids that completed or failed since last poll |
| 48 | - Completed transfers are automatically cleaned up |
| 49 | 4. Cancel: cancel(transfer_ids, mode="immediate" | "wait") |
| 50 | - Best-effort cancellation of inflight transfers |
| 51 | - mode="wait" returns ids still in PROC/PEND so the caller can poll |
| 52 | them to completion |
| 53 | |
| 54 | Implementor contracts |
| 55 | --------------------- |
| 56 | |
| 57 | - write_blocks() must not block. The transfer runs asynchronously. |
| 58 | - poll() must be called periodically. It drives completion checking. |
| 59 | - transfer_ids are unique across the lifetime of the transport. |
| 60 | - add_remote_peer() must be called before write_blocks() to that peer. |
| 61 | - get_agent_metadata() returns opaque bytes that the remote peer |
| 62 | needs to call add_remote_peer() (e.g., RDMA connection info). |
| 63 | - config_fingerprint is a content hash of the model configuration. |
| 64 | Peers with different fingerprints are incompatible and must not |
| 65 | exchange blocks (validated during the control-plane handshake). |
| 66 | - close() releases all resources (memory registrations, handles). |
| 67 | After close(), no other methods may be called. |
| 68 | |
| 69 | Threading model: no background threads. All I/O driven by poll(). |
| 70 | """ |
| 71 | |
| 72 | from __future__ import annotations |
| 73 | |
| 74 | import ctypes |
| 75 | import hashlib |
| 76 | import json |
| 77 | from abc import ABC, abstractmethod |
| 78 | from collections.abc import Iterable, Sequence |
| 79 | from typing import Literal, NamedTuple |
| 80 | |
| 81 | CancelMode = Literal["immediate", "wait"] |
| 82 | |
| 83 | |
| 84 | class PollResult(NamedTuple): |
| 85 | """Result of polling inflight transfers. |
| 86 | |
| 87 | Attributes: |
| 88 | done: Transfer IDs that completed successfully. |
| 89 | failed: Transfer IDs that failed (error, timeout, etc.). |
| 90 | |
| 91 | """ |
| 92 | |
| 93 | done: Sequence[int] |
| 94 | failed: Sequence[int] |
| 95 | |
| 96 | |
| 97 | class DataTransport(ABC): |
| 98 | """Abstract data-plane transport for RDMA-style block transfers. |
| 99 | |
| 100 | Owns the local KV block memory region and manages transfers to/from |
| 101 | registered remote peers. |
| 102 | |
| 103 | Construction: |
| 104 | view: A 2D memoryview (num_blocks × block_len bytes) over the |
| 105 | local KV cache block storage. |
| 106 | config_fields: Dict of model config values used to compute the |
| 107 | compatibility fingerprint. None → empty fingerprint |
| 108 | (compatible with any peer). |
| 109 | """ |
| 110 | |
| 111 | def __init__(self, view: memoryview, config_fields: dict | None = None) -> None: |
| 112 | assert view.shape is not None |
| 113 | self._view = view |
| 114 | self._base_addr = ctypes.addressof(ctypes.c_char.from_buffer(view)) |
| 115 | self._num_blocks = view.shape[0] |
| 116 | self._block_len = view.shape[1] |
| 117 | self._config_fingerprint = self._compute_fingerprint(config_fields) |
| 118 | |
| 119 | @property |
| 120 | def base_addr(self) -> int: |
| 121 | """Base address of the local block memory region.""" |
| 122 | return self._base_addr |
| 123 | |
| 124 | @property |
| 125 | def num_blocks(self) -> int: |
| 126 | """Number of blocks in the local region.""" |
| 127 | return self._num_blocks |
| 128 | |
| 129 | @property |
| 130 | def block_len(self) -> int: |
| 131 | """Size of each block in bytes.""" |
| 132 | return self._block_len |
| 133 | |
| 134 | @property |
| 135 | def config_fingerprint(self) -> str: |
| 136 | """Content-hash of the model configuration (hex string). |
| 137 | |
| 138 | Peers must have matching fingerprints to exchange blocks. |
| 139 | Empty string means no fingerprint (always compatible). |
| 140 | """ |
| 141 | return self._config_fingerprint |
| 142 | |
| 143 | @staticmethod |
| 144 | def _compute_fingerprint(config_fields: dict | None) -> str: |
| 145 | if not config_fields: |
| 146 | return "" |
| 147 | canonical = json.dumps(config_fields, sort_keys=True, separators=(",", ":")) |
| 148 | return hashlib.sha256(canonical.encode()).hexdigest()[:16] |
| 149 | |
| 150 | @abstractmethod |
| 151 | def get_agent_metadata(self) -> bytes: |
| 152 | """Return opaque metadata needed by remote peers to connect. |
| 153 | |
| 154 | The returned bytes are sent during the control-plane handshake |
| 155 | and passed to the remote peer's add_remote_peer(). |
| 156 | """ |
| 157 | ... |
| 158 | |
| 159 | @abstractmethod |
| 160 | def add_remote_peer( |
| 161 | self, |
| 162 | peer_id: str, |
| 163 | agent_metadata: bytes, |
| 164 | base_addr: int, |
| 165 | num_blocks: int, |
| 166 | block_len: int, |
| 167 | ) -> None: |
| 168 | """Register a remote peer for block transfers. |
| 169 | |
| 170 | Must be called before write_blocks() to this peer. |
| 171 | |
| 172 | Args: |
| 173 | peer_id: Unique identifier for the remote peer. |
| 174 | agent_metadata: Opaque bytes from the peer's get_agent_metadata(). |
| 175 | base_addr: Base address of the peer's block memory region. |
| 176 | num_blocks: Number of blocks in the peer's region. |
| 177 | block_len: Size of each block (must match local block_len). |
| 178 | |
| 179 | """ |
| 180 | ... |
| 181 | |
| 182 | @abstractmethod |
| 183 | def remove_remote_peer(self, peer_id: str) -> None: |
| 184 | """Unregister a remote peer and release associated resources. |
| 185 | |
| 186 | Inflight transfers to this peer should be cancelled first. |
| 187 | """ |
| 188 | ... |
| 189 | |
| 190 | @abstractmethod |
| 191 | def write_blocks( |
| 192 | self, |
| 193 | peer_id: str, |
| 194 | local_idxs: list[int], |
| 195 | remote_idxs: list[int], |
| 196 | ) -> int | None: |
| 197 | """Submit a WRITE transfer: local blocks → remote peer's blocks. |
| 198 | |
| 199 | Args: |
| 200 | peer_id: Target peer (must be registered via add_remote_peer). |
| 201 | local_idxs: Indexes of local blocks to read from. |
| 202 | remote_idxs: Indexes of remote blocks to write to. |
| 203 | Must be same length as local_idxs. |
| 204 | |
| 205 | Returns: |
| 206 | A unique transfer_id (int) to track this transfer, or |
| 207 | None if the peer is not registered or submission failed. |
| 208 | |
| 209 | """ |
| 210 | ... |
| 211 | |
| 212 | @abstractmethod |
| 213 | def poll(self, peer_id: str | None = None) -> PollResult: |
| 214 | """Poll inflight transfers for completion. |
| 215 | |
| 216 | Args: |
| 217 | peer_id: If given, only poll (and drain) transfers submitted for |
| 218 | this peer_id — the value passed to ``write_blocks``. This is |
| 219 | required when a single transport is shared across multiple |
| 220 | peer sessions: ``poll()`` pops completed handles, so an |
| 221 | unscoped poll by one session would consume and discard the |
| 222 | completions of its siblings, starving them. ``None`` polls |
| 223 | every peer's transfers (used only for the shutdown drain). |
| 224 | |
| 225 | Returns: |
| 226 | PollResult with lists of completed and failed transfer_ids. |
| 227 | Completed/failed transfers are removed from the inflight set. |
| 228 | |
| 229 | Must be called periodically to drive progress checking. |
| 230 | |
| 231 | """ |
| 232 | ... |
| 233 | |
| 234 | @abstractmethod |
| 235 | def cancel( |
| 236 | self, |
| 237 | transfer_ids: Iterable[int], |
| 238 | mode: CancelMode = "immediate", |
| 239 | ) -> list[int]: |
| 240 | """Cancel inflight transfers by their IDs. |
| 241 | |
| 242 | Best-effort: transfers that already completed are ignored. |
| 243 | |
| 244 | Args: |
| 245 | transfer_ids: IDs to cancel. Unknown IDs are ignored. |
| 246 | mode: |
| 247 | "immediate" (default): pop and release each handle and |
| 248 | return []. Matches the legacy fire-and-forget |
| 249 | behavior — the caller does not wait for the |
| 250 | underlying transfer to drain. |
| 251 | "wait": attempt to release each handle. If the release |
| 252 | cannot complete because the transfer is still |
| 253 | PROC/PEND, the entry stays in the inflight set and |
| 254 | its id is included in the returned list. The |
| 255 | caller is expected to keep calling poll() until |
| 256 | every returned id surfaces in done/failed. |
| 257 | |
| 258 | Returns: |
| 259 | For mode="wait", the subset of *transfer_ids* still |
| 260 | tracked as inflight after the cancel attempt. For |
| 261 | mode="immediate", always []. |
| 262 | |
| 263 | """ |
| 264 | ... |
| 265 | |
| 266 | @abstractmethod |
| 267 | def close(self) -> None: |
| 268 | """Release all resources (registrations, handles, memory). |
| 269 | |
| 270 | Cancels any remaining inflight transfers. Idempotent. |
| 271 | After close(), no other methods may be called. |
| 272 | """ |
| 273 | ... |