[Bugfix][KV Offload] Unlink /dev/shm region after all workers map it (barrier variant of #51317) (#52596)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: NickLucche <nicolo.lucchesi@mistral.ai>
Itay Etelis committed
Aug 31, 2026 at 15:53 UTC
4c58a0c398b056b135b98bd93c644945be7e3109
3 files changed
+247
-34
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+155
@@ -42,6 +42,7 @@ def _make_region(
42
cpu_page_size: int = PAGE_SIZE,
43
num_workers: int = 1,
44
rank: int = 0,
45
+ barrier=None,
46
) -> SharedOffloadRegion:
47
assert cpu_page_size % PAGE_SIZE == 0
48
return SharedOffloadRegion(
@@ -50,6 +51,7 @@ def _make_region(
51
rank=rank,
52
kv_bytes_per_block=num_workers * cpu_page_size,
53
cpu_page_size=cpu_page_size,
54
+ barrier=barrier,
55
)
56
57
@@ -183,6 +185,40 @@ def _mp_race_construct_and_write(
185
done_queue.put({"rank": rank, "error": repr(e)})
186
187
188
+def _mp_barrier_construct_and_hold(
189
+ engine_id: str,
190
+ rank: int,
191
+ num_workers: int,
192
+ barrier,
193
+ fill_value: int,
194
+ done_queue,
195
+) -> None:
196
+ """Construct with a real cross-process barrier, write, then hold the
197
+ mapping (no cleanup) until the parent SIGKILLs this process."""
198
+ try:
199
+ region = SharedOffloadRegion(
200
+ engine_id=engine_id,
201
+ num_blocks=2,
202
+ rank=rank,
203
+ kv_bytes_per_block=num_workers * PAGE_SIZE,
204
+ cpu_page_size=PAGE_SIZE,
205
+ barrier=lambda: barrier.wait(30),
206
+ )
207
+ t = region.create_next_worker_view(PAGE_SIZE)
208
+ t[:, :] = fill_value
209
+ done_queue.put(
210
+ {
211
+ "rank": rank,
212
+ "error": None,
213
+ "inode": os.fstat(region.fd).st_ino,
214
+ "path_exists": os.path.exists(region.mmap_path),
215
+ }
216
+ )
217
+ time.sleep(60) # hold the mapping until killed
218
+ except Exception as e:
219
+ done_queue.put({"rank": rank, "error": repr(e)})
220
+
221
+
222
@pytest.fixture
223
def iid():
224
"""Fresh instance ID for each test."""
@@ -840,3 +876,122 @@ def test_ftruncate_failure_cleans_up_creator(monkeypatch):
876
877
mock_unlink.assert_called_once_with(mmap_path)
878
mock_close.assert_called_once_with(9999)
879
+
880
+
881
+# ---------------------------------------------------------------------------
882
+# Unlink-after-barrier — leak immunity to hard kills
883
+# ---------------------------------------------------------------------------
884
+
885
+
886
+def test_backing_file_unlinked_after_barrier(iid):
887
+ """The file must exist until the barrier releases (late joiners need the
888
+ name) and be gone right after, so no exit path can leak it."""
889
+ path = f"/dev/shm/vllm_offload_{iid}.mmap"
890
+ seen_at_barrier = []
891
+
892
+ region = _make_region(
893
+ iid, barrier=lambda: seen_at_barrier.append(os.path.exists(path))
894
+ )
895
+ try:
896
+ assert seen_at_barrier == [True], "file must exist during rendezvous"
897
+ assert not os.path.exists(path), "name must be dropped after the barrier"
898
+ assert region._creator is False, "nothing left for cleanup() to unlink"
899
+ t = region.create_next_worker_view(PAGE_SIZE)
900
+ t[:, :] = 7
901
+ assert memoryview(region.mmap_obj)[0] == 7, "mapping must stay valid"
902
+ del t
903
+ finally:
904
+ region.cleanup()
905
+ _cleanup_file(path)
906
+
907
+
908
+def test_barrier_failure_unlinks_creator_and_raises(iid):
909
+ """A failed rendezvous must remove the creator's file and re-raise, not
910
+ leave a stub that wedges the next start in _wait_for_file_size."""
911
+ path = f"/dev/shm/vllm_offload_{iid}.mmap"
912
+ with pytest.raises(RuntimeError, match="peer died"):
913
+ _make_region(iid, barrier=MagicMock(side_effect=RuntimeError("peer died")))
914
+ assert not os.path.exists(path)
915
+
916
+
917
+def test_mp_barrier_unlinks_file_and_survives_sigkill(iid):
918
+ """With a real cross-process barrier every worker maps one shared inode,
919
+ the name is gone while they run, and SIGKILL leaks nothing."""
920
+ num_workers = 2
921
+ path = f"/dev/shm/vllm_offload_{iid}.mmap"
922
+ ctx = get_mp_context()
923
+ barrier = ctx.Barrier(num_workers)
924
+ done_queue = ctx.Queue()
925
+ procs = [
926
+ ctx.Process(
927
+ target=_mp_barrier_construct_and_hold,
928
+ args=(iid, rank, num_workers, barrier, rank + 1, done_queue),
929
+ )
930
+ for rank in range(num_workers)
931
+ ]
932
+ try:
933
+ for p in procs:
934
+ p.start()
935
+ results = [done_queue.get(timeout=30) for _ in range(num_workers)]
936
+ for r in results:
937
+ assert r["error"] is None, f"rank {r['rank']}: {r['error']}"
938
+ assert len({r["inode"] for r in results}) == 1, "workers split onto two files"
939
+ assert not any(r["path_exists"] for r in results)
940
+ assert not os.path.exists(path)
941
+ finally:
942
+ for p in procs:
943
+ p.kill() # SIGKILL: no cleanup() runs
944
+ p.join(timeout=10)
945
+
946
+ assert not os.path.exists(path), "SIGKILL must not leak the file"
947
+ with _region(iid) as restarted:
948
+ assert restarted._creator is True, "restart must be able to create anew"
949
+
950
+
951
+def test_setup_failure_before_barrier_releases_peers(iid, monkeypatch):
952
+ """A worker that dies before the rendezvous must still arrive at the
953
+ barrier, or its peers block in the collective until it times out."""
954
+ import vllm.v1.kv_offload.cpu.shared_offload_region as region
955
+
956
+ monkeypatch.setattr(
957
+ region,
958
+ "check_shm_free_space",
959
+ MagicMock(side_effect=RuntimeError("Insufficient space")),
960
+ )
961
+ barrier = MagicMock()
962
+
963
+ with pytest.raises(RuntimeError, match="Insufficient space"):
964
+ _make_region(iid, barrier=barrier)
965
+
966
+ barrier.assert_called_once_with()
967
+ assert not os.path.exists(f"/dev/shm/vllm_offload_{iid}.mmap")
968
+
969
+
970
+def test_mmap_failure_unlinks_creator_before_releasing_peers(iid, monkeypatch):
971
+ """A creator that fails after sizing the file must drop it before arriving
972
+ at the barrier, so the next start does not land on a stale file."""
973
+ path = f"/dev/shm/vllm_offload_{iid}.mmap"
974
+ monkeypatch.setattr("mmap.mmap", MagicMock(side_effect=OSError("mmap")))
975
+ seen_at_barrier = []
976
+
977
+ with pytest.raises(OSError, match="mmap"):
978
+ _make_region(iid, barrier=lambda: seen_at_barrier.append(os.path.exists(path)))
979
+
980
+ assert seen_at_barrier == [False], "file must be gone before peers release"
981
+
982
+
983
+def test_barrier_release_failure_keeps_original_error(iid, monkeypatch):
984
+ """When releasing the peers fails too, the setup error that explains the
985
+ failure must be the one that propagates."""
986
+ import vllm.v1.kv_offload.cpu.shared_offload_region as region
987
+
988
+ monkeypatch.setattr(
989
+ region,
990
+ "check_shm_free_space",
991
+ MagicMock(side_effect=RuntimeError("Insufficient space")),
992
+ )
993
+
994
+ with pytest.raises(RuntimeError, match="Insufficient space"):
995
+ _make_region(iid, barrier=MagicMock(side_effect=TimeoutError("barrier")))
996
+
997
+ assert not os.path.exists(f"/dev/shm/vllm_offload_{iid}.mmap")
vllm/v1/kv_offload/cpu/shared_offload_region.py
+74
-34
@@ -71,7 +71,10 @@ class SharedOffloadRegion:
71
the rest open the existing file and wait until it reaches the expected
72
size. Each worker then mmap()s the full file.
73
74
- File path: /dev/shm/vllm_offload_{engine_id}.mmap
74
+ File path: /dev/shm/vllm_offload_{engine_id}.mmap. When a barrier is
75
+ given, the path is unlinked once every worker has mapped the file, so
76
+ the kernel reclaims the memory when the last worker exits, no matter
77
+ how it exits; mappings taken before the unlink stay valid.
78
"""
79
80
BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE
@@ -83,6 +86,7 @@ class SharedOffloadRegion:
86
rank: int | None,
87
kv_bytes_per_block: int,
88
cpu_page_size: int,
89
+ barrier: Callable[[], None] | None = None,
90
) -> None:
91
self.page_size = mmap.PAGESIZE
92
assert kv_bytes_per_block % self.page_size == 0
@@ -100,44 +104,80 @@ class SharedOffloadRegion:
104
# exclusive upper bound for this worker's area within each row
105
self._worker_area_end = (rank + 1) * cpu_page_size
106
try:
103
- self.fd: int | None = os.open(
104
- self.mmap_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600
107
+ try:
108
+ self.fd: int | None = os.open(
109
+ self.mmap_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600
110
+ )
111
+ except FileExistsError:
112
+ # Joiner path — another worker won O_EXCL. Reopen and wait
113
+ # for the file to reach expected size.
114
+ self.fd = os.open(self.mmap_path, os.O_RDWR)
115
+ try:
116
+ _wait_for_file_size(self.fd, self.total_size_bytes)
117
+ except (TimeoutError, OSError):
118
+ os.close(self.fd)
119
+ raise
120
+ logger.info("Opened existing mmap file %s", self.mmap_path)
121
+ else:
122
+ # Creator path. We won O_EXCL, so we own the file: any
123
+ # failure here must clean up so concurrent joiners don't
124
+ # land on a 0-byte stub and spin in _wait_for_file_size
125
+ # for the full 30 s timeout.
126
+ try:
127
+ check_shm_free_space(self.total_size_bytes)
128
+ os.ftruncate(self.fd, self.total_size_bytes)
129
+ except (RuntimeError, OSError):
130
+ os.unlink(self.mmap_path)
131
+ os.close(self.fd)
132
+ raise
133
+ self._creator = True
134
+ logger.info(
135
+ "Created mmap file %s (%.2f GB)",
136
+ self.mmap_path,
137
+ self.total_size_bytes / 1e9,
138
+ )
139
+
140
+ self.mmap_obj: mmap.mmap | None = mmap.mmap(
141
+ self.fd,
142
+ self.total_size_bytes,
143
+ flags=mmap.MAP_SHARED,
144
+ prot=mmap.PROT_READ | mmap.PROT_WRITE,
145
)
106
- except FileExistsError:
107
- # Joiner path — another worker won O_EXCL. Reopen and wait
108
- # for the file to reach expected size.
109
- self.fd = os.open(self.mmap_path, os.O_RDWR)
146
+ except Exception:
147
+ if self._creator:
148
+ os.unlink(self.mmap_path)
149
+ self._creator = False
150
+ # Peers block inside the barrier until the collective times out if
151
+ # we die before reaching it. Arrive anyway so every worker calls
152
+ # barrier() exactly once and they fail on their own errors instead
153
+ # of hanging; a failure here must not replace ours.
154
+ if barrier is not None:
155
+ try:
156
+ barrier()
157
+ except Exception:
158
+ logger.warning(
159
+ "Failed to release peers waiting at the mmap barrier",
160
+ exc_info=True,
161
+ )
162
+ raise
163
+
164
+ if barrier is not None:
165
+ # Every worker has mapped the file once the barrier releases, so
166
+ # its name is no longer needed and dropping it here means no exit
167
+ # path — including SIGKILL — can leak the file.
168
try:
111
- _wait_for_file_size(self.fd, self.total_size_bytes)
112
- except (TimeoutError, OSError):
169
+ barrier()
170
+ except Exception:
171
+ if self._creator:
172
+ os.unlink(self.mmap_path)
173
+ self._creator = False
174
+ self.mmap_obj.close()
175
os.close(self.fd)
176
raise
115
- logger.info("Opened existing mmap file %s", self.mmap_path)
116
- else:
117
- # Creator path. We won O_EXCL, so we own the file: any
118
- # failure here must clean up so concurrent joiners don't
119
- # land on a 0-byte stub and spin in _wait_for_file_size
120
- # for the full 30 s timeout.
121
- try:
122
- check_shm_free_space(self.total_size_bytes)
123
- os.ftruncate(self.fd, self.total_size_bytes)
124
- except (RuntimeError, OSError):
177
+ if self._creator:
178
os.unlink(self.mmap_path)
126
- os.close(self.fd)
127
- raise
128
- self._creator = True
129
- logger.info(
130
- "Created mmap file %s (%.2f GB)",
131
- self.mmap_path,
132
- self.total_size_bytes / 1e9,
133
- )
134
-
135
- self.mmap_obj: mmap.mmap | None = mmap.mmap(
136
- self.fd,
137
- self.total_size_bytes,
138
- flags=mmap.MAP_SHARED,
139
- prot=mmap.PROT_READ | mmap.PROT_WRITE,
140
- )
179
+ self._creator = False
180
+ logger.info("Unlinked mmap file %s", self.mmap_path)
181
182
populate_write_fn = _get_populate_write_fn(self.mmap_obj)
183
vllm/v1/kv_offload/cpu/spec.py
+18
@@ -24,6 +24,23 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
24
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
25
26
27
+def _all_workers_barrier() -> None:
28
+ """Block until every worker rank has reached this point (gloo cpu group).
29
+
30
+ A superset of the node-local mmap openers suffices: once the barrier
31
+ releases, every worker sharing the region file has mapped it."""
32
+ from vllm.distributed.parallel_state import (
33
+ get_inner_dp_world_group,
34
+ get_world_group,
35
+ )
36
+
37
+ try:
38
+ group = get_inner_dp_world_group()
39
+ except AssertionError:
40
+ group = get_world_group()
41
+ group.barrier()
42
+
43
+
44
class CPUOffloadingSpec(OffloadingSpec):
45
BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
46
@@ -165,6 +182,7 @@ class CPUOffloadingSpec(OffloadingSpec):
182
rank=rank,
183
kv_bytes_per_block=self.kv_bytes_per_chunk,
184
cpu_page_size=self.cpu_page_size_per_worker,
185
+ barrier=_all_workers_barrier,
186
)
187
try:
188
return CPUOffloadingWorker(