| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Abstract base class for the EC data-plane transport (NIXL/UCX).""" |
| 4 | |
| 5 | from abc import ABC, abstractmethod |
| 6 | from typing import Any |
| 7 | |
| 8 | |
| 9 | class DataTransport(ABC): |
| 10 | """Interface wrapping the NIXL data-plane for EC block transfers. |
| 11 | |
| 12 | Peer lifecycle: |
| 13 | add_remote_peer() — register a peer and store its dlist handle internally |
| 14 | remove_remote_peer() — deregister and release the stored handle |
| 15 | post_read() — issue a READ using the stored handle for agent_name |
| 16 | |
| 17 | The dlist handle (prepped xfer descriptor list) is an implementation detail |
| 18 | of the transport and is never exposed to callers. |
| 19 | """ |
| 20 | |
| 21 | @abstractmethod |
| 22 | def get_agent_metadata(self) -> bytes: |
| 23 | """Return this agent's serialized NIXL metadata blob.""" |
| 24 | |
| 25 | @abstractmethod |
| 26 | def get_mem_descriptor(self) -> bytes: |
| 27 | """Return the msgpack-encoded block descriptor list for this agent.""" |
| 28 | |
| 29 | @abstractmethod |
| 30 | def add_remote_peer(self, metadata: bytes, mem_descriptor: bytes) -> str: |
| 31 | """Register a remote peer, prep its READ-source dlist, and return agent_name. |
| 32 | |
| 33 | The dlist handle is stored internally and used by post_read(). |
| 34 | """ |
| 35 | |
| 36 | @abstractmethod |
| 37 | def remove_remote_peer(self, agent_name: str) -> None: |
| 38 | """Deregister a previously registered remote peer and release its handle.""" |
| 39 | |
| 40 | @abstractmethod |
| 41 | def post_read( |
| 42 | self, |
| 43 | local_indices: list[int], |
| 44 | agent_name: str, |
| 45 | remote_indices: list[int], |
| 46 | notif_msg: bytes, |
| 47 | ) -> Any: |
| 48 | """Issue a consumer-initiated READ using the stored dlist for agent_name. |
| 49 | |
| 50 | Returns a transfer handle for use with check_xfer_state / release_xfer_handle. |
| 51 | """ |
| 52 | |
| 53 | @abstractmethod |
| 54 | def get_new_notifs(self) -> dict[str, list[bytes]]: |
| 55 | """Drain NIXL completion notifications addressed to this agent.""" |
| 56 | |
| 57 | @abstractmethod |
| 58 | def check_xfer_state(self, handle: Any) -> str: |
| 59 | """Return the NIXL state string: 'DONE', 'PROC', or an error string.""" |
| 60 | |
| 61 | @abstractmethod |
| 62 | def release_xfer_handle(self, handle: Any) -> None: |
| 63 | """Release a transfer handle returned by post_read.""" |
| 64 | |
| 65 | @abstractmethod |
| 66 | def deregister(self) -> None: |
| 67 | """Deregister all memory and tear down the NIXL agent.""" |