@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 283 lines 10.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for resource management in hf3fs_client.py: constructor failure cleanup
4 and idempotent close(). Tests use mock to replace real I/O operations
5 (hf3fs_fuse.io, SharedMemory, os, CUDA).
6 Requires hf3fs_fuse.io to be installed; skipped otherwise.
7 """
8
9 from typing import Any
10 from unittest.mock import MagicMock, patch
11
12 import pytest
13
14 HF3FS_AVAILABLE = True
15 try:
16 from hf3fs_fuse.io import ( # noqa: F401
17 deregister_fd,
18 extract_mount_point,
19 make_ioring,
20 make_iovec,
21 register_fd,
22 )
23
24 from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_client import (
25 Hf3fsClient,
26 )
27 except Exception:
28 HF3FS_AVAILABLE = False
29
30 requires_hf3fs = pytest.mark.skipif(
31 not HF3FS_AVAILABLE,
32 reason="hf3fs_fuse.io is not available on this machine",
33 )
34
35
36 # ---------------------------------------------------------------------------
37 # Helpers
38 # ---------------------------------------------------------------------------
39
40
41 class _FakeShm:
42 """Shared-memory stub matching the multiprocessing.shared_memory.SharedMemory
43 interface used by Hf3fsClient:
44
45 Attributes accessed by the constructor:
46 .buf – memoryview / buffer-protocol object consumed by torch.frombuffer
47 Methods called during normal lifetime:
48 .unlink() – called right after the iovec is set up
49 .close() – called in _release_resources()
50 """
51
52 def __init__(self, size: int = 1024):
53 self._data = bytearray(size)
54 self.buf = memoryview(self._data)
55 self.closed = False
56 self.close_call_count = 0
57 self.unlink_call_count = 0
58
59 def close(self):
60 self.closed = True
61 self.close_call_count += 1
62
63 def unlink(self):
64 self.unlink_call_count += 1
65
66
67 # ===========================================================================
68 # TestHf3fsClientResourceManagement
69 # ===========================================================================
70
71
72 @requires_hf3fs
73 class TestHf3fsClientResourceManagement:
74 """Tests for constructor failure cleanup and idempotent close()."""
75
76 _MOD = "vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_client"
77
78 # ------------------------------------------------------------------
79 # Helper: build a minimal Hf3fsClient bypassing all real I/O so that
80 # we can fully control its internal state.
81 # ------------------------------------------------------------------
82
83 def _make_client(self, tmp_path):
84 """Return a fully-mocked Hf3fsClient with controllable internals."""
85 fake_shm_r = _FakeShm()
86 fake_shm_w = _FakeShm()
87
88 patcher_list: list[Any] = [
89 patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
90 patch(f"{self._MOD}.register_fd"),
91 patch(f"{self._MOD}.deregister_fd"),
92 patch(f"{self._MOD}.extract_mount_point", return_value="/mnt/hf3fs"),
93 patch(f"{self._MOD}.make_ioring", return_value=MagicMock()),
94 patch(f"{self._MOD}.make_iovec", return_value=MagicMock()),
95 patch(
96 "multiprocessing.shared_memory.SharedMemory",
97 side_effect=[fake_shm_r, fake_shm_w],
98 ),
99 patch("os.open", return_value=99),
100 patch("os.ftruncate"),
101 patch("os.close"),
102 patch("os.fsync"),
103 patch("torch.cuda.Stream", return_value=MagicMock()),
104 patch("torch.frombuffer", return_value=MagicMock()),
105 patch("torch.empty", return_value=MagicMock()),
106 ]
107 for p in patcher_list:
108 p.start()
109
110 try:
111 client = Hf3fsClient(
112 path=str(tmp_path / "test.bin"),
113 size=1024,
114 bytes_per_page=256,
115 entries=4,
116 )
117 finally:
118 for p in patcher_list:
119 p.stop()
120
121 # Manually point internal handles to our controllable fakes so that
122 # assertions after close() can inspect them directly.
123 client.shm_r = fake_shm_r
124 client.shm_w = fake_shm_w
125 client.file = 99
126 return client, fake_shm_r, fake_shm_w
127
128 # ------------------------------------------------------------------
129 # close() idempotency
130 # ------------------------------------------------------------------
131
132 def test_close_idempotent_and_handles_cleared(self, tmp_path):
133 """Multiple close() calls must not raise; deregister_fd called exactly
134 once, all handles set to None, shm.close() invoked."""
135 client, shm_r, shm_w = self._make_client(tmp_path)
136
137 with (
138 patch(f"{self._MOD}.deregister_fd") as mock_dereg,
139 patch("os.close"),
140 ):
141 client.close() # first close
142 client.close() # second close — must be no-op
143 client.close() # third close — must be no-op
144
145 assert client._closed is True
146 assert mock_dereg.call_count == 1, (
147 f"deregister_fd called {mock_dereg.call_count} times; expected 1"
148 )
149 for attr in ("iov_r", "iov_w", "ior_r", "ior_w", "shm_r", "shm_w", "file"):
150 assert getattr(client, attr) is None, f"{attr} should be None after close()"
151 assert shm_r.closed is True
152 assert shm_w.closed is True
153
154 def test_flush_after_close_is_noop(self, tmp_path):
155 """flush() after close() must silently do nothing (no fsync call)."""
156 client, _, _ = self._make_client(tmp_path)
157
158 with (
159 patch(f"{self._MOD}.deregister_fd"),
160 patch("os.close"),
161 patch("os.fsync") as mock_fsync,
162 ):
163 client.close()
164 client.flush()
165
166 mock_fsync.assert_not_called()
167
168 # ------------------------------------------------------------------
169 # Constructor failure leaves no leaked resources
170 # ------------------------------------------------------------------
171
172 def test_constructor_failure_after_file_open_cleans_file(self, tmp_path):
173 """If the constructor raises after os.open(), the fd must be closed."""
174 with (
175 patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
176 patch(f"{self._MOD}.register_fd"),
177 patch(f"{self._MOD}.deregister_fd"),
178 patch(
179 f"{self._MOD}.extract_mount_point",
180 side_effect=RuntimeError("mount point not found"),
181 ),
182 patch("os.open", return_value=55),
183 patch("os.ftruncate"),
184 patch("os.close") as mock_os_close,
185 patch("torch.cuda.Stream", return_value=MagicMock()),
186 pytest.raises(RuntimeError, match="mount point not found"),
187 ):
188 Hf3fsClient(
189 path=str(tmp_path / "fail.bin"),
190 size=1024,
191 bytes_per_page=256,
192 entries=4,
193 )
194
195 mock_os_close.assert_called_once_with(55)
196
197 def test_constructor_failure_after_shm_alloc_closes_shm(self, tmp_path):
198 """Constructor raises after SharedMemory creation → both shm objects closed."""
199 fake_shm_r = _FakeShm()
200 fake_shm_w = _FakeShm()
201
202 with (
203 patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
204 patch(f"{self._MOD}.register_fd"),
205 patch(f"{self._MOD}.deregister_fd"),
206 patch(f"{self._MOD}.extract_mount_point", return_value="/mnt/hf3fs"),
207 patch(
208 "multiprocessing.shared_memory.SharedMemory",
209 side_effect=[fake_shm_r, fake_shm_w],
210 ),
211 patch("os.open", return_value=66),
212 patch("os.ftruncate"),
213 patch("os.close"),
214 patch("torch.frombuffer", return_value=MagicMock()),
215 patch("torch.empty", return_value=MagicMock()),
216 patch(
217 f"{self._MOD}.make_ioring",
218 side_effect=RuntimeError("ioring init failed"),
219 ),
220 patch(f"{self._MOD}.make_iovec", return_value=MagicMock()),
221 patch("torch.cuda.Stream", return_value=MagicMock()),
222 pytest.raises(RuntimeError, match="ioring init failed"),
223 ):
224 Hf3fsClient(
225 path=str(tmp_path / "fail2.bin"),
226 size=1024,
227 bytes_per_page=256,
228 entries=4,
229 )
230
231 assert fake_shm_r.closed is True, (
232 "shm_r was not closed after constructor failure"
233 )
234 assert fake_shm_w.closed is True, (
235 "shm_w was not closed after constructor failure"
236 )
237
238 def test_constructor_failure_does_not_close_unallocated_shm(self, tmp_path):
239 """Failure before SharedMemory is created must not raise AttributeError
240 or TypeError from cleanup."""
241 with (
242 patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
243 patch(f"{self._MOD}.register_fd"),
244 patch(f"{self._MOD}.deregister_fd"),
245 patch(
246 f"{self._MOD}.extract_mount_point",
247 side_effect=RuntimeError("early failure"),
248 ),
249 patch("os.open", return_value=77),
250 patch("os.ftruncate"),
251 patch("os.close"),
252 patch("torch.cuda.Stream", return_value=MagicMock()),
253 pytest.raises(RuntimeError, match="early failure"),
254 ):
255 Hf3fsClient(
256 path=str(tmp_path / "early_fail.bin"),
257 size=1024,
258 bytes_per_page=256,
259 entries=4,
260 )
261
262 # ------------------------------------------------------------------
263 # _release_resources on already-cleared state must be a no-op
264 # ------------------------------------------------------------------
265
266 def test_release_resources_on_empty_state_is_safe(self, tmp_path):
267 """_release_resources() on a fully-cleared client must not raise."""
268 client, _, _ = self._make_client(tmp_path)
269
270 with (
271 patch(f"{self._MOD}.deregister_fd"),
272 patch("os.close"),
273 ):
274 client.close() # clears all handles
275
276 with (
277 patch(f"{self._MOD}.deregister_fd") as mock_dereg2,
278 patch("os.close") as mock_os_close2,
279 ):
280 client._release_resources() # must not raise
281
282 mock_dereg2.assert_not_called()
283 mock_os_close2.assert_not_called()