[Bugfix][EC Connector] Fall back when MADV_POPULATE_WRITE is unsupported (#53190)
Signed-off-by: wentian-byte <2990624738@qq.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
xiaolinchen committed
Sep 1, 2026 at 11:17 UTC
188716ace7475593dc22f8a72bf1d692f5c9b5ee
2 files changed
+235
-2
tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py
+198
@@ -2,6 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Unit tests for ECSharedRegion (the mmap substrate)."""
4
5
+import contextlib
6
+import ctypes
7
+import errno
8
+import mmap
9
import os
10
import uuid
11
from unittest.mock import MagicMock, patch
@@ -30,6 +34,39 @@ def region() -> ECSharedRegion:
34
r.cleanup()
35
36
37
+@contextlib.contextmanager
38
+def _region(**kwargs):
39
+ """Context manager: create one region, clean up on exit."""
40
+ r = _make_region(**kwargs)
41
+ try:
42
+ yield r
43
+ finally:
44
+ r.cleanup()
45
+
46
+
47
+def _page_residency(mmap_obj: mmap.mmap, length: int) -> list[bool]:
48
+ """Return Linux page-residency bits for a writable mmap."""
49
+ try:
50
+ libc = ctypes.CDLL(None, use_errno=True)
51
+ mincore = libc.mincore
52
+ except AttributeError:
53
+ pytest.skip("mincore is unavailable")
54
+
55
+ page_count = (length + mmap.PAGESIZE - 1) // mmap.PAGESIZE
56
+ mincore.argtypes = [
57
+ ctypes.c_void_p,
58
+ ctypes.c_size_t,
59
+ ctypes.POINTER(ctypes.c_ubyte),
60
+ ]
61
+ mincore.restype = ctypes.c_int
62
+ vector = (ctypes.c_ubyte * page_count)()
63
+ address = ctypes.addressof(ctypes.c_ubyte.from_buffer(mmap_obj))
64
+ result = mincore(ctypes.c_void_p(address), length, vector)
65
+ if result != 0:
66
+ raise OSError(ctypes.get_errno())
67
+ return [bool(value & 1) for value in vector]
68
+
69
+
70
# ── mmap sharing between two instances ───────────────────────────────────────
71
72
@@ -109,6 +146,167 @@ def test_wait_for_file_size_times_out_when_file_stays_empty(tmp_path):
146
os.close(fd)
147
148
149
+# ── MADV_POPULATE_WRITE pre-faulting / fallback ──────────────────────────────
150
+
151
+
152
+def test_madvise_success_selects_madvise_population(monkeypatch):
153
+ """A successful probe must keep using MADV_POPULATE_WRITE — the fallback
154
+ helper must never be invoked on a kernel that accepts the advice.
155
+
156
+ Spies on both helpers so we can verify call counts: 1 probe call over a
157
+ single page + 1 populate call over the whole flat region, fallback 0 times.
158
+ """
159
+ from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
160
+
161
+ madvise_calls: list[tuple[int, int, int]] = []
162
+ fallback_calls: list[tuple[int, int]] = []
163
+
164
+ def _spy_madvise(mm, off, ln):
165
+ madvise_calls.append((off, ln, id(mm)))
166
+
167
+ def _spy_fallback(mm, off, ln):
168
+ fallback_calls.append((off, ln))
169
+
170
+ monkeypatch.setattr(esr, "_madvise_populate_write", _spy_madvise)
171
+ monkeypatch.setattr(esr, "_fallback_populate_write", _spy_fallback)
172
+
173
+ num_blocks, block_size_bytes = 4, mmap.PAGESIZE
174
+ with _region(num_blocks=num_blocks, block_size_bytes=block_size_bytes):
175
+ total = num_blocks * block_size_bytes
176
+ mmap_id = madvise_calls[0][2]
177
+ assert madvise_calls == [
178
+ (0, mmap.PAGESIZE, mmap_id), # probe
179
+ (0, total, mmap_id), # flat pre-fault of the whole region
180
+ ]
181
+ assert fallback_calls == [], (
182
+ "native-path constructor must not invoke the fallback helper"
183
+ )
184
+
185
+
186
+def test_madvise_einval_selects_fallback_for_whole_region(monkeypatch):
187
+ """An EINVAL probe must select the fallback for the whole flat region.
188
+
189
+ This is the regression under test: on Linux < 5.14 the constant is absent
190
+ from the `mmap` module, so advice=23 reaches a kernel that rejects it and
191
+ every ECSharedRegion construction used to abort with EINVAL.
192
+ """
193
+ from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
194
+
195
+ fallback_calls: list[tuple[int, int]] = []
196
+ real_fallback = esr._fallback_populate_write
197
+
198
+ def _raise_einval(mm, off, ln):
199
+ raise OSError(errno.EINVAL, "simulated unsupported kernel")
200
+
201
+ def _spy_fallback(mm, off, ln):
202
+ fallback_calls.append((off, ln))
203
+ return real_fallback(mm, off, ln)
204
+
205
+ monkeypatch.setattr(esr, "_madvise_populate_write", _raise_einval)
206
+ monkeypatch.setattr(esr, "_fallback_populate_write", _spy_fallback)
207
+
208
+ num_blocks, block_size_bytes = 4, mmap.PAGESIZE
209
+ with _region(num_blocks=num_blocks, block_size_bytes=block_size_bytes) as r:
210
+ assert fallback_calls == [(0, num_blocks * block_size_bytes)]
211
+ # Region is still fully usable after taking the fallback path.
212
+ r.blocks[0, :4] = torch.tensor([1, 2, 3, 4], dtype=torch.int8)
213
+ assert r.blocks[0, :4].tolist() == [1, 2, 3, 4]
214
+
215
+
216
+def test_non_creator_does_not_probe_or_populate(monkeypatch):
217
+ """Only the creator pre-faults. A second opener must not probe the advice
218
+ nor re-touch pages the creator already populated."""
219
+ from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
220
+
221
+ instance_id = str(uuid.uuid4())
222
+ r1 = ECSharedRegion(
223
+ engine_id=instance_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE
224
+ )
225
+ try:
226
+ calls: list[str] = []
227
+ monkeypatch.setattr(
228
+ esr,
229
+ "_madvise_populate_write",
230
+ lambda mm, off, ln: calls.append("madvise"),
231
+ )
232
+ monkeypatch.setattr(
233
+ esr,
234
+ "_fallback_populate_write",
235
+ lambda mm, off, ln: calls.append("fallback"),
236
+ )
237
+ r2 = ECSharedRegion(
238
+ engine_id=instance_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE
239
+ )
240
+ try:
241
+ assert not r2._is_creator
242
+ assert calls == []
243
+ finally:
244
+ r2.cleanup()
245
+ finally:
246
+ r1.cleanup()
247
+
248
+
249
+def test_fallback_populate_write_preserves_bytes_and_faults_pages():
250
+ """Fallback preserves existing bytes and touches every target page.
251
+
252
+ `|= 0` rather than `= 0` matters here: a peer worker may already have
253
+ written EC data into the shared mmap before we pre-fault it.
254
+ """
255
+ from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
256
+
257
+ size = 3 * mmap.PAGESIZE
258
+ if not hasattr(mmap, "MADV_DONTNEED"):
259
+ pytest.skip("MADV_DONTNEED is unavailable")
260
+
261
+ mmap_obj = mmap.mmap(
262
+ -1,
263
+ size,
264
+ flags=mmap.MAP_SHARED,
265
+ prot=mmap.PROT_READ | mmap.PROT_WRITE,
266
+ )
267
+ try:
268
+ mmap_obj.madvise(mmap.MADV_DONTNEED, 0, size)
269
+ if any(_page_residency(mmap_obj, size)):
270
+ pytest.skip("kernel did not discard anonymous mmap pages")
271
+
272
+ mmap_obj[0] = 0xAB
273
+ assert _page_residency(mmap_obj, size) == [True, False, False]
274
+
275
+ esr._fallback_populate_write(mmap_obj, 0, size)
276
+
277
+ assert _page_residency(mmap_obj, size) == [True, True, True]
278
+ assert mmap_obj[0] == 0xAB
279
+ finally:
280
+ mmap_obj.close()
281
+
282
+
283
+def test_madvise_unexpected_oserror_propagates(monkeypatch):
284
+ """Only EINVAL triggers the fallback. Other OSErrors (e.g. EIO) must
285
+ propagate out of __init__, not be silently masked by the fallback branch.
286
+ """
287
+ from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
288
+
289
+ monkeypatch.setattr(
290
+ esr,
291
+ "_madvise_populate_write",
292
+ lambda mm, off, ln: (_ for _ in ()).throw(
293
+ OSError(errno.EIO, "simulated I/O failure")
294
+ ),
295
+ )
296
+ engine_id = str(uuid.uuid4())
297
+ try:
298
+ with pytest.raises(OSError) as exc_info:
299
+ ECSharedRegion(
300
+ engine_id=engine_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE
301
+ )
302
+ assert exc_info.value.errno == errno.EIO
303
+ finally:
304
+ # __init__ raised past the unlink-owning cleanup(), so drop the
305
+ # backing file here rather than leaking it into /dev/shm.
306
+ with contextlib.suppress(FileNotFoundError):
307
+ os.unlink(f"/dev/shm/vllm_ec_{engine_id}.mmap")
308
+
309
+
310
# ── pin_memory ───────────────────────────────────────────────────────────────
311
312
vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py
+37
-2
@@ -7,16 +7,51 @@ Modeled after SharedOffloadRegion (vllm/v1/kv_offload/cpu/) but simplified
7
for EC: flat shared layout, no multi-tensor cursor, no block_size_factor.
8
"""
9
10
+import errno
11
import mmap
12
import os
13
import time
14
+from collections.abc import Callable
15
16
+import numpy as np
17
import torch
18
19
from vllm.logger import init_logger
20
21
logger = init_logger(__name__)
22
23
+# MADV_POPULATE_WRITE was added in Linux 5.14 (value 23).
24
+_MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23)
25
+
26
+
27
+def _madvise_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None:
28
+ mmap_obj.madvise(_MADV_POPULATE_WRITE, offset, length)
29
+
30
+
31
+def _fallback_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None:
32
+ # Touch one byte per page via a read-modify-write so existing bytes are
33
+ # preserved — a peer worker may have already written EC data into this
34
+ # shared mmap by the time we run on a kernel without MADV_POPULATE_WRITE.
35
+ arr = np.frombuffer(mmap_obj, dtype=np.uint8)
36
+ arr[offset : offset + length : mmap.PAGESIZE] |= 0
37
+
38
+
39
+def _get_populate_write_fn(
40
+ mmap_obj: mmap.mmap,
41
+) -> Callable[[mmap.mmap, int, int], None]:
42
+ """Select the pre-faulting method once for this mmap."""
43
+ try:
44
+ _madvise_populate_write(mmap_obj, 0, mmap.PAGESIZE)
45
+ except OSError as e:
46
+ if e.errno != errno.EINVAL:
47
+ raise
48
+ logger.warning(
49
+ "MADV_POPULATE_WRITE is not supported; falling back to per-page "
50
+ "writes for EC mmap pre-population. Startup may be slower."
51
+ )
52
+ return _fallback_populate_write
53
+ return _madvise_populate_write
54
+
55
56
def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0):
57
"""Spin-wait until the file reaches expected_size (creator truncated it)."""
@@ -93,8 +128,8 @@ class ECSharedRegion:
128
)
129
130
if self._is_creator:
96
- _MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23)
97
- self._mmap_obj.madvise(_MADV_POPULATE_WRITE, 0, total_size_bytes)
131
+ populate_write_fn = _get_populate_write_fn(self._mmap_obj)
132
+ populate_write_fn(self._mmap_obj, 0, total_size_bytes)
133
134
# (num_blocks, block_size_bytes) int8 tensor over the mmap buffer.
135
self.blocks: torch.Tensor = torch.frombuffer(