| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """NIXL-backed data-plane transport for ECCPUConnector.""" |
| 4 | |
| 5 | from typing import Any |
| 6 | |
| 7 | from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import ( |
| 8 | DataTransport, |
| 9 | ) |
| 10 | from vllm.distributed.ec_transfer.ec_connector.cpu.utils import ( |
| 11 | build_block_descs, |
| 12 | deserialize_mem_descriptor, |
| 13 | serialize_mem_descriptor, |
| 14 | ) |
| 15 | from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config |
| 16 | from vllm.logger import init_logger |
| 17 | |
| 18 | logger = init_logger(__name__) |
| 19 | _NIXL_DRAM = "DRAM" |
| 20 | |
| 21 | |
| 22 | class NixlDataTransport(DataTransport): |
| 23 | """Adapts NixlWrapper to the DataTransport interface. |
| 24 | |
| 25 | Dlist handles (prepped xfer descriptor lists for remote peers) are stored |
| 26 | internally keyed by agent_name. Callers never touch raw NIXL handles — |
| 27 | post_read() resolves the handle from agent_name internally. |
| 28 | |
| 29 | NixlWrapper is not thread-safe, and neither is this class. Every call site |
| 30 | runs on the scheduler thread, which serializes the transport by |
| 31 | construction; transfers are async inside NIXL and are drained from that |
| 32 | same thread via get_new_notifs() / check_xfer_state(). |
| 33 | """ |
| 34 | |
| 35 | def __init__( |
| 36 | self, |
| 37 | agent_name: str, |
| 38 | base_ptr: int, |
| 39 | num_blocks: int, |
| 40 | block_size_bytes: int, |
| 41 | total_size_bytes: int, |
| 42 | num_threads: int = 1, |
| 43 | ) -> None: |
| 44 | if NixlWrapper is None or nixl_agent_config is None: |
| 45 | raise RuntimeError( |
| 46 | "ECCPUConnector requires NIXL; " |
| 47 | "install the `nixl` package or set a different ec_connector." |
| 48 | ) |
| 49 | # agent_name → remote dlist handle (from add_remote_peer, used by post_read) |
| 50 | self._peer_handles: dict[str, int] = {} |
| 51 | # None until the local region is registered, so deregister() is safe on |
| 52 | # a transport whose registration failed part-way. |
| 53 | self._reg_descs: Any = None |
| 54 | |
| 55 | self._nixl = NixlWrapper( |
| 56 | agent_name, |
| 57 | nixl_agent_config(num_threads=num_threads, capture_telemetry=True), |
| 58 | ) |
| 59 | self._register_local_memory( |
| 60 | base_ptr, num_blocks, block_size_bytes, total_size_bytes |
| 61 | ) |
| 62 | |
| 63 | def _register_local_memory( |
| 64 | self, |
| 65 | base_ptr: int, |
| 66 | num_blocks: int, |
| 67 | block_size_bytes: int, |
| 68 | total_size_bytes: int, |
| 69 | ) -> None: |
| 70 | """Register the local region with NIXL and prep its transfer dlist. |
| 71 | |
| 72 | Unwinds the registration if a later step fails: __init__ propagates the |
| 73 | error, so the caller never receives an object it could deregister with. |
| 74 | """ |
| 75 | block_descs = build_block_descs(base_ptr, num_blocks, block_size_bytes) |
| 76 | self._reg_descs = self._nixl.get_reg_descs( |
| 77 | [(base_ptr, total_size_bytes, 0, "")], _NIXL_DRAM |
| 78 | ) |
| 79 | self._nixl.register_memory(self._reg_descs, backends=["UCX"]) |
| 80 | try: |
| 81 | xfer_descs = self._nixl.get_xfer_descs(block_descs, _NIXL_DRAM) |
| 82 | self._local_xfer_handle = self._nixl.prep_xfer_dlist( |
| 83 | "NIXL_INIT_AGENT", xfer_descs |
| 84 | ) |
| 85 | self._agent_metadata = self._nixl.get_agent_metadata() |
| 86 | self._mem_descriptor_bytes = serialize_mem_descriptor(block_descs) |
| 87 | except Exception: |
| 88 | self.deregister() |
| 89 | raise |
| 90 | |
| 91 | def get_agent_metadata(self) -> bytes: |
| 92 | return self._agent_metadata |
| 93 | |
| 94 | def get_mem_descriptor(self) -> bytes: |
| 95 | return self._mem_descriptor_bytes |
| 96 | |
| 97 | def add_remote_peer(self, metadata: bytes, mem_descriptor: bytes) -> str: |
| 98 | """Register remote peer, prep its dlist, store handle internally. |
| 99 | |
| 100 | Returns agent_name for use in post_read() and remove_remote_peer(). |
| 101 | """ |
| 102 | agent_name = self._nixl.add_remote_agent(metadata) |
| 103 | remote_blocks = deserialize_mem_descriptor(mem_descriptor) |
| 104 | remote_xfer_descs = self._nixl.get_xfer_descs(remote_blocks, _NIXL_DRAM) |
| 105 | self._peer_handles[agent_name] = self._nixl.prep_xfer_dlist( |
| 106 | agent_name, remote_xfer_descs |
| 107 | ) |
| 108 | return agent_name |
| 109 | |
| 110 | def remove_remote_peer(self, agent_name: str) -> None: |
| 111 | self._peer_handles.pop(agent_name, None) |
| 112 | try: |
| 113 | self._nixl.remove_remote_agent(agent_name) |
| 114 | except Exception: |
| 115 | logger.warning( |
| 116 | "EC: remove_remote_peer failed for %s", agent_name, exc_info=True |
| 117 | ) |
| 118 | |
| 119 | def post_read( |
| 120 | self, |
| 121 | local_indices: list[int], |
| 122 | agent_name: str, |
| 123 | remote_indices: list[int], |
| 124 | notif_msg: bytes, |
| 125 | ) -> Any: |
| 126 | if len(local_indices) != len(remote_indices): |
| 127 | raise ValueError( |
| 128 | f"EC: local/remote block count mismatch " |
| 129 | f"({len(local_indices)} vs {len(remote_indices)})" |
| 130 | ) |
| 131 | remote_handle = self._peer_handles[agent_name] |
| 132 | handle = self._nixl.make_prepped_xfer( |
| 133 | "READ", |
| 134 | self._local_xfer_handle, |
| 135 | local_indices, |
| 136 | remote_handle, |
| 137 | remote_indices, |
| 138 | notif_msg=notif_msg, |
| 139 | ) |
| 140 | self._nixl.transfer(handle) |
| 141 | return handle |
| 142 | |
| 143 | def get_new_notifs(self) -> dict[str, list[bytes]]: |
| 144 | return self._nixl.get_new_notifs() |
| 145 | |
| 146 | def check_xfer_state(self, handle: Any) -> str: |
| 147 | return self._nixl.check_xfer_state(handle) |
| 148 | |
| 149 | def release_xfer_handle(self, handle: Any) -> None: |
| 150 | try: |
| 151 | self._nixl.release_xfer_handle(handle) |
| 152 | except Exception: |
| 153 | logger.warning("EC: release_xfer_handle failed", exc_info=True) |
| 154 | |
| 155 | def deregister(self) -> None: |
| 156 | if self._reg_descs is None: |
| 157 | return |
| 158 | try: |
| 159 | self._nixl.deregister_memory(self._reg_descs) |
| 160 | except Exception: |
| 161 | logger.warning("EC: deregister failed", exc_info=True) |
| 162 | self._reg_descs = None |