@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 184 lines 6.9 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for DataTransport ABC and NixlDataTransport."""
4
5 from unittest.mock import MagicMock, patch
6
7 import pytest
8
9 from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import (
10 DataTransport,
11 )
12 from vllm.distributed.ec_transfer.ec_connector.cpu.utils import (
13 serialize_mem_descriptor,
14 )
15
16 _NIXL_PATH = "vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl"
17
18
19 # ── helpers ───────────────────────────────────────────────────────────────────
20
21
22 def _make_nixl_mock(prep_side_effect=None) -> MagicMock:
23 nixl = MagicMock()
24 nixl.get_agent_metadata.return_value = b"agent-meta"
25 nixl.get_reg_descs.return_value = MagicMock()
26 nixl.get_xfer_descs.return_value = MagicMock()
27 nixl.prep_xfer_dlist.side_effect = prep_side_effect or [77, 88]
28 nixl.add_remote_agent.return_value = "remote-agent-1"
29 nixl.get_new_notifs.return_value = {}
30 nixl.check_xfer_state.return_value = "PROC"
31 return nixl
32
33
34 def _make_transport(nixl_mock=None):
35 from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
36 NixlDataTransport,
37 )
38
39 mock = nixl_mock or _make_nixl_mock()
40 with (
41 patch(f"{_NIXL_PATH}.NixlWrapper", return_value=mock),
42 patch(f"{_NIXL_PATH}.nixl_agent_config", return_value=MagicMock()),
43 ):
44 t = NixlDataTransport("agent", 0, 4, 64, 256)
45 t._nixl = mock
46 return t
47
48
49 # ── ABC ───────────────────────────────────────────────────────────────────────
50
51
52 def test_data_transport_abc_cannot_be_instantiated():
53 with pytest.raises(TypeError):
54 DataTransport() # type: ignore[abstract]
55
56
57 # ── NixlDataTransport construction ───────────────────────────────────────────
58
59
60 def test_nixl_data_transport_registers_memory_on_init():
61 nixl = _make_nixl_mock()
62 _make_transport(nixl)
63 nixl.register_memory.assert_called_once()
64
65
66 def test_nixl_data_transport_unwinds_registration_if_init_fails():
67 """A failure after register_memory must not leave the region pinned.
68
69 __init__ propagates, so the caller never gets a transport to deregister.
70 """
71 nixl = _make_nixl_mock(prep_side_effect=RuntimeError("no dlist"))
72 from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
73 NixlDataTransport,
74 )
75
76 with (
77 patch(f"{_NIXL_PATH}.NixlWrapper", return_value=nixl),
78 patch(f"{_NIXL_PATH}.nixl_agent_config", return_value=MagicMock()),
79 pytest.raises(RuntimeError, match="no dlist"),
80 ):
81 NixlDataTransport("agent", 0, 4, 64, 256)
82
83 nixl.deregister_memory.assert_called_once_with(nixl.get_reg_descs.return_value)
84
85
86 def test_nixl_data_transport_raises_if_nixl_unavailable():
87 from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
88 NixlDataTransport,
89 )
90
91 with (
92 patch(f"{_NIXL_PATH}.NixlWrapper", None),
93 patch(f"{_NIXL_PATH}.nixl_agent_config", None),
94 pytest.raises(RuntimeError, match="requires NIXL"),
95 ):
96 NixlDataTransport("a", 0, 1, 64, 64)
97
98
99 # ── public accessors ──────────────────────────────────────────────────────────
100
101
102 def test_get_agent_metadata_returns_nixl_blob():
103 assert _make_transport().get_agent_metadata() == b"agent-meta"
104
105
106 def test_get_mem_descriptor_roundtrips_block_descs():
107 t = _make_transport()
108 raw = t.get_mem_descriptor()
109 assert isinstance(raw, bytes) and len(raw) > 0
110
111
112 # ── add_remote_peer ───────────────────────────────────────────────────────────
113
114
115 def test_add_remote_peer_returns_agent_name_and_stores_handle_internally():
116 nixl = _make_nixl_mock(prep_side_effect=[77, 88])
117 t = _make_transport(nixl)
118 mem_desc = serialize_mem_descriptor([(0, 64, 0)])
119 agent_name = t.add_remote_peer(b"meta", mem_desc)
120 assert agent_name == "remote-agent-1"
121 # Dlist handle stored internally; accessible via _peer_handles.
122 assert t._peer_handles["remote-agent-1"] == 88
123
124
125 def test_add_remote_peer_registers_remote_agent():
126 nixl = _make_nixl_mock(prep_side_effect=[77, 88])
127 t = _make_transport(nixl)
128 mem_desc = serialize_mem_descriptor([(0, 64, 0)])
129 t.add_remote_peer(b"fresh-meta", mem_desc)
130 nixl.add_remote_agent.assert_called_once_with(b"fresh-meta")
131
132
133 # ── remove_remote_peer ────────────────────────────────────────────────────────
134
135
136 def test_remove_remote_peer_delegates_to_nixl():
137 t = _make_transport()
138 t.remove_remote_peer("agent-x")
139 t._nixl.remove_remote_agent.assert_called_once_with("agent-x")
140
141
142 def test_remove_remote_peer_swallows_exceptions():
143 t = _make_transport()
144 t._nixl.remove_remote_agent.side_effect = RuntimeError("gone")
145 t.remove_remote_peer("agent-x") # must not raise
146
147
148 # ── post_read ─────────────────────────────────────────────────────────────────
149
150
151 def test_post_read_raises_on_index_count_mismatch():
152 nixl = _make_nixl_mock(prep_side_effect=[77, 88])
153 t = _make_transport(nixl)
154 mem_desc = serialize_mem_descriptor([(0, 64, 0)])
155 agent = t.add_remote_peer(b"meta", mem_desc)
156 with pytest.raises(ValueError, match="mismatch"):
157 t.post_read([0, 1], agent, [5], notif_msg=b"sid")
158
159
160 def test_post_read_calls_make_prepped_xfer_and_transfer():
161 nixl = _make_nixl_mock(prep_side_effect=[77, 88])
162 t = _make_transport(nixl)
163 mem_desc = serialize_mem_descriptor([(0, 64, 0)])
164 agent = t.add_remote_peer(b"meta", mem_desc)
165 nixl.make_prepped_xfer.return_value = 999
166 handle = t.post_read([0, 1], agent, [5, 6], notif_msg=b"sid")
167 nixl.make_prepped_xfer.assert_called_once()
168 nixl.transfer.assert_called_once_with(999)
169 assert handle == 999
170
171
172 # ── release_xfer_handle / deregister ─────────────────────────────────────────
173
174
175 def test_release_xfer_handle_swallows_exceptions():
176 t = _make_transport()
177 t._nixl.release_xfer_handle.side_effect = RuntimeError("stale")
178 t.release_xfer_handle(42) # must not raise
179
180
181 def test_deregister_swallows_exceptions():
182 t = _make_transport()
183 t._nixl.deregister_memory.side_effect = RuntimeError("already gone")
184 t.deregister() # must not raise