[Bugfix][KV Offload] Check cgroup memory before SHM allocation (#54014)
Signed-off-by: Alex <jihui.huang@daocloud.io> Signed-off-by: AlexHuang <alex.tech.lab@outlook.com> Signed-off-by: Alex <alex.tech.lab@outlook.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com>
AlexHuang committed
Sep 16, 2026 at 14:17 UTC
0f8fa53acf43a9a6a7a1f7e6b196950c78cd6d7a
7 files changed
+343
-34
tests/distributed/test_shm_broadcast.py
+32
-3
@@ -694,14 +694,43 @@ def test_check_shm_free_space_raises_when_insufficient(tmp_path):
694
695
696
def test_check_shm_free_space_passes_when_sufficient(tmp_path):
697
- with mock.patch.object(
698
- shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(512 << 20)
697
+ with (
698
+ mock.patch.object(
699
+ shm_broadcast.shutil,
700
+ "disk_usage",
701
+ return_value=_fake_disk_usage(512 << 20),
702
+ ),
703
+ mock.patch.object(shm_broadcast, "check_cgroup_memory_available"),
704
):
705
check_shm_free_space(240 << 20, shm_path=str(tmp_path))
706
707
708
def test_check_shm_free_space_skipped_when_path_missing(tmp_path):
704
- check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist"))
709
+ with mock.patch.object(shm_broadcast, "check_cgroup_memory_available"):
710
+ check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist"))
711
+
712
+
713
+def test_check_shm_free_space_checks_cgroup(tmp_path):
714
+ with (
715
+ mock.patch.object(
716
+ shm_broadcast.shutil,
717
+ "disk_usage",
718
+ return_value=_fake_disk_usage(512 << 20),
719
+ ),
720
+ mock.patch.object(
721
+ shm_broadcast, "check_cgroup_memory_available"
722
+ ) as check_cgroup,
723
+ ):
724
+ check_shm_free_space(
725
+ 240 << 20,
726
+ shm_path=str(tmp_path),
727
+ allocation_name="SHM mmap",
728
+ )
729
+
730
+ check_cgroup.assert_called_once_with(
731
+ 240 << 20,
732
+ "SHM mmap",
733
+ )
734
735
736
def test_shm_ring_buffer_creation_checks_free_space():
tests/utils_/test_cpu_resource_utils.py
new
+189
@@ -0,0 +1,189 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Cgroup memory readers and SHM allocation preflight checks."""
4
+
5
+import sys
6
+from io import StringIO
7
+from unittest import mock
8
+
9
+import pytest
10
+
11
+from vllm.utils import cpu_resource_utils as cru
12
+from vllm.utils.mem_constants import GiB_bytes
13
+
14
+_V2_LIMIT_PATH = "/sys/fs/cgroup/memory.max"
15
+_V2_USAGE_PATH = "/sys/fs/cgroup/memory.current"
16
+_V1_LIMIT_PATH = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
17
+_V1_USAGE_PATH = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
18
+
19
+
20
+@pytest.fixture(autouse=True)
21
+def _clear_cgroup_cache():
22
+ cru.get_cgroup_memory_limit.cache_clear()
23
+ yield
24
+ cru.get_cgroup_memory_limit.cache_clear()
25
+
26
+
27
+def _stub_files(monkeypatch, files: dict):
28
+ """Stub ``open()`` for a fixed set of paths; ``None`` -> OSError."""
29
+ real_open = open
30
+
31
+ def fake_open(path, *args, **kwargs):
32
+ if path in files:
33
+ content = files[path]
34
+ if content is None:
35
+ raise OSError(f"no such file: {path}")
36
+ return StringIO(content)
37
+ return real_open(path, *args, **kwargs)
38
+
39
+ monkeypatch.setattr("builtins.open", fake_open)
40
+
41
+
42
+def test_cgroup_v2_limit_and_usage(monkeypatch):
43
+ monkeypatch.setattr(sys, "platform", "linux")
44
+ _stub_files(
45
+ monkeypatch,
46
+ {
47
+ _V2_LIMIT_PATH: f"{20 * GiB_bytes}\n",
48
+ _V2_USAGE_PATH: f"{5 * GiB_bytes}\n",
49
+ },
50
+ )
51
+
52
+ assert cru.get_cgroup_memory_limit() == 20 * GiB_bytes
53
+ assert cru.get_cgroup_memory_usage() == 5 * GiB_bytes
54
+
55
+
56
+def test_cgroup_v2_unlimited_falls_back_to_v1(monkeypatch):
57
+ monkeypatch.setattr(sys, "platform", "linux")
58
+ _stub_files(
59
+ monkeypatch,
60
+ {
61
+ _V2_LIMIT_PATH: "max\n",
62
+ _V1_LIMIT_PATH: f"{8 * GiB_bytes}\n",
63
+ _V1_USAGE_PATH: f"{1 * GiB_bytes}\n",
64
+ },
65
+ )
66
+
67
+ assert cru.get_cgroup_memory_limit() == 8 * GiB_bytes
68
+ assert cru.get_cgroup_memory_usage() == 1 * GiB_bytes
69
+
70
+
71
+def test_cgroup_v1_limit_and_usage(monkeypatch):
72
+ monkeypatch.setattr(sys, "platform", "linux")
73
+ _stub_files(
74
+ monkeypatch,
75
+ {
76
+ _V2_LIMIT_PATH: None,
77
+ _V1_LIMIT_PATH: f"{8 * GiB_bytes}\n",
78
+ _V1_USAGE_PATH: f"{2 * GiB_bytes}\n",
79
+ },
80
+ )
81
+
82
+ assert cru.get_cgroup_memory_limit() == 8 * GiB_bytes
83
+ assert cru.get_cgroup_memory_usage() == 2 * GiB_bytes
84
+
85
+
86
+def test_cgroup_v1_unlimited_sentinel_is_ignored(monkeypatch):
87
+ """An unlimited cgroup v1 sentinel must not be treated as a real limit."""
88
+ monkeypatch.setattr(sys, "platform", "linux")
89
+ _stub_files(
90
+ monkeypatch,
91
+ {_V2_LIMIT_PATH: None, _V1_LIMIT_PATH: f"{(1 << 63) - 1}\n"},
92
+ )
93
+
94
+ assert cru.get_cgroup_memory_limit() is None
95
+ assert cru.get_cgroup_memory_usage() is None
96
+
97
+
98
+def test_cgroup_no_limit_files_present(monkeypatch):
99
+ monkeypatch.setattr(sys, "platform", "linux")
100
+ _stub_files(monkeypatch, {_V2_LIMIT_PATH: None, _V1_LIMIT_PATH: None})
101
+
102
+ assert cru.get_cgroup_memory_limit() is None
103
+ assert cru.get_cgroup_memory_usage() is None
104
+
105
+
106
+def test_cgroup_skipped_on_non_linux(monkeypatch):
107
+ monkeypatch.setattr(sys, "platform", "darwin")
108
+ _stub_files(monkeypatch, {_V2_LIMIT_PATH: f"{1 * GiB_bytes}\n"})
109
+
110
+ assert cru.get_cgroup_memory_limit() is None
111
+ assert cru.get_cgroup_memory_usage() is None
112
+
113
+
114
+def test_cgroup_limit_is_cached_until_explicitly_cleared(monkeypatch):
115
+ monkeypatch.setattr(sys, "platform", "linux")
116
+ files = {
117
+ _V2_LIMIT_PATH: f"{1 * GiB_bytes}\n",
118
+ _V2_USAGE_PATH: f"{512 << 20}\n",
119
+ }
120
+ _stub_files(monkeypatch, files)
121
+
122
+ assert cru.get_cgroup_memory_limit() == 1 * GiB_bytes
123
+ files[_V2_LIMIT_PATH] = f"{2 * GiB_bytes}\n"
124
+ assert cru.get_cgroup_memory_limit() == 1 * GiB_bytes
125
+
126
+ cru.get_cgroup_memory_limit.cache_clear()
127
+ assert cru.get_cgroup_memory_limit() == 2 * GiB_bytes
128
+
129
+
130
+def test_cgroup_usage_is_read_without_cache(monkeypatch):
131
+ monkeypatch.setattr(sys, "platform", "linux")
132
+ files = {
133
+ _V2_LIMIT_PATH: f"{1 * GiB_bytes}\n",
134
+ _V2_USAGE_PATH: f"{512 << 20}\n",
135
+ }
136
+ _stub_files(monkeypatch, files)
137
+
138
+ assert cru.get_cgroup_memory_usage() == 512 << 20
139
+ files[_V2_USAGE_PATH] = f"{768 << 20}\n"
140
+ assert cru.get_cgroup_memory_usage() == 768 << 20
141
+
142
+
143
+def test_check_cgroup_memory_available_warns_on_low_headroom(monkeypatch):
144
+ monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: 1 << 30)
145
+ monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: 512 << 20)
146
+
147
+ with (
148
+ mock.patch.object(cru.logger, "debug") as log_debug,
149
+ mock.patch.object(cru.logger, "warning") as log_warning,
150
+ ):
151
+ cru.check_cgroup_memory_available(600 << 20, "mmap")
152
+
153
+ log_debug.assert_not_called()
154
+ log_warning.assert_called_once()
155
+ assert "current headroom is below" in log_warning.call_args.args[-1]
156
+
157
+
158
+def test_check_cgroup_memory_available_logs_success_at_debug(monkeypatch):
159
+ monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: 1 << 30)
160
+ monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: 512 << 20)
161
+
162
+ with (
163
+ mock.patch.object(cru.logger, "debug") as log_debug,
164
+ mock.patch.object(cru.logger, "warning") as log_warning,
165
+ ):
166
+ cru.check_cgroup_memory_available(256 << 20, "mmap")
167
+
168
+ log_debug.assert_called_once()
169
+ log_warning.assert_not_called()
170
+
171
+
172
+@pytest.mark.parametrize(
173
+ ("limit", "usage"),
174
+ [(None, 512 << 20), (1 << 30, None)],
175
+)
176
+def test_check_cgroup_memory_available_skips_without_snapshot(
177
+ monkeypatch, limit, usage
178
+):
179
+ monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: limit)
180
+ monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: usage)
181
+
182
+ with (
183
+ mock.patch.object(cru.logger, "debug") as log_debug,
184
+ mock.patch.object(cru.logger, "warning") as log_warning,
185
+ ):
186
+ cru.check_cgroup_memory_available(1 << 60, "mmap")
187
+
188
+ log_debug.assert_not_called()
189
+ log_warning.assert_not_called()
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+8
-3
@@ -906,12 +906,13 @@ def test_insufficient_space_raises_clear_error(monkeypatch):
906
monkeypatch.setattr(region.os, "open", mock_open)
907
monkeypatch.setattr(region.os, "unlink", mock_unlink)
908
monkeypatch.setattr(region.os, "close", mock_close)
909
+ mock_check = MagicMock(
910
+ side_effect=RuntimeError("Insufficient space in /dev/shm: 30 GB required.")
911
+ )
912
monkeypatch.setattr(
913
region,
914
"check_shm_free_space",
912
- lambda *a, **kw: (_ for _ in ()).throw(
913
- RuntimeError("Insufficient space in /dev/shm: 30 GB required.")
914
- ),
915
+ mock_check,
916
)
917
918
with pytest.raises(RuntimeError, match="Insufficient space"):
@@ -925,6 +926,10 @@ def test_insufficient_space_raises_clear_error(monkeypatch):
926
927
mock_unlink.assert_called_once_with(mmap_path)
928
mock_close.assert_called_once_with(9999)
929
+ mock_check.assert_called_once_with(
930
+ 4 * PAGE_SIZE,
931
+ allocation_name="CPU KV offload shared region in /dev/shm",
932
+ )
933
934
935
def test_ftruncate_failure_cleans_up_creator(monkeypatch):
vllm/distributed/device_communicators/shm_broadcast.py
+26
-14
@@ -34,6 +34,7 @@ import vllm.envs as envs
34
from vllm.distributed.utils import StatelessProcessGroup, sched_yield
35
from vllm.logger import init_logger
36
from vllm.platforms import current_platform
37
+from vllm.utils.cpu_resource_utils import check_cgroup_memory_available
38
from vllm.utils.network_utils import (
39
get_ip,
40
get_open_zmq_inproc_path,
@@ -224,26 +225,37 @@ class SpinCondition:
225
SHM_PATH = "/dev/shm"
226
227
227
-def check_shm_free_space(required_bytes: int, shm_path: str = SHM_PATH) -> None:
228
- """Raise if ``shm_path`` cannot fit a ``required_bytes`` shared segment.
228
+def check_shm_free_space(
229
+ required_bytes: int,
230
+ shm_path: str = SHM_PATH,
231
+ *,
232
+ allocation_name: str = "shared-memory allocation",
233
+) -> None:
234
+ """Raise if SHM cannot fit a shared segment and log cgroup headroom.
235
236
Args:
237
required_bytes: Size of the shared-memory segment to be created.
232
- shm_path: Mount point backing POSIX shared memory; skipped if absent.
238
+ shm_path: Mount point backing POSIX shared memory; its filesystem
239
+ check is skipped if absent.
240
+ allocation_name: Human-readable name used in errors and logs.
241
242
Raises:
235
- RuntimeError: If ``required_bytes`` exceeds the free space.
243
+ RuntimeError: If the SHM filesystem has insufficient space.
244
"""
237
- if not os.path.isdir(shm_path):
238
- return
239
- free_bytes = shutil.disk_usage(shm_path).free
240
- if required_bytes <= free_bytes:
241
- return
242
- mib = 1 << 20
243
- raise RuntimeError(
244
- f"Insufficient space in {shm_path}: {required_bytes / mib:.0f} MiB "
245
- f"required, {free_bytes / mib:.0f} MiB free. Increase {shm_path} "
246
- "(e.g. --shm-size or --ipc=host)."
245
+ if os.path.isdir(shm_path):
246
+ free_bytes = shutil.disk_usage(shm_path).free
247
+ if required_bytes > free_bytes:
248
+ mib = 1 << 20
249
+ raise RuntimeError(
250
+ f"Insufficient space in {shm_path} for {allocation_name}: "
251
+ f"{required_bytes / mib:.0f} MiB required, "
252
+ f"{free_bytes / mib:.0f} MiB free. Increase {shm_path} "
253
+ "(e.g. --shm-size or --ipc=host)."
254
+ )
255
+
256
+ check_cgroup_memory_available(
257
+ required_bytes,
258
+ allocation_name,
259
)
260
261
vllm/model_executor/model_loader/weight_utils.py
+6
-2
@@ -727,11 +727,15 @@ def _get_available_ram_bytes() -> int:
727
728
host_available = psutil.virtual_memory().available
729
730
- from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit
730
+ from vllm.utils.cpu_resource_utils import (
731
+ get_cgroup_memory_limit,
732
+ get_cgroup_memory_usage,
733
+ )
734
732
- cgroup_limit, cgroup_usage = get_cgroup_memory_limit()
735
+ cgroup_limit = get_cgroup_memory_limit()
736
if cgroup_limit is None:
737
return host_available
738
+ cgroup_usage = get_cgroup_memory_usage()
739
cgroup_available = (
740
cgroup_limit if cgroup_usage is None else max(0, cgroup_limit - cgroup_usage)
741
)
vllm/utils/cpu_resource_utils.py
+78
-11
@@ -12,8 +12,12 @@ from functools import cache
12
import psutil
13
import regex as re
14
15
+from vllm.logger import init_logger
16
+
17
DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES"
18
19
+logger = init_logger(__name__)
20
+
21
22
@dataclass
23
class LogicalCPUInfo:
@@ -63,21 +67,20 @@ def _read_int_file(path: str) -> int | None:
67
68
69
@cache
66
-def get_cgroup_memory_limit() -> tuple[int | None, int | None]:
67
- """Return (limit, usage) in bytes from cgroup, or (None, None).
70
+def get_cgroup_memory_limit() -> int | None:
71
+ """Return the cgroup memory limit in bytes, or None.
72
69
- Supports both cgroup v2 (unified) and v1. Returns (None, None) when
73
+ Supports both cgroup v2 (unified) and v1. Returns None when
74
not running under a constrained cgroup (e.g. bare metal, or limit
75
reported as `max`/an unrealistically large value).
76
"""
77
if sys.platform != "linux":
74
- return None, None
78
+ return None
79
80
# cgroup v2 unified hierarchy
81
v2_limit = _read_int_file("/sys/fs/cgroup/memory.max")
82
if v2_limit is not None:
79
- v2_usage = _read_int_file("/sys/fs/cgroup/memory.current")
80
- return v2_limit, v2_usage
83
+ return v2_limit
84
85
# cgroup v1
86
v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes")
@@ -85,11 +88,74 @@ def get_cgroup_memory_limit() -> tuple[int | None, int | None]:
88
# cgroup v1 reports a huge sentinel (close to PAGE_COUNTER_MAX)
89
# when unlimited. Treat absurdly large values as "no limit".
90
if v1_limit >= (1 << 62):
88
- return None, None
89
- v1_usage = _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes")
90
- return v1_limit, v1_usage
91
+ return None
92
+ return v1_limit
93
+
94
+ return None
95
+
96
92
- return None, None
97
+def get_cgroup_memory_usage() -> int | None:
98
+ """Return the current cgroup memory usage in bytes, or None.
99
+
100
+ The usage value is intentionally read on every call because cgroup
101
+ memory usage changes while the process is running.
102
+ """
103
+ if sys.platform != "linux":
104
+ return None
105
+
106
+ # cgroup v2 unified hierarchy
107
+ if _read_int_file("/sys/fs/cgroup/memory.max") is not None:
108
+ return _read_int_file("/sys/fs/cgroup/memory.current")
109
+
110
+ # cgroup v1
111
+ v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes")
112
+ if v1_limit is not None and v1_limit < (1 << 62):
113
+ return _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes")
114
+
115
+ return None
116
+
117
+
118
+def check_cgroup_memory_available(
119
+ required_bytes: int,
120
+ allocation_name: str,
121
+) -> None:
122
+ """Log cgroup memory headroom for an upcoming allocation.
123
+
124
+ Args:
125
+ required_bytes: Bytes required by the allocation.
126
+ allocation_name: Human-readable name used in log messages.
127
+
128
+ Low headroom logs a warning, but does not reject the allocation because
129
+ cgroup usage can include reclaimable memory. If the cgroup limit or usage
130
+ cannot be read, the check is skipped.
131
+ """
132
+ cgroup_limit = get_cgroup_memory_limit()
133
+ cgroup_usage = get_cgroup_memory_usage()
134
+ if cgroup_limit is None or cgroup_usage is None:
135
+ return
136
+
137
+ cgroup_available = max(0, cgroup_limit - cgroup_usage)
138
+ mib = 1 << 20
139
+ remaining_bytes = cgroup_available - required_bytes
140
+ log_fn = logger.debug if remaining_bytes >= 0 else logger.warning
141
+ status = (
142
+ "current headroom meets the requested allocation"
143
+ if remaining_bytes >= 0
144
+ else "current headroom is below the requested allocation; allocation "
145
+ "will still be attempted because cgroup usage may be reclaimable"
146
+ )
147
+ log_fn(
148
+ "Cgroup memory preflight for %s: %.0f MiB required, %.0f MiB current "
149
+ "usage, %.0f MiB available under %.0f MiB limit, %.0f MiB remaining "
150
+ "after allocation based on current usage; %s.",
151
+ allocation_name,
152
+ required_bytes / mib,
153
+ cgroup_usage / mib,
154
+ cgroup_available / mib,
155
+ cgroup_limit / mib,
156
+ remaining_bytes / mib,
157
+ status,
158
+ )
159
160
161
def get_memory_affinity(pid: int = 0) -> list[int]:
@@ -161,7 +227,8 @@ def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:
227
# would be applied to host RAM instead of the pod's limit. cgroup
228
# does not expose per-NUMA-node limits, so we just clamp the totals
229
# against the pod-wide limit here.
164
- cgroup_limit, cgroup_usage = get_cgroup_memory_limit()
230
+ cgroup_limit = get_cgroup_memory_limit()
231
+ cgroup_usage = get_cgroup_memory_usage()
232
if cgroup_limit is not None and cgroup_limit < total_memory:
233
total_memory = cgroup_limit
234
cgroup_available = cgroup_limit - (cgroup_usage or 0)
vllm/v1/kv_offload/cpu/shared_offload_region.py
+4
-1
@@ -135,7 +135,10 @@ class SharedOffloadRegion:
135
self._creator = True
136
if creator_memory_check is not None:
137
creator_memory_check(self.total_size_bytes)
138
- check_shm_free_space(self.total_size_bytes)
138
+ check_shm_free_space(
139
+ self.total_size_bytes,
140
+ allocation_name="CPU KV offload shared region in /dev/shm",
141
+ )
142
os.ftruncate(self.fd, self.total_size_bytes)
143
logger.info(
144
"Created mmap file %s (%.2f GB)",