@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 189 lines 5.88 KB
Raw
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()