[Bugfix] Fail fast with a clear error when CPU offload region exceeds available space (#50358)
Signed-off-by: Alex <jihui.huang@daocloud.io> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
AlexHuang committed
Aug 6, 2026 at 01:59 UTC
9833aa53d5a9559b4cc8d174ae49d6e282c30423
2 files changed
+95
-6
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+70
@@ -8,6 +8,7 @@ import os
8
import threading
9
import time
10
import uuid
11
+from unittest.mock import MagicMock
12
13
import pytest
14
@@ -610,3 +611,72 @@ def test_wait_for_file_size_timeout(tmp_path):
611
_wait_for_file_size(fd, PAGE_SIZE, timeout=0.1)
612
finally:
613
os.close(fd)
614
+
615
+
616
+# ---------------------------------------------------------------------------
617
+# Constructor — capacity validation
618
+# ---------------------------------------------------------------------------
619
+
620
+
621
+def test_insufficient_space_raises_clear_error(monkeypatch):
622
+ """A failed creator capacity check must clean up and give a clear error."""
623
+ import vllm.v1.kv_offload.cpu.shared_offload_region as region
624
+
625
+ engine_id = str(uuid.uuid4())
626
+ mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap"
627
+ mock_open = MagicMock(return_value=9999)
628
+ mock_unlink = MagicMock()
629
+ mock_close = MagicMock()
630
+ monkeypatch.setattr(region.os, "open", mock_open)
631
+ monkeypatch.setattr(region.os, "unlink", mock_unlink)
632
+ monkeypatch.setattr(region.os, "close", mock_close)
633
+ monkeypatch.setattr(
634
+ region,
635
+ "check_shm_free_space",
636
+ lambda *a, **kw: (_ for _ in ()).throw(
637
+ RuntimeError("Insufficient space in /dev/shm: 30 GB required.")
638
+ ),
639
+ )
640
+
641
+ with pytest.raises(RuntimeError, match="Insufficient space"):
642
+ SharedOffloadRegion(
643
+ engine_id=engine_id,
644
+ num_blocks=4,
645
+ rank=0,
646
+ kv_bytes_per_block=PAGE_SIZE,
647
+ cpu_page_size=PAGE_SIZE,
648
+ )
649
+
650
+ mock_unlink.assert_called_once_with(mmap_path)
651
+ mock_close.assert_called_once_with(9999)
652
+
653
+
654
+def test_ftruncate_failure_cleans_up_creator(monkeypatch):
655
+ """A failed creator ftruncate must close and unlink before re-raising."""
656
+ import vllm.v1.kv_offload.cpu.shared_offload_region as region
657
+
658
+ engine_id = str(uuid.uuid4())
659
+ mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap"
660
+ mock_unlink = MagicMock()
661
+ mock_close = MagicMock()
662
+ monkeypatch.setattr(region.os, "open", MagicMock(return_value=9999))
663
+ monkeypatch.setattr(region.os, "unlink", mock_unlink)
664
+ monkeypatch.setattr(region.os, "close", mock_close)
665
+ monkeypatch.setattr(region, "check_shm_free_space", MagicMock())
666
+ monkeypatch.setattr(
667
+ region.os,
668
+ "ftruncate",
669
+ MagicMock(side_effect=OSError("ftruncate failed")),
670
+ )
671
+
672
+ with pytest.raises(OSError, match="ftruncate failed"):
673
+ SharedOffloadRegion(
674
+ engine_id=engine_id,
675
+ num_blocks=4,
676
+ rank=0,
677
+ kv_bytes_per_block=PAGE_SIZE,
678
+ cpu_page_size=PAGE_SIZE,
679
+ )
680
+
681
+ mock_unlink.assert_called_once_with(mmap_path)
682
+ mock_close.assert_called_once_with(9999)
vllm/v1/kv_offload/cpu/shared_offload_region.py
+25
-6
@@ -6,6 +6,9 @@ import time
6
7
import torch
8
9
+from vllm.distributed.device_communicators.shm_broadcast import (
10
+ check_shm_free_space,
11
+)
12
from vllm.logger import init_logger
13
from vllm.platforms import current_platform
14
@@ -62,21 +65,37 @@ class SharedOffloadRegion:
65
# exclusive upper bound for this worker's area within each row
66
self._worker_area_end = (rank + 1) * cpu_page_size
67
try:
65
- # Exclusive create — only one worker succeeds
68
self.fd: int | None = os.open(
69
self.mmap_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600
70
)
69
- os.ftruncate(self.fd, self.total_size_bytes)
71
+ except FileExistsError:
72
+ # Joiner path — another worker won O_EXCL. Reopen and wait
73
+ # for the file to reach expected size.
74
+ self.fd = os.open(self.mmap_path, os.O_RDWR)
75
+ try:
76
+ _wait_for_file_size(self.fd, self.total_size_bytes)
77
+ except (TimeoutError, OSError):
78
+ os.close(self.fd)
79
+ raise
80
+ logger.info("Opened existing mmap file %s", self.mmap_path)
81
+ else:
82
+ # Creator path. We won O_EXCL, so we own the file: any
83
+ # failure here must clean up so concurrent joiners don't
84
+ # land on a 0-byte stub and spin in _wait_for_file_size
85
+ # for the full 30 s timeout.
86
+ try:
87
+ check_shm_free_space(self.total_size_bytes)
88
+ os.ftruncate(self.fd, self.total_size_bytes)
89
+ except (RuntimeError, OSError):
90
+ os.unlink(self.mmap_path)
91
+ os.close(self.fd)
92
+ raise
93
self._creator = True
94
logger.info(
95
"Created mmap file %s (%.2f GB)",
96
self.mmap_path,
97
self.total_size_bytes / 1e9,
98
)
76
- except FileExistsError:
77
- self.fd = os.open(self.mmap_path, os.O_RDWR)
78
- _wait_for_file_size(self.fd, self.total_size_bytes)
79
- logger.info("Opened existing mmap file %s", self.mmap_path)
99
100
self.mmap_obj: mmap.mmap | None = mmap.mmap(
101
self.fd,