| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for DataTransport base class and NixlTransport.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import ctypes |
| 8 | from unittest.mock import MagicMock, patch |
| 9 | |
| 10 | import numpy as np |
| 11 | |
| 12 | from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult |
| 13 | from vllm.v1.kv_offload.tiering.p2p.data.nixl import NixlTransport |
| 14 | |
| 15 | # --------------------------------------------------------------------------- |
| 16 | # DataTransport base class tests |
| 17 | # --------------------------------------------------------------------------- |
| 18 | |
| 19 | |
| 20 | class TestDataTransportBase: |
| 21 | """Tests for the DataTransport abstract base properties.""" |
| 22 | |
| 23 | def _make_view(self, num_blocks: int = 8, block_len: int = 1024) -> memoryview: |
| 24 | """Create a memoryview with the given shape.""" |
| 25 | buf = np.zeros((num_blocks, block_len), dtype=np.uint8) |
| 26 | return memoryview(buf) |
| 27 | |
| 28 | def test_properties(self): |
| 29 | """base_addr, num_blocks, block_len are set from memoryview shape.""" |
| 30 | view = self._make_view(num_blocks=4, block_len=2048) |
| 31 | |
| 32 | # Use NixlTransport (concrete) with NIXL mocked away |
| 33 | with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None): |
| 34 | transport = NixlTransport("test:1", view) |
| 35 | |
| 36 | assert transport.num_blocks == 4 |
| 37 | assert transport.block_len == 2048 |
| 38 | assert transport.base_addr == ctypes.addressof(ctypes.c_char.from_buffer(view)) |
| 39 | |
| 40 | def test_config_fingerprint_empty_when_no_fields(self): |
| 41 | """No config fields → empty fingerprint.""" |
| 42 | view = self._make_view() |
| 43 | with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None): |
| 44 | transport = NixlTransport("test:1", view, config_fields=None) |
| 45 | assert transport.config_fingerprint == "" |
| 46 | |
| 47 | def test_config_fingerprint_deterministic(self): |
| 48 | """Same config fields → same fingerprint.""" |
| 49 | view = self._make_view() |
| 50 | fields = {"model": "llama", "dtype": "float16", "blocks_per_chunk": 1} |
| 51 | with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None): |
| 52 | t1 = NixlTransport("test:1", view, config_fields=fields) |
| 53 | t2 = NixlTransport("test:2", view, config_fields=fields) |
| 54 | assert t1.config_fingerprint == t2.config_fingerprint |
| 55 | assert len(t1.config_fingerprint) == 16 |
| 56 | |
| 57 | def test_config_fingerprint_differs_for_different_fields(self): |
| 58 | """Different config fields → different fingerprint.""" |
| 59 | view = self._make_view() |
| 60 | with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None): |
| 61 | t1 = NixlTransport("test:1", view, config_fields={"model": "a"}) |
| 62 | t2 = NixlTransport("test:2", view, config_fields={"model": "b"}) |
| 63 | assert t1.config_fingerprint != t2.config_fingerprint |
| 64 | |
| 65 | |
| 66 | # --------------------------------------------------------------------------- |
| 67 | # NixlTransport tests (with mocked NIXL agent) |
| 68 | # --------------------------------------------------------------------------- |
| 69 | |
| 70 | |
| 71 | class TestNixlTransportWithMockedAgent: |
| 72 | """Tests for NixlTransport logic with a mocked NIXL agent.""" |
| 73 | |
| 74 | def _make_transport(self) -> NixlTransport: |
| 75 | """Create a NixlTransport with mocked NIXL internals.""" |
| 76 | view = memoryview(np.zeros((8, 1024), dtype=np.uint8)) |
| 77 | |
| 78 | with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None): |
| 79 | transport = NixlTransport("test:1", view) |
| 80 | |
| 81 | # Manually set up a mock agent after construction |
| 82 | agent = MagicMock() |
| 83 | agent.add_remote_agent.return_value = "nixl-peer-name" |
| 84 | agent.get_xfer_descs.return_value = MagicMock() |
| 85 | agent.prep_xfer_dlist.return_value = MagicMock() |
| 86 | agent.make_prepped_xfer.return_value = MagicMock(name="handle") |
| 87 | agent.transfer.return_value = None |
| 88 | agent.check_xfer_state.return_value = "PROC" |
| 89 | agent.get_agent_metadata.return_value = b"test-metadata" |
| 90 | |
| 91 | transport._agent = agent |
| 92 | transport._local_dlist = MagicMock() |
| 93 | return transport |
| 94 | |
| 95 | def test_available_false_without_nixl(self): |
| 96 | """Without NIXL installed, available is False.""" |
| 97 | view = memoryview(np.zeros((4, 512), dtype=np.uint8)) |
| 98 | with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None): |
| 99 | transport = NixlTransport("test:1", view) |
| 100 | assert transport.available is False |
| 101 | |
| 102 | def test_available_true_with_agent(self): |
| 103 | transport = self._make_transport() |
| 104 | assert transport.available is True |
| 105 | |
| 106 | def test_get_agent_metadata(self): |
| 107 | transport = self._make_transport() |
| 108 | assert transport.get_agent_metadata() == b"test-metadata" |
| 109 | |
| 110 | def test_write_blocks_returns_none_for_unknown_peer(self): |
| 111 | """write_blocks returns None if peer not registered.""" |
| 112 | transport = self._make_transport() |
| 113 | result = transport.write_blocks("unknown:1", [0, 1], [2, 3]) |
| 114 | assert result is None |
| 115 | |
| 116 | def test_write_blocks_returns_transfer_id(self): |
| 117 | """write_blocks returns an integer transfer_id on success.""" |
| 118 | transport = self._make_transport() |
| 119 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 120 | |
| 121 | tid = transport.write_blocks("peer:1", [0, 1], [2, 3]) |
| 122 | assert tid is not None |
| 123 | assert isinstance(tid, int) |
| 124 | |
| 125 | def test_write_blocks_increments_transfer_id(self): |
| 126 | """Each write_blocks call gets a unique transfer_id.""" |
| 127 | transport = self._make_transport() |
| 128 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 129 | |
| 130 | tid1 = transport.write_blocks("peer:1", [0], [1]) |
| 131 | tid2 = transport.write_blocks("peer:1", [2], [3]) |
| 132 | assert tid1 != tid2 |
| 133 | |
| 134 | def test_poll_empty_when_no_inflight(self): |
| 135 | """Poll returns empty when nothing is inflight.""" |
| 136 | transport = self._make_transport() |
| 137 | result = transport.poll() |
| 138 | assert result == PollResult(done=(), failed=()) |
| 139 | |
| 140 | def test_poll_returns_done_when_transfer_completes(self): |
| 141 | """Completed transfer appears in poll().done.""" |
| 142 | transport = self._make_transport() |
| 143 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 144 | |
| 145 | tid = transport.write_blocks("peer:1", [0], [1]) |
| 146 | |
| 147 | # Simulate completion |
| 148 | transport._agent.check_xfer_state.return_value = "DONE" |
| 149 | result = transport.poll() |
| 150 | |
| 151 | assert tid in result.done |
| 152 | assert result.failed == () |
| 153 | # Handle released |
| 154 | transport._agent.release_xfer_handle.assert_called() |
| 155 | |
| 156 | def test_poll_returns_failed_for_error_state(self): |
| 157 | """Transfer in error state appears in poll().failed.""" |
| 158 | transport = self._make_transport() |
| 159 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 160 | |
| 161 | tid = transport.write_blocks("peer:1", [0], [1]) |
| 162 | |
| 163 | transport._agent.check_xfer_state.return_value = "ERR" |
| 164 | result = transport.poll() |
| 165 | |
| 166 | assert result.done == () |
| 167 | assert tid in result.failed |
| 168 | |
| 169 | def test_poll_ignores_in_progress(self): |
| 170 | """Transfers in PROC/PEND state stay inflight.""" |
| 171 | transport = self._make_transport() |
| 172 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 173 | |
| 174 | transport.write_blocks("peer:1", [0], [1]) |
| 175 | |
| 176 | transport._agent.check_xfer_state.return_value = "PROC" |
| 177 | result = transport.poll() |
| 178 | assert result.done == () |
| 179 | assert result.failed == () |
| 180 | |
| 181 | transport._agent.check_xfer_state.return_value = "PEND" |
| 182 | result = transport.poll() |
| 183 | assert result.done == () |
| 184 | assert result.failed == () |
| 185 | |
| 186 | def test_poll_peer_id_scopes_to_peer(self): |
| 187 | """poll(peer_id) drains only that peer's transfers. |
| 188 | |
| 189 | Regression: the transport is shared across peer sessions (e.g. a |
| 190 | single prefiller serving a DP>1 decoder). An unscoped poll by one |
| 191 | session used to consume and discard sibling sessions' completions, |
| 192 | starving them until timeout. poll(peer_id) must leave other peers' |
| 193 | transfers inflight. |
| 194 | """ |
| 195 | transport = self._make_transport() |
| 196 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 197 | transport.add_remote_peer("peer:2", b"meta", 0x2000, 8, 1024) |
| 198 | |
| 199 | tid1 = transport.write_blocks("peer:1", [0], [1]) |
| 200 | tid2 = transport.write_blocks("peer:2", [2], [3]) |
| 201 | transport._agent.check_xfer_state.return_value = "DONE" |
| 202 | |
| 203 | # Polling peer:1 must not consume peer:2's completed transfer. |
| 204 | result = transport.poll(peer_id="peer:1") |
| 205 | assert tid1 in result.done |
| 206 | assert tid2 not in result.done |
| 207 | assert tid1 not in transport._inflight |
| 208 | assert tid2 in transport._inflight |
| 209 | |
| 210 | # peer:2 sees its own completion when it polls. |
| 211 | result2 = transport.poll(peer_id="peer:2") |
| 212 | assert tid2 in result2.done |
| 213 | assert tid2 not in transport._inflight |
| 214 | |
| 215 | def test_cancel_removes_inflight(self): |
| 216 | """Cancel removes transfers and releases handles.""" |
| 217 | transport = self._make_transport() |
| 218 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 219 | |
| 220 | tid = transport.write_blocks("peer:1", [0], [1]) |
| 221 | assert tid in transport._inflight |
| 222 | |
| 223 | result = transport.cancel([tid]) |
| 224 | assert result == [] |
| 225 | assert tid not in transport._inflight |
| 226 | transport._agent.release_xfer_handle.assert_called() |
| 227 | |
| 228 | def test_cancel_ignores_unknown_ids(self): |
| 229 | """Cancel with unknown IDs doesn't crash.""" |
| 230 | transport = self._make_transport() |
| 231 | assert transport.cancel([999, 1000]) == [] |
| 232 | assert transport.cancel([999, 1000], mode="wait") == [] |
| 233 | |
| 234 | def test_cancel_wait_release_succeeds(self): |
| 235 | """wait-mode cancel that succeeds pops the entry and returns [].""" |
| 236 | transport = self._make_transport() |
| 237 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 238 | |
| 239 | tid = transport.write_blocks("peer:1", [0], [1]) |
| 240 | assert tid in transport._inflight |
| 241 | |
| 242 | result = transport.cancel([tid], mode="wait") |
| 243 | assert result == [] |
| 244 | assert tid not in transport._inflight |
| 245 | transport._agent.release_xfer_handle.assert_called_once() |
| 246 | |
| 247 | def test_cancel_wait_release_raises(self): |
| 248 | """wait-mode cancel keeps the entry and returns the tid on raise.""" |
| 249 | transport = self._make_transport() |
| 250 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 251 | |
| 252 | tid = transport.write_blocks("peer:1", [0], [1]) |
| 253 | transport._agent.release_xfer_handle.side_effect = RuntimeError( |
| 254 | "NIXL_ERR_REPOST_ACTIVE" |
| 255 | ) |
| 256 | |
| 257 | result = transport.cancel([tid], mode="wait") |
| 258 | assert result == [tid] |
| 259 | assert tid in transport._inflight |
| 260 | |
| 261 | def test_cancel_wait_then_poll_completes(self): |
| 262 | """A wait-cancel that left a tid pending later completes via poll.""" |
| 263 | transport = self._make_transport() |
| 264 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 265 | |
| 266 | tid = transport.write_blocks("peer:1", [0], [1]) |
| 267 | transport._agent.release_xfer_handle.side_effect = RuntimeError("busy") |
| 268 | assert transport.cancel([tid], mode="wait") == [tid] |
| 269 | assert tid in transport._inflight |
| 270 | |
| 271 | transport._agent.release_xfer_handle.side_effect = None |
| 272 | transport._agent.check_xfer_state.return_value = "DONE" |
| 273 | |
| 274 | result = transport.poll() |
| 275 | assert tid in result.done |
| 276 | assert tid not in transport._inflight |
| 277 | |
| 278 | def test_add_and_remove_remote_peer(self): |
| 279 | transport = self._make_transport() |
| 280 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 281 | assert "peer:1" in transport._remote_dlists |
| 282 | |
| 283 | transport.remove_remote_peer("peer:1") |
| 284 | assert "peer:1" not in transport._remote_dlists |
| 285 | transport._agent.release_dlist_handle.assert_called() |
| 286 | transport._agent.remove_remote_agent.assert_called() |
| 287 | |
| 288 | def test_close_releases_everything(self): |
| 289 | """Close releases all handles and clears state.""" |
| 290 | transport = self._make_transport() |
| 291 | transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) |
| 292 | transport.write_blocks("peer:1", [0], [1]) |
| 293 | |
| 294 | transport.close() |
| 295 | assert transport._agent is None |
| 296 | assert transport._inflight == {} |
| 297 | assert transport._remote_dlists == {} |
| 298 | |
| 299 | |
| 300 | # --------------------------------------------------------------------------- |
| 301 | # NIXL agent-config selection |
| 302 | # --------------------------------------------------------------------------- |
| 303 | |
| 304 | |
| 305 | class TestNixlAgentConfigSelection: |
| 306 | """Tests that backends/num_threads pick the right nixl_agent_config call. |
| 307 | |
| 308 | Mirrors the conditional in |
| 309 | vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py:325-329. |
| 310 | """ |
| 311 | |
| 312 | def _make_view(self) -> memoryview: |
| 313 | return memoryview(np.zeros((4, 512), dtype=np.uint8)) |
| 314 | |
| 315 | def test_non_ucx_backends_passes_backends_kwarg(self): |
| 316 | """When any non-UCX backend is requested, pass backends + telemetry.""" |
| 317 | agent_cls = MagicMock() |
| 318 | config_fn = MagicMock(return_value=MagicMock(name="cfg")) |
| 319 | with ( |
| 320 | patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", agent_cls), |
| 321 | patch( |
| 322 | "vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig", config_fn |
| 323 | ), |
| 324 | ): |
| 325 | NixlTransport("test:1", self._make_view(), backends=["MOONCAKE"]) |
| 326 | |
| 327 | config_fn.assert_called_once_with(backends=["MOONCAKE"], capture_telemetry=True) |
| 328 | # num_threads must NOT be passed on the non-UCX branch. |
| 329 | assert "num_threads" not in config_fn.call_args.kwargs |
| 330 | |
| 331 | def test_ucx_only_passes_num_threads(self): |
| 332 | """UCX-only configuration passes num_threads + telemetry, no backends.""" |
| 333 | agent_cls = MagicMock() |
| 334 | config_fn = MagicMock(return_value=MagicMock(name="cfg")) |
| 335 | with ( |
| 336 | patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", agent_cls), |
| 337 | patch( |
| 338 | "vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig", config_fn |
| 339 | ), |
| 340 | ): |
| 341 | NixlTransport("test:1", self._make_view(), num_threads=8) |
| 342 | |
| 343 | config_fn.assert_called_once_with(num_threads=8, capture_telemetry=True) |
| 344 | assert "backends" not in config_fn.call_args.kwargs |
| 345 | |
| 346 | def test_default_backends_is_ucx_only(self): |
| 347 | """No backends arg → defaults to UCX-only branch.""" |
| 348 | agent_cls = MagicMock() |
| 349 | config_fn = MagicMock(return_value=MagicMock(name="cfg")) |
| 350 | with ( |
| 351 | patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", agent_cls), |
| 352 | patch( |
| 353 | "vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig", config_fn |
| 354 | ), |
| 355 | ): |
| 356 | NixlTransport("test:1", self._make_view()) |
| 357 | |
| 358 | # Default num_threads=4, no backends kwarg. |
| 359 | config_fn.assert_called_once_with(num_threads=4, capture_telemetry=True) |