@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,182 lines 41.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for SharedOffloadRegion."""
4
5 import contextlib
6 import ctypes
7 import errno
8 import mmap
9 import os
10 import threading
11 import time
12 import uuid
13 from unittest.mock import MagicMock, call
14
15 import pytest
16
17 from vllm.utils.system_utils import get_mp_context
18 from vllm.v1.kv_offload.cpu import gpu_worker
19 from vllm.v1.kv_offload.cpu import shared_offload_region as region_module
20 from vllm.v1.kv_offload.cpu.shared_offload_region import (
21 SharedOffloadRegion,
22 _wait_for_file_size,
23 )
24
25 PAGE_SIZE = mmap.PAGESIZE
26
27
28 # ---------------------------------------------------------------------------
29 # Helpers / fixtures
30 # ---------------------------------------------------------------------------
31
32
33 @pytest.fixture(autouse=True)
34 def _set_spawn_method(monkeypatch):
35 # On WSL, NVML is not compatible with fork so vLLM auto-overrides the
36 # multiprocessing start method to 'spawn' with a warning. Set it explicitly
37 # here so the override is a no-op and the warning is suppressed.
38 monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
39
40
41 def _make_region(
42 engine_id: str,
43 num_chunks: int = 4,
44 cpu_page_size: int = PAGE_SIZE,
45 num_workers: int = 1,
46 rank: int = 0,
47 barrier=None,
48 ) -> SharedOffloadRegion:
49 assert cpu_page_size % PAGE_SIZE == 0
50 return SharedOffloadRegion(
51 engine_id=engine_id,
52 num_chunks=num_chunks,
53 rank=rank,
54 kv_bytes_per_chunk=num_workers * cpu_page_size,
55 cpu_page_size=cpu_page_size,
56 barrier=barrier,
57 )
58
59
60 def _cleanup_file(path: str) -> None:
61 """Best-effort file removal for test teardown."""
62 with contextlib.suppress(FileNotFoundError):
63 os.unlink(path)
64
65
66 def _page_residency(mmap_obj: mmap.mmap, length: int) -> list[bool]:
67 """Return Linux page-residency bits for a writable mmap."""
68 try:
69 libc = ctypes.CDLL(None, use_errno=True)
70 mincore = libc.mincore
71 except AttributeError:
72 pytest.skip("mincore is unavailable")
73
74 page_count = (length + PAGE_SIZE - 1) // PAGE_SIZE
75 mincore.argtypes = [
76 ctypes.c_void_p,
77 ctypes.c_size_t,
78 ctypes.POINTER(ctypes.c_ubyte),
79 ]
80 mincore.restype = ctypes.c_int
81 vector = (ctypes.c_ubyte * page_count)()
82 address = ctypes.addressof(ctypes.c_ubyte.from_buffer(mmap_obj))
83 result = mincore(ctypes.c_void_p(address), length, vector)
84 if result != 0:
85 raise OSError(ctypes.get_errno())
86 return [bool(value & 1) for value in vector]
87
88
89 @contextlib.contextmanager
90 def _region(engine_id: str, **kwargs):
91 """Context manager: create one region, clean up on exit."""
92 r = _make_region(engine_id, **kwargs)
93 try:
94 yield r
95 finally:
96 r.cleanup()
97 _cleanup_file(r.mmap_path)
98
99
100 @contextlib.contextmanager
101 def _multi_region(
102 engine_id: str,
103 num_workers: int,
104 num_chunks: int = 4,
105 cpu_page_size: int = PAGE_SIZE,
106 ):
107 """Context manager: create one SharedOffloadRegion per rank, clean up on exit."""
108 regions = [
109 SharedOffloadRegion(
110 engine_id=engine_id,
111 num_chunks=num_chunks,
112 rank=rank,
113 kv_bytes_per_chunk=num_workers * cpu_page_size,
114 cpu_page_size=cpu_page_size,
115 )
116 for rank in range(num_workers)
117 ]
118 try:
119 yield regions
120 finally:
121 for r in regions:
122 r.cleanup()
123 _cleanup_file(regions[0].mmap_path)
124
125
126 def _race_construct(
127 engine_id: str,
128 num_workers: int,
129 num_chunks: int = 4,
130 cpu_page_size: int = PAGE_SIZE,
131 ) -> tuple[list[SharedOffloadRegion], list[Exception]]:
132 """Spawn num_workers threads that all race to construct SharedOffloadRegion."""
133 regions: list[SharedOffloadRegion | None] = [None] * num_workers
134 errors: list[Exception] = []
135 barrier = threading.Barrier(num_workers)
136
137 def worker(rank: int) -> None:
138 barrier.wait() # all threads start at the same instant
139 try:
140 regions[rank] = SharedOffloadRegion(
141 engine_id=engine_id,
142 num_chunks=num_chunks,
143 rank=rank,
144 kv_bytes_per_chunk=num_workers * cpu_page_size,
145 cpu_page_size=cpu_page_size,
146 )
147 except Exception as e:
148 errors.append(e)
149
150 threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_workers)]
151 for t in threads:
152 t.start()
153 for t in threads:
154 t.join()
155
156 return [r for r in regions if r is not None], errors
157
158
159 def _mp_race_construct_and_write(
160 engine_id: str,
161 num_chunks: int,
162 rank: int,
163 num_workers: int,
164 cpu_page_size: int,
165 fill_value: int,
166 done_queue,
167 cleanup_queue,
168 ) -> None:
169 """Race to construct a SharedOffloadRegion, write fill_value, then wait
170 for the parent's cleanup signal before tearing down. The wait gives the
171 parent a window to read the raw mmap before the creator removes the file."""
172 try:
173 region = SharedOffloadRegion(
174 engine_id=engine_id,
175 num_chunks=num_chunks,
176 rank=rank,
177 kv_bytes_per_chunk=num_workers * cpu_page_size,
178 cpu_page_size=cpu_page_size,
179 )
180 t = region.create_next_worker_view(cpu_page_size)
181 t[:, :] = fill_value
182 done_queue.put({"rank": rank, "error": None})
183 cleanup_queue.get() # wait for parent's verification to finish
184 del t # release view before cleanup to avoid BufferError
185 region.cleanup()
186 except Exception as e:
187 done_queue.put({"rank": rank, "error": repr(e)})
188
189
190 def _mp_barrier_construct_and_hold(
191 engine_id: str,
192 rank: int,
193 num_workers: int,
194 barrier,
195 fill_value: int,
196 done_queue,
197 replicated: bool = False,
198 ) -> None:
199 """Construct with a real cross-process barrier, write, then hold the
200 mapping (no cleanup) until the parent SIGKILLs this process."""
201 try:
202 populate_calls = 0
203 get_populate_write_fn = region_module._get_populate_write_fn
204
205 def track_population(mmap_obj):
206 nonlocal populate_calls
207 populate_calls += 1
208 return get_populate_write_fn(mmap_obj)
209
210 region_module._get_populate_write_fn = track_population
211 region = SharedOffloadRegion(
212 engine_id=engine_id,
213 num_chunks=2,
214 rank=0 if replicated else rank,
215 kv_bytes_per_chunk=PAGE_SIZE if replicated else num_workers * PAGE_SIZE,
216 cpu_page_size=PAGE_SIZE,
217 barrier=lambda: barrier.wait(30),
218 populate_only_on_creator=replicated,
219 )
220 # The constructor's barrier precedes the creator's unlink.
221 barrier.wait(30)
222 t = region.create_next_worker_view(PAGE_SIZE)
223 if not replicated or rank == 0:
224 t[:, :] = fill_value
225 barrier.wait(30)
226 done_queue.put(
227 {
228 "rank": rank,
229 "error": None,
230 "inode": os.fstat(region.fd).st_ino,
231 "path_exists": os.path.exists(region.mmap_path),
232 "populate_calls": populate_calls,
233 "contents_match": bool((t == (1 if replicated else fill_value)).all()),
234 }
235 )
236 time.sleep(60) # hold the mapping until killed
237 except Exception as e:
238 done_queue.put({"rank": rank, "error": repr(e)})
239
240
241 @pytest.fixture
242 def iid():
243 """Fresh instance ID for each test."""
244 return str(uuid.uuid4())
245
246
247 # ---------------------------------------------------------------------------
248 # create_next_worker_view — shape, stride and storage offset
249 # ---------------------------------------------------------------------------
250
251
252 def test_create_next_worker_view_shape_and_stride(iid):
253 """Returned tensor must have shape (num_chunks, tensor_page_size) and
254 stride (row_stride, 1) where row_stride = cpu_page_size * num_workers."""
255 with _region(iid, num_chunks=4, cpu_page_size=2 * PAGE_SIZE) as r:
256 t = r.create_next_worker_view(PAGE_SIZE)
257 assert t.shape == (4, PAGE_SIZE)
258 # num_workers=1 → row_stride = cpu_page_size
259 assert t.stride() == (2 * PAGE_SIZE, 1)
260 del t
261
262
263 def test_create_next_worker_view_storage_offset_rank0(iid):
264 """rank=0 worker's first tensor must start at byte 0 of the mmap."""
265 with _region(iid, cpu_page_size=PAGE_SIZE, num_workers=2, rank=0) as r:
266 t = r.create_next_worker_view(PAGE_SIZE)
267 assert t.data_ptr() == r._base.data_ptr() # storage_offset == 0
268 del t
269
270
271 def test_create_next_worker_view_storage_offset_rank1(iid):
272 """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap."""
273 with _multi_region(iid, num_workers=2, num_chunks=4) as (r0, r1):
274 t1 = r1.create_next_worker_view(PAGE_SIZE)
275 assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE
276 del t1
277
278
279 def test_create_next_worker_view_row_stride_with_multiple_workers(iid):
280 """With num_workers=4, row_stride must be 4 * cpu_page_size."""
281 with _region(iid, num_chunks=2, num_workers=4) as r:
282 t = r.create_next_worker_view(PAGE_SIZE)
283 assert t.stride(0) == 4 * PAGE_SIZE
284 del t
285
286
287 # ---------------------------------------------------------------------------
288 # create_next_worker_view — cursor advancement
289 # ---------------------------------------------------------------------------
290
291
292 def test_create_next_worker_view_cursor_advances(iid):
293 """Each create_next_worker_view call must advance _worker_offset by
294 tensor_page_size."""
295 with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r:
296 assert r._worker_offset == 0
297 r.create_next_worker_view(PAGE_SIZE)
298 assert r._worker_offset == PAGE_SIZE
299 r.create_next_worker_view(PAGE_SIZE)
300 assert r._worker_offset == 2 * PAGE_SIZE
301 r.create_next_worker_view(PAGE_SIZE)
302 assert r._worker_offset == 3 * PAGE_SIZE # exactly at area end
303
304
305 def test_create_next_worker_view_exact_fill_succeeds(iid):
306 """Allocations whose total exactly equals cpu_page_size must all succeed."""
307 with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
308 r.create_next_worker_view(PAGE_SIZE) # first half
309 r.create_next_worker_view(PAGE_SIZE) # fills to area end — must not raise
310
311
312 # ---------------------------------------------------------------------------
313 # create_next_worker_view — overflow guard
314 # ---------------------------------------------------------------------------
315
316
317 def test_create_next_worker_view_single_overflow_raises(iid):
318 """A single allocation larger than cpu_page_size must raise AssertionError."""
319 with (
320 _region(iid) as r,
321 pytest.raises(AssertionError, match="exceeds worker area end"),
322 ):
323 r.create_next_worker_view(PAGE_SIZE + 1)
324
325
326 def test_create_next_worker_view_cumulative_overflow_raises(iid):
327 """Successive allocations that cumulatively exceed cpu_page_size must raise."""
328 with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
329 r.create_next_worker_view(PAGE_SIZE) # ok — half used
330 r.create_next_worker_view(PAGE_SIZE) # ok — full
331 with pytest.raises(AssertionError, match="exceeds worker area end"):
332 r.create_next_worker_view(1) # one byte too many
333
334
335 def test_create_next_worker_view_overflow_does_not_mutate_cursor(iid):
336 """A failed create_next_worker_view must leave _worker_offset unchanged."""
337 with _region(iid) as r:
338 offset_before = r._worker_offset
339 with pytest.raises(AssertionError):
340 r.create_next_worker_view(PAGE_SIZE + 1)
341 assert r._worker_offset == offset_before
342
343
344 # ---------------------------------------------------------------------------
345 # create_next_worker_view — data correctness and layout
346 # ---------------------------------------------------------------------------
347
348
349 def test_create_next_worker_view_write_visible_in_raw_mmap(iid):
350 """Writes into a create_next_worker_view view must appear at the correct
351 raw mmap offset"""
352 with _region(iid, num_chunks=4) as r:
353 t = r.create_next_worker_view(PAGE_SIZE)
354 t[2, :] = 42 # write to chunk row 2
355
356 raw = memoryview(r.mmap_obj)
357 # num_workers=1 → row_stride = PAGE_SIZE; chunk 2 starts at byte 2*PAGE_SIZE
358 chunk = bytes(raw[2 * PAGE_SIZE : 3 * PAGE_SIZE])
359 assert all(b == 42 for b in chunk)
360 del raw, t
361
362
363 def test_create_next_worker_view_multi_tensor_layout(iid):
364 """Two tensors from the same worker land at consecutive byte offsets per row."""
365 with _region(iid, num_chunks=2, cpu_page_size=2 * PAGE_SIZE) as r:
366 ta = r.create_next_worker_view(PAGE_SIZE)
367 tb = r.create_next_worker_view(PAGE_SIZE)
368
369 ta[:, :] = 1
370 tb[:, :] = 2
371
372 raw = memoryview(r.mmap_obj)
373 for chunk in range(2):
374 row_offset = chunk * 2 * PAGE_SIZE # num_workers=1
375 assert all(b == 1 for b in raw[row_offset : row_offset + PAGE_SIZE])
376 assert all(
377 b == 2 for b in raw[row_offset + PAGE_SIZE : row_offset + 2 * PAGE_SIZE]
378 )
379 del raw, ta, tb
380
381
382 def test_create_next_worker_view_multiprocess_slots(iid):
383 """Each worker process calls create_next_worker_view and writes distinct data;
384 the parent verifies each slot lands at the correct interleaved offset."""
385 num_workers = 2
386 num_chunks = 4
387
388 ctx = get_mp_context()
389 done_queue = ctx.Queue()
390 cleanup_queue = ctx.Queue()
391
392 # Parent is rank 0 (creator); child is rank 1 (joiner).
393 region = SharedOffloadRegion(
394 engine_id=iid,
395 num_chunks=num_chunks,
396 rank=0,
397 kv_bytes_per_chunk=num_workers * PAGE_SIZE,
398 cpu_page_size=PAGE_SIZE,
399 )
400 try:
401 child = ctx.Process(
402 target=_mp_race_construct_and_write,
403 args=(
404 iid,
405 num_chunks,
406 1,
407 num_workers,
408 PAGE_SIZE,
409 22,
410 done_queue,
411 cleanup_queue,
412 ),
413 )
414 child.start()
415
416 t0 = region.create_next_worker_view(PAGE_SIZE)
417 t0[:, :] = 11
418
419 result = done_queue.get(timeout=30)
420 assert result["error"] is None, result["error"]
421
422 raw = memoryview(region.mmap_obj)
423 for chunk in range(num_chunks):
424 row_start = chunk * num_workers * PAGE_SIZE
425 w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
426 w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
427 assert all(b == 11 for b in w0), f"chunk {chunk}: rank0 slot wrong"
428 assert all(b == 22 for b in w1), f"chunk {chunk}: rank1 slot wrong"
429
430 del raw, t0 # release before finally triggers cleanup
431 cleanup_queue.put(True)
432 child.join(timeout=10)
433 assert child.exitcode == 0
434 finally:
435 region.cleanup()
436 _cleanup_file(region.mmap_path)
437
438
439 def test_create_next_worker_view_worker_isolation(iid):
440 """Writes by worker 0 must not affect worker 1's slot and vice versa."""
441 num_workers = 2
442 num_chunks = 4
443 with _multi_region(iid, num_workers=num_workers, num_chunks=num_chunks) as regions:
444 t0 = regions[0].create_next_worker_view(PAGE_SIZE)
445 t1 = regions[1].create_next_worker_view(PAGE_SIZE)
446
447 t0[:, :] = 11
448 t1[:, :] = 22
449
450 raw = memoryview(regions[0].mmap_obj)
451 for chunk in range(num_chunks):
452 row_start = chunk * num_workers * PAGE_SIZE
453 w0 = bytes(raw[row_start : row_start + PAGE_SIZE])
454 w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE])
455 assert all(b == 11 for b in w0), f"chunk {chunk}: worker0 slot corrupted"
456 assert all(b == 22 for b in w1), f"chunk {chunk}: worker1 slot corrupted"
457 del raw, t0, t1 # release before finally triggers cleanup
458
459
460 # ---------------------------------------------------------------------------
461 # Constructor — creator vs joiner semantics
462 # ---------------------------------------------------------------------------
463
464
465 def test_creator_flag_set_on_first_open(iid):
466 """The first worker to open the file must have _creator == True."""
467 with _region(iid) as r:
468 assert r._creator is True
469
470
471 def test_joiner_flag_not_set(iid):
472 """A second worker opening the same file must have _creator == False."""
473 with _multi_region(iid, num_workers=2) as (r0, r1):
474 assert r0._creator is True
475 assert r1._creator is False
476
477
478 def test_file_exists_after_construction(iid):
479 """The mmap file must be present on disk after __init__ completes."""
480 with _region(iid) as r:
481 assert os.path.exists(r.mmap_path)
482
483
484 def test_file_has_correct_size(iid):
485 """The mmap file size on disk must equal total_size_bytes."""
486 with _region(iid, num_chunks=4) as r:
487 assert os.path.getsize(r.mmap_path) == 4 * PAGE_SIZE
488
489
490 def test_madvise_success_selects_madvise_population(iid, monkeypatch):
491 """A successful probe must keep using MADV_POPULATE_WRITE — the fallback
492 helper must never be invoked on a kernel that accepts the advice.
493
494 Spies on both helpers so we can verify call counts: 1 probe call +
495 N populate calls, fallback 0 times.
496 """
497 from vllm.v1.kv_offload.cpu import shared_offload_region as sor
498
499 madvise_calls: list[tuple[int, int, int]] = []
500 fallback_calls: list[tuple[int, int]] = []
501
502 def _spy_madvise(mm, off, ln):
503 madvise_calls.append((off, ln, id(mm)))
504
505 def _spy_fallback(mm, off, ln):
506 fallback_calls.append((off, ln))
507
508 monkeypatch.setattr(sor, "_madvise_populate_write", _spy_madvise)
509 monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
510
511 num_chunks = 3
512 num_workers = 2
513 with _region(iid, num_chunks=num_chunks, num_workers=num_workers, rank=1):
514 # 1 probe (PAGESIZE) + N populate calls (one per chunk per worker column).
515 expected_populate = num_chunks # ranked path: one call per chunk
516 mmap_id = madvise_calls[0][2]
517 assert madvise_calls == [
518 (0, mmap.PAGESIZE, mmap_id),
519 (mmap.PAGESIZE, mmap.PAGESIZE, mmap_id),
520 (3 * mmap.PAGESIZE, mmap.PAGESIZE, mmap_id),
521 (5 * mmap.PAGESIZE, mmap.PAGESIZE, mmap_id),
522 ]
523 assert len(madvise_calls) == 1 + expected_populate
524 assert fallback_calls == [], (
525 "native-path constructor must not invoke the fallback helper"
526 )
527
528
529 def test_madvise_einval_selects_fallback_for_ranked_region(iid, monkeypatch):
530 """An EINVAL probe must select fallback for every ranked chunk."""
531 from vllm.v1.kv_offload.cpu import shared_offload_region as sor
532
533 fallback_calls: list[tuple[int, int]] = []
534 real_fallback = sor._fallback_populate_write
535
536 def _raise_einval(mm, off, ln):
537 raise OSError(errno.EINVAL, "simulated unsupported kernel")
538
539 def _spy_fallback(mm, off, ln):
540 fallback_calls.append((off, ln))
541 return real_fallback(mm, off, ln)
542
543 monkeypatch.setattr(sor, "_madvise_populate_write", _raise_einval)
544 monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
545
546 with _region(iid, num_chunks=3, num_workers=2, rank=1):
547 assert fallback_calls == [
548 (mmap.PAGESIZE, mmap.PAGESIZE),
549 (3 * mmap.PAGESIZE, mmap.PAGESIZE),
550 (5 * mmap.PAGESIZE, mmap.PAGESIZE),
551 ]
552
553
554 def test_madvise_einval_selects_fallback_for_unranked_region(iid, monkeypatch):
555 """An EINVAL probe must select fallback for the whole unranked region."""
556 from vllm.v1.kv_offload.cpu import shared_offload_region as sor
557
558 fallback_calls: list[tuple[int, int]] = []
559 real_fallback = sor._fallback_populate_write
560
561 def _raise_einval(mm, off, ln):
562 raise OSError(errno.EINVAL, "simulated unsupported kernel")
563
564 def _spy_fallback(mm, off, ln):
565 fallback_calls.append((off, ln))
566 return real_fallback(mm, off, ln)
567
568 monkeypatch.setattr(sor, "_madvise_populate_write", _raise_einval)
569 monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
570
571 with _region(iid, num_chunks=3, num_workers=2, rank=None):
572 assert fallback_calls == [(0, 6 * mmap.PAGESIZE)]
573
574
575 def test_fallback_populate_write_preserves_bytes_and_faults_pages():
576 """Fallback preserves existing bytes and touches every target page."""
577 from vllm.v1.kv_offload.cpu import shared_offload_region as sor
578
579 size = 3 * mmap.PAGESIZE
580 if not hasattr(mmap, "MADV_DONTNEED"):
581 pytest.skip("MADV_DONTNEED is unavailable")
582
583 mmap_obj = mmap.mmap(
584 -1,
585 size,
586 flags=mmap.MAP_SHARED,
587 prot=mmap.PROT_READ | mmap.PROT_WRITE,
588 )
589 try:
590 mmap_obj.madvise(mmap.MADV_DONTNEED, 0, size)
591 if any(_page_residency(mmap_obj, size)):
592 pytest.skip("kernel did not discard anonymous mmap pages")
593
594 mmap_obj[0] = 0xAB
595 assert _page_residency(mmap_obj, size) == [True, False, False]
596
597 sor._fallback_populate_write(mmap_obj, 0, size)
598
599 assert _page_residency(mmap_obj, size) == [True, True, True]
600 assert mmap_obj[0] == 0xAB
601 finally:
602 mmap_obj.close()
603
604
605 def test_madvise_unexpected_oserror_propagates(iid, monkeypatch):
606 """Only EINVAL triggers the fallback. Other OSErrors (e.g. EIO) must
607 propagate out of __init__, not be silently masked by the fallback branch.
608 """
609 from vllm.v1.kv_offload.cpu import shared_offload_region as sor
610
611 monkeypatch.setattr(
612 sor,
613 "_madvise_populate_write",
614 lambda mm, off, ln: (_ for _ in ()).throw(
615 OSError(errno.EIO, "simulated I/O failure")
616 ),
617 )
618 with pytest.raises(OSError) as exc_info:
619 _make_region(iid)
620 assert exc_info.value.errno == errno.EIO
621
622
623 # ---------------------------------------------------------------------------
624 # Multi-worker race — concurrent construction
625 # ---------------------------------------------------------------------------
626
627
628 def test_multi_worker_race_exactly_one_creator(iid):
629 """When N threads race to create the same region, exactly one becomes creator."""
630 num_workers = 8
631 regions, errors = _race_construct(iid, num_workers=num_workers)
632 try:
633 assert not errors, f"Workers raised: {errors}"
634 assert len(regions) == num_workers, "Some workers failed to construct"
635
636 creators = [r for r in regions if r._creator]
637 assert len(creators) == 1, f"Expected 1 creator, got {len(creators)}"
638 assert sum(1 for r in regions if not r._creator) == num_workers - 1, (
639 f"Expected {num_workers - 1} non-creators, got "
640 f"{sum(1 for r in regions if not r._creator)}"
641 )
642
643 for r in regions:
644 assert not r.mmap_obj.closed
645 assert r.total_size_bytes == 4 * num_workers * PAGE_SIZE
646 finally:
647 for r in regions:
648 r.cleanup()
649 _cleanup_file(regions[0].mmap_path)
650
651
652 def test_multi_worker_race_shared_memory_visible(iid):
653 """After a concurrent construction race, MAP_SHARED is intact across all workers."""
654 num_workers = 4
655 regions, errors = _race_construct(iid, num_workers=num_workers)
656 assert not errors
657 try:
658 regions[0].mmap_obj[0:1] = b"\xab"
659 for r in regions[1:]:
660 assert memoryview(r.mmap_obj)[0:1] == b"\xab"
661 finally:
662 for r in regions:
663 r.cleanup()
664 _cleanup_file(regions[0].mmap_path)
665
666
667 def test_multiprocess_race_construct_and_write(iid):
668 """N processes race to construct the same SharedOffloadRegion, each writes
669 fill_value = rank+1 into their slot; parent verifies interleaved layout."""
670 num_workers = 4
671 num_chunks = 3
672
673 ctx = get_mp_context()
674 done_queue = ctx.Queue()
675 cleanup_queue = ctx.Queue()
676
677 procs = [
678 ctx.Process(
679 target=_mp_race_construct_and_write,
680 args=(
681 iid,
682 num_chunks,
683 rank,
684 num_workers,
685 PAGE_SIZE,
686 rank + 1,
687 done_queue,
688 cleanup_queue,
689 ),
690 )
691 for rank in range(num_workers)
692 ]
693 for p in procs:
694 p.start()
695
696 results = {}
697 for _ in range(num_workers):
698 r = done_queue.get(timeout=30)
699 results[r["rank"]] = r
700
701 for rank, r in results.items():
702 assert r["error"] is None, f"rank {rank}: {r['error']}"
703
704 # Read the raw file while all workers still hold it open.
705 mmap_path = f"/dev/shm/vllm_offload_{iid}.mmap"
706 with open(mmap_path, "rb") as f:
707 raw = f.read()
708
709 for chunk in range(num_chunks):
710 for w in range(num_workers):
711 slot_start = (chunk * num_workers + w) * PAGE_SIZE
712 slot = raw[slot_start : slot_start + PAGE_SIZE]
713 expected = w + 1 # fill_value = rank + 1
714 assert all(b == expected for b in slot), (
715 f"chunk {chunk}, worker {w}: expected {expected} but got wrong bytes"
716 )
717
718 # Unblock all workers to clean up.
719 for _ in range(num_workers):
720 cleanup_queue.put(True)
721 for p in procs:
722 p.join(timeout=10)
723 assert p.exitcode == 0
724
725
726 # ---------------------------------------------------------------------------
727 # Cleanup
728 # ---------------------------------------------------------------------------
729
730
731 def test_cleanup_creator_all_effects(iid):
732 """cleanup() on the creator closes mmap, closes fd, and removes the file."""
733 r = _make_region(iid)
734 path = r.mmap_path
735 fd = r.fd
736 mmap_obj = r.mmap_obj
737
738 r.cleanup()
739
740 assert mmap_obj.closed, "mmap should be closed after cleanup"
741 assert not os.path.exists(path), "creator should remove the file"
742 with pytest.raises(OSError):
743 os.fstat(fd) # fd should be closed
744
745
746 def test_cleanup_non_creator_all_effects(iid):
747 """cleanup() on a non-creator closes mmap and fd, but leaves the file on disk."""
748 r0 = _make_region(iid) # creator
749 r1 = _make_region(iid) # joiner
750 path = r0.mmap_path
751 fd1 = r1.fd
752 mmap_obj1 = r1.mmap_obj
753 try:
754 r1.cleanup()
755
756 assert mmap_obj1.closed, "mmap should be closed after cleanup"
757 assert os.path.exists(path), "non-creator must not remove the file"
758 with pytest.raises(OSError):
759 os.fstat(fd1) # fd should be closed
760 finally:
761 r0.cleanup()
762 _cleanup_file(path)
763
764
765 def test_cleanup_idempotent(iid):
766 """Calling cleanup() twice must not raise any exception."""
767 r = _make_region(iid)
768 r.cleanup()
769 r.cleanup() # must be a no-op
770
771
772 def test_cleanup_unregisters_every_pinned_chunk(iid, monkeypatch):
773 """cleanup() must release every chunk from a split host registration."""
774 r = _make_region(iid)
775 cudart = MagicMock()
776 cudart.cudaHostUnregister.return_value = MagicMock(value=0)
777 monkeypatch.setattr(region_module, "current_platform", MagicMock())
778 region_module.current_platform.is_cuda_alike.return_value = True
779 monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: cudart)
780 r.pinned_addresses = [0x100000, 0x200000, 0x300000]
781 r.is_pinned = True
782
783 r.cleanup()
784
785 assert cudart.cudaHostUnregister.call_args_list == [
786 call(0x300000),
787 call(0x200000),
788 call(0x100000),
789 ]
790
791
792 @pytest.fixture
793 def host_register(monkeypatch):
794 """Pretend to run on CUDA with host registrations capped at seven pages.
795
796 Registration goes to a mock CudaRTLibrary and the region's cleanup to a
797 separate mock torch runtime. Both stay installed through the cleanup that
798 ``_region`` runs on exit.
799 """
800 cudart = MagicMock()
801 cudart.cudaHostRegister.return_value = 0
802 cudart.cudaHostUnregister.return_value = 0
803 cudart.cudaGetLastError.return_value = 2
804 torch_cudart = MagicMock()
805 torch_cudart.cudaHostUnregister.return_value = MagicMock(value=0)
806 monkeypatch.setattr(gpu_worker.current_platform, "is_cuda_alike", lambda: True)
807 monkeypatch.setattr(gpu_worker, "MAX_HOST_REGISTER_CHUNK_BYTES", 7 * PAGE_SIZE)
808 monkeypatch.setattr(gpu_worker, "CudaRTLibrary", lambda: cudart)
809 monkeypatch.setattr(region_module.torch.cuda, "cudart", lambda: torch_cudart)
810 return cudart, torch_cudart
811
812
813 def test_pin_mmap_region_registers_row_aligned_chunks(iid, host_register):
814 """Chunks end on row boundaries: 3-page rows under a 7-page cap register
815 as 6 + 6 + 3 pages, where a raw byte cap would give 7 + 7 + 1."""
816 cudart, torch_cudart = host_register
817 with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region:
818 gpu_worker.pin_mmap_region(region)
819 base = region._base.data_ptr()
820 assert cudart.mock_calls == [
821 call.cudaHostRegister(base, 6 * PAGE_SIZE),
822 call.cudaHostRegister(base + 6 * PAGE_SIZE, 6 * PAGE_SIZE),
823 call.cudaHostRegister(base + 12 * PAGE_SIZE, 3 * PAGE_SIZE),
824 ]
825 assert region.is_pinned
826
827 assert torch_cudart.cudaHostUnregister.call_args_list == [
828 call(base + 12 * PAGE_SIZE),
829 call(base + 6 * PAGE_SIZE),
830 call(base),
831 ]
832
833
834 @pytest.mark.parametrize("fail_at", [0, 1, 2])
835 def test_pin_mmap_region_failure_leaves_region_pageable(iid, host_register, fail_at):
836 """A failed chunk drains its pending error and unregisters the chunks
837 before it on the same handle, so the region is pinned whole or not at all."""
838 cudart, torch_cudart = host_register
839 cudart.cudaHostRegister.side_effect = [0] * fail_at + [2]
840 with _region(iid, num_chunks=5, cpu_page_size=3 * PAGE_SIZE) as region:
841 gpu_worker.pin_mmap_region(region)
842 base = region._base.data_ptr()
843 assert not region.is_pinned
844 assert region.pinned_addresses == []
845
846 registered = [base + i * 6 * PAGE_SIZE for i in range(fail_at)]
847 failed_size = min(6, 15 - 6 * fail_at) * PAGE_SIZE
848 assert cudart.mock_calls == [
849 *(call.cudaHostRegister(address, 6 * PAGE_SIZE) for address in registered),
850 call.cudaHostRegister(base + fail_at * 6 * PAGE_SIZE, failed_size),
851 call.cudaGetLastError(),
852 *(call.cudaHostUnregister(address) for address in reversed(registered)),
853 ]
854 torch_cudart.cudaHostUnregister.assert_not_called()
855
856
857 def test_pin_mmap_region_without_cuda_runtime_stays_pageable(
858 iid, host_register, monkeypatch
859 ):
860 """A CUDA runtime that cannot be loaded leaves the region pageable instead
861 of failing worker startup."""
862
863 def missing_runtime():
864 raise AssertionError("libcudart is not loaded in the current process")
865
866 monkeypatch.setattr(gpu_worker, "CudaRTLibrary", missing_runtime)
867 with _region(iid) as region:
868 gpu_worker.pin_mmap_region(region)
869 assert not region.is_pinned
870 assert region.pinned_addresses == []
871
872
873 def test_cleanup_after_create_next_worker_view_releases_mmap(iid):
874 """cleanup() must close the mmap even after create_next_worker_view was called.
875 create_next_worker_view returns a view that shares storage with _base; both must be
876 released before mmap.close() can succeed."""
877 r = _make_region(iid)
878 mmap_obj = r.mmap_obj
879
880 t = r.create_next_worker_view(PAGE_SIZE)
881 del t
882
883 r.cleanup()
884
885 assert mmap_obj.closed, "mmap should be closed after releasing the tensor"
886
887
888 # ---------------------------------------------------------------------------
889 # _wait_for_file_size
890 # ---------------------------------------------------------------------------
891
892
893 def test_wait_for_file_size_already_large_enough(tmp_path):
894 """_wait_for_file_size must return immediately when file is already big enough."""
895 fd = os.open(str(tmp_path / "ready.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
896 try:
897 os.ftruncate(fd, PAGE_SIZE)
898 start = time.monotonic()
899 _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0)
900 assert time.monotonic() - start < 0.5
901 finally:
902 os.close(fd)
903
904
905 def test_wait_for_file_size_waits_for_grow(tmp_path):
906 """_wait_for_file_size must return once a background thread grows the file."""
907 fd = os.open(str(tmp_path / "grow.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
908 try:
909
910 def grow():
911 time.sleep(0.05)
912 os.ftruncate(fd, PAGE_SIZE)
913
914 t = threading.Thread(target=grow)
915 t.start()
916 _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0) # must not raise
917 t.join()
918 finally:
919 os.close(fd)
920
921
922 def test_wait_for_file_size_timeout(tmp_path):
923 """_wait_for_file_size must raise TimeoutError when the file never grows."""
924 fd = os.open(str(tmp_path / "stuck.mmap"), os.O_CREAT | os.O_RDWR, 0o600)
925 try:
926 with pytest.raises(TimeoutError):
927 _wait_for_file_size(fd, PAGE_SIZE, timeout=0.1)
928 finally:
929 os.close(fd)
930
931
932 @pytest.mark.skip_global_cleanup
933 def test_wait_for_file_size_rejects_unlinked_file(tmp_path):
934 path = tmp_path / "unlinked.mmap"
935 fd = os.open(str(path), os.O_CREAT | os.O_RDWR, 0o600)
936 try:
937 os.unlink(path)
938 with pytest.raises(RuntimeError, match="creator failed"):
939 _wait_for_file_size(fd, PAGE_SIZE, timeout=1.0)
940 finally:
941 os.close(fd)
942
943
944 # ---------------------------------------------------------------------------
945 # Constructor — capacity validation
946 # ---------------------------------------------------------------------------
947
948
949 @pytest.mark.skip_global_cleanup
950 @pytest.mark.skipif(not os.path.isdir("/dev/shm"), reason="requires /dev/shm")
951 def test_creator_memory_check_runs_only_for_creator(iid):
952 checked_sizes: list[int] = []
953 creator = SharedOffloadRegion(
954 engine_id=iid,
955 num_chunks=4,
956 rank=0,
957 kv_bytes_per_chunk=PAGE_SIZE,
958 cpu_page_size=PAGE_SIZE,
959 creator_memory_check=checked_sizes.append,
960 )
961 joiner: SharedOffloadRegion | None = None
962 try:
963 joiner = SharedOffloadRegion(
964 engine_id=iid,
965 num_chunks=4,
966 rank=0,
967 kv_bytes_per_chunk=PAGE_SIZE,
968 cpu_page_size=PAGE_SIZE,
969 creator_memory_check=checked_sizes.append,
970 )
971 assert checked_sizes == [4 * PAGE_SIZE]
972 finally:
973 if joiner is not None:
974 joiner.cleanup()
975 creator.cleanup()
976 _cleanup_file(creator.mmap_path)
977
978
979 def test_insufficient_space_raises_clear_error(monkeypatch):
980 """A failed creator capacity check must clean up and give a clear error."""
981 import vllm.v1.kv_offload.cpu.shared_offload_region as region
982
983 engine_id = str(uuid.uuid4())
984 mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap"
985 mock_open = MagicMock(return_value=9999)
986 mock_unlink = MagicMock()
987 mock_close = MagicMock()
988 monkeypatch.setattr(region.os, "open", mock_open)
989 monkeypatch.setattr(region.os, "unlink", mock_unlink)
990 monkeypatch.setattr(region.os, "close", mock_close)
991 mock_check = MagicMock(
992 side_effect=RuntimeError("Insufficient space in /dev/shm: 30 GB required.")
993 )
994 monkeypatch.setattr(
995 region,
996 "check_shm_free_space",
997 mock_check,
998 )
999
1000 with pytest.raises(RuntimeError, match="Insufficient space"):
1001 SharedOffloadRegion(
1002 engine_id=engine_id,
1003 num_chunks=4,
1004 rank=0,
1005 kv_bytes_per_chunk=PAGE_SIZE,
1006 cpu_page_size=PAGE_SIZE,
1007 )
1008
1009 mock_unlink.assert_called_once_with(mmap_path)
1010 mock_close.assert_called_once_with(9999)
1011 mock_check.assert_called_once_with(
1012 4 * PAGE_SIZE,
1013 allocation_name="CPU KV offload shared region in /dev/shm",
1014 )
1015
1016
1017 def test_ftruncate_failure_cleans_up_creator(monkeypatch):
1018 """A failed creator ftruncate must close and unlink before re-raising."""
1019 import vllm.v1.kv_offload.cpu.shared_offload_region as region
1020
1021 engine_id = str(uuid.uuid4())
1022 mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap"
1023 mock_unlink = MagicMock()
1024 mock_close = MagicMock()
1025 monkeypatch.setattr(region.os, "open", MagicMock(return_value=9999))
1026 monkeypatch.setattr(region.os, "unlink", mock_unlink)
1027 monkeypatch.setattr(region.os, "close", mock_close)
1028 monkeypatch.setattr(region, "check_shm_free_space", MagicMock())
1029 monkeypatch.setattr(
1030 region.os,
1031 "ftruncate",
1032 MagicMock(side_effect=OSError("ftruncate failed")),
1033 )
1034
1035 with pytest.raises(OSError, match="ftruncate failed"):
1036 SharedOffloadRegion(
1037 engine_id=engine_id,
1038 num_chunks=4,
1039 rank=0,
1040 kv_bytes_per_chunk=PAGE_SIZE,
1041 cpu_page_size=PAGE_SIZE,
1042 )
1043
1044 mock_unlink.assert_called_once_with(mmap_path)
1045 mock_close.assert_called_once_with(9999)
1046
1047
1048 # ---------------------------------------------------------------------------
1049 # Unlink-after-barrier — leak immunity to hard kills
1050 # ---------------------------------------------------------------------------
1051
1052
1053 def test_backing_file_unlinked_after_barrier(iid):
1054 """The file must exist until the barrier releases (late joiners need the
1055 name) and be gone right after, so no exit path can leak it."""
1056 path = f"/dev/shm/vllm_offload_{iid}.mmap"
1057 seen_at_barrier = []
1058
1059 region = _make_region(
1060 iid, barrier=lambda: seen_at_barrier.append(os.path.exists(path))
1061 )
1062 try:
1063 assert seen_at_barrier == [True], "file must exist during rendezvous"
1064 assert not os.path.exists(path), "name must be dropped after the barrier"
1065 assert region._creator is False, "nothing left for cleanup() to unlink"
1066 t = region.create_next_worker_view(PAGE_SIZE)
1067 t[:, :] = 7
1068 assert memoryview(region.mmap_obj)[0] == 7, "mapping must stay valid"
1069 del t
1070 finally:
1071 region.cleanup()
1072 _cleanup_file(path)
1073
1074
1075 def test_barrier_failure_unlinks_creator_and_raises(iid):
1076 """A failed rendezvous must remove the creator's file and re-raise, not
1077 leave a stub that wedges the next start in _wait_for_file_size."""
1078 path = f"/dev/shm/vllm_offload_{iid}.mmap"
1079 with pytest.raises(RuntimeError, match="peer died"):
1080 _make_region(iid, barrier=MagicMock(side_effect=RuntimeError("peer died")))
1081 assert not os.path.exists(path)
1082
1083
1084 @pytest.mark.parametrize("replicated", [False, True])
1085 def test_mp_barrier_unlinks_file_and_survives_sigkill(iid, replicated):
1086 """With a real cross-process barrier every worker maps one shared inode,
1087 the name is gone while they run, and SIGKILL leaks nothing."""
1088 num_workers = 2
1089 path = f"/dev/shm/vllm_offload_{iid}.mmap"
1090 ctx = get_mp_context()
1091 barrier = ctx.Barrier(num_workers)
1092 done_queue = ctx.Queue()
1093 procs = [
1094 ctx.Process(
1095 target=_mp_barrier_construct_and_hold,
1096 args=(iid, rank, num_workers, barrier, rank + 1, done_queue, replicated),
1097 )
1098 for rank in range(num_workers)
1099 ]
1100 try:
1101 for p in procs:
1102 p.start()
1103 results = [done_queue.get(timeout=30) for _ in range(num_workers)]
1104 for r in results:
1105 assert r["error"] is None, f"rank {r['rank']}: {r['error']}"
1106 assert len({r["inode"] for r in results}) == 1, "workers split onto two files"
1107 assert not any(r["path_exists"] for r in results)
1108 assert all(r["contents_match"] for r in results)
1109 assert sum(r["populate_calls"] for r in results) == (
1110 1 if replicated else num_workers
1111 )
1112 assert not os.path.exists(path)
1113 finally:
1114 for p in procs:
1115 p.kill() # SIGKILL: no cleanup() runs
1116 p.join(timeout=10)
1117
1118 assert not os.path.exists(path), "SIGKILL must not leak the file"
1119 with _region(iid) as restarted:
1120 assert restarted._creator is True, "restart must be able to create anew"
1121
1122
1123 @pytest.mark.parametrize("failure_phase", ["shm", "memory", "population"])
1124 def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch, failure_phase):
1125 """A worker that dies before the rendezvous must still arrive at the
1126 barrier, or its peers block in the collective until it times out."""
1127 failure = MagicMock(side_effect=ValueError("Insufficient space"))
1128 memory_check = failure if failure_phase == "memory" else None
1129 if failure_phase == "shm":
1130 monkeypatch.setattr(region_module, "check_shm_free_space", failure)
1131 elif failure_phase == "population":
1132 monkeypatch.setattr(region_module, "_get_populate_write_fn", failure)
1133 path = f"/dev/shm/vllm_offload_{iid}.mmap"
1134 seen_at_barrier = []
1135 barrier = MagicMock(
1136 side_effect=lambda: seen_at_barrier.append(os.path.exists(path))
1137 )
1138
1139 with pytest.raises(ValueError, match="Insufficient space"):
1140 SharedOffloadRegion(
1141 engine_id=iid,
1142 num_chunks=1,
1143 rank=0,
1144 kv_bytes_per_chunk=PAGE_SIZE,
1145 cpu_page_size=PAGE_SIZE,
1146 barrier=barrier,
1147 creator_memory_check=memory_check,
1148 populate_only_on_creator=True,
1149 )
1150
1151 barrier.assert_called_once_with()
1152 assert seen_at_barrier == [False]
1153
1154
1155 def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch):
1156 """A creator that fails after sizing the file must drop it before arriving
1157 at the barrier, so the next start does not land on a stale file."""
1158 path = f"/dev/shm/vllm_offload_{iid}.mmap"
1159 monkeypatch.setattr("mmap.mmap", MagicMock(side_effect=OSError("mmap")))
1160 seen_at_barrier = []
1161
1162 with pytest.raises(OSError, match="mmap"):
1163 _make_region(iid, barrier=lambda: seen_at_barrier.append(os.path.exists(path)))
1164
1165 assert seen_at_barrier == [False], "file must be gone before peers release"
1166
1167
1168 def test_barrier_release_failure_keeps_original_error(iid, monkeypatch):
1169 """When releasing the peers fails too, the setup error that explains the
1170 failure must be the one that propagates."""
1171 import vllm.v1.kv_offload.cpu.shared_offload_region as region
1172
1173 monkeypatch.setattr(
1174 region,
1175 "check_shm_free_space",
1176 MagicMock(side_effect=RuntimeError("Insufficient space")),
1177 )
1178
1179 with pytest.raises(RuntimeError, match="Insufficient space"):
1180 _make_region(iid, barrier=MagicMock(side_effect=TimeoutError("barrier")))
1181
1182 assert not os.path.exists(f"/dev/shm/vllm_offload_{iid}.mmap")