diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py index 8224db31bf..722c0d89fa 100644 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py @@ -2,6 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for ECSharedRegion (the mmap substrate).""" +import contextlib +import ctypes +import errno +import mmap import os import uuid from unittest.mock import MagicMock, patch @@ -30,6 +34,39 @@ def region() -> ECSharedRegion: r.cleanup() +@contextlib.contextmanager +def _region(**kwargs): + """Context manager: create one region, clean up on exit.""" + r = _make_region(**kwargs) + try: + yield r + finally: + r.cleanup() + + +def _page_residency(mmap_obj: mmap.mmap, length: int) -> list[bool]: + """Return Linux page-residency bits for a writable mmap.""" + try: + libc = ctypes.CDLL(None, use_errno=True) + mincore = libc.mincore + except AttributeError: + pytest.skip("mincore is unavailable") + + page_count = (length + mmap.PAGESIZE - 1) // mmap.PAGESIZE + mincore.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_ubyte), + ] + mincore.restype = ctypes.c_int + vector = (ctypes.c_ubyte * page_count)() + address = ctypes.addressof(ctypes.c_ubyte.from_buffer(mmap_obj)) + result = mincore(ctypes.c_void_p(address), length, vector) + if result != 0: + raise OSError(ctypes.get_errno()) + return [bool(value & 1) for value in vector] + + # ── mmap sharing between two instances ─────────────────────────────────────── @@ -109,6 +146,167 @@ def test_wait_for_file_size_times_out_when_file_stays_empty(tmp_path): os.close(fd) +# ── MADV_POPULATE_WRITE pre-faulting / fallback ────────────────────────────── + + +def test_madvise_success_selects_madvise_population(monkeypatch): + """A successful probe must keep using MADV_POPULATE_WRITE — the fallback + helper must never be invoked on a kernel that accepts the advice. + + Spies on both helpers so we can verify call counts: 1 probe call over a + single page + 1 populate call over the whole flat region, fallback 0 times. + """ + from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr + + madvise_calls: list[tuple[int, int, int]] = [] + fallback_calls: list[tuple[int, int]] = [] + + def _spy_madvise(mm, off, ln): + madvise_calls.append((off, ln, id(mm))) + + def _spy_fallback(mm, off, ln): + fallback_calls.append((off, ln)) + + monkeypatch.setattr(esr, "_madvise_populate_write", _spy_madvise) + monkeypatch.setattr(esr, "_fallback_populate_write", _spy_fallback) + + num_blocks, block_size_bytes = 4, mmap.PAGESIZE + with _region(num_blocks=num_blocks, block_size_bytes=block_size_bytes): + total = num_blocks * block_size_bytes + mmap_id = madvise_calls[0][2] + assert madvise_calls == [ + (0, mmap.PAGESIZE, mmap_id), # probe + (0, total, mmap_id), # flat pre-fault of the whole region + ] + assert fallback_calls == [], ( + "native-path constructor must not invoke the fallback helper" + ) + + +def test_madvise_einval_selects_fallback_for_whole_region(monkeypatch): + """An EINVAL probe must select the fallback for the whole flat region. + + This is the regression under test: on Linux < 5.14 the constant is absent + from the `mmap` module, so advice=23 reaches a kernel that rejects it and + every ECSharedRegion construction used to abort with EINVAL. + """ + from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr + + fallback_calls: list[tuple[int, int]] = [] + real_fallback = esr._fallback_populate_write + + def _raise_einval(mm, off, ln): + raise OSError(errno.EINVAL, "simulated unsupported kernel") + + def _spy_fallback(mm, off, ln): + fallback_calls.append((off, ln)) + return real_fallback(mm, off, ln) + + monkeypatch.setattr(esr, "_madvise_populate_write", _raise_einval) + monkeypatch.setattr(esr, "_fallback_populate_write", _spy_fallback) + + num_blocks, block_size_bytes = 4, mmap.PAGESIZE + with _region(num_blocks=num_blocks, block_size_bytes=block_size_bytes) as r: + assert fallback_calls == [(0, num_blocks * block_size_bytes)] + # Region is still fully usable after taking the fallback path. + r.blocks[0, :4] = torch.tensor([1, 2, 3, 4], dtype=torch.int8) + assert r.blocks[0, :4].tolist() == [1, 2, 3, 4] + + +def test_non_creator_does_not_probe_or_populate(monkeypatch): + """Only the creator pre-faults. A second opener must not probe the advice + nor re-touch pages the creator already populated.""" + from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr + + instance_id = str(uuid.uuid4()) + r1 = ECSharedRegion( + engine_id=instance_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE + ) + try: + calls: list[str] = [] + monkeypatch.setattr( + esr, + "_madvise_populate_write", + lambda mm, off, ln: calls.append("madvise"), + ) + monkeypatch.setattr( + esr, + "_fallback_populate_write", + lambda mm, off, ln: calls.append("fallback"), + ) + r2 = ECSharedRegion( + engine_id=instance_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE + ) + try: + assert not r2._is_creator + assert calls == [] + finally: + r2.cleanup() + finally: + r1.cleanup() + + +def test_fallback_populate_write_preserves_bytes_and_faults_pages(): + """Fallback preserves existing bytes and touches every target page. + + `|= 0` rather than `= 0` matters here: a peer worker may already have + written EC data into the shared mmap before we pre-fault it. + """ + from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr + + size = 3 * mmap.PAGESIZE + if not hasattr(mmap, "MADV_DONTNEED"): + pytest.skip("MADV_DONTNEED is unavailable") + + mmap_obj = mmap.mmap( + -1, + size, + flags=mmap.MAP_SHARED, + prot=mmap.PROT_READ | mmap.PROT_WRITE, + ) + try: + mmap_obj.madvise(mmap.MADV_DONTNEED, 0, size) + if any(_page_residency(mmap_obj, size)): + pytest.skip("kernel did not discard anonymous mmap pages") + + mmap_obj[0] = 0xAB + assert _page_residency(mmap_obj, size) == [True, False, False] + + esr._fallback_populate_write(mmap_obj, 0, size) + + assert _page_residency(mmap_obj, size) == [True, True, True] + assert mmap_obj[0] == 0xAB + finally: + mmap_obj.close() + + +def test_madvise_unexpected_oserror_propagates(monkeypatch): + """Only EINVAL triggers the fallback. Other OSErrors (e.g. EIO) must + propagate out of __init__, not be silently masked by the fallback branch. + """ + from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr + + monkeypatch.setattr( + esr, + "_madvise_populate_write", + lambda mm, off, ln: (_ for _ in ()).throw( + OSError(errno.EIO, "simulated I/O failure") + ), + ) + engine_id = str(uuid.uuid4()) + try: + with pytest.raises(OSError) as exc_info: + ECSharedRegion( + engine_id=engine_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE + ) + assert exc_info.value.errno == errno.EIO + finally: + # __init__ raised past the unlink-owning cleanup(), so drop the + # backing file here rather than leaking it into /dev/shm. + with contextlib.suppress(FileNotFoundError): + os.unlink(f"/dev/shm/vllm_ec_{engine_id}.mmap") + + # ── pin_memory ─────────────────────────────────────────────────────────────── diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py b/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py index 093292142e..413bf28571 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py @@ -7,16 +7,51 @@ Modeled after SharedOffloadRegion (vllm/v1/kv_offload/cpu/) but simplified for EC: flat shared layout, no multi-tensor cursor, no block_size_factor. """ +import errno import mmap import os import time +from collections.abc import Callable +import numpy as np import torch from vllm.logger import init_logger logger = init_logger(__name__) +# MADV_POPULATE_WRITE was added in Linux 5.14 (value 23). +_MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23) + + +def _madvise_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None: + mmap_obj.madvise(_MADV_POPULATE_WRITE, offset, length) + + +def _fallback_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None: + # Touch one byte per page via a read-modify-write so existing bytes are + # preserved — a peer worker may have already written EC data into this + # shared mmap by the time we run on a kernel without MADV_POPULATE_WRITE. + arr = np.frombuffer(mmap_obj, dtype=np.uint8) + arr[offset : offset + length : mmap.PAGESIZE] |= 0 + + +def _get_populate_write_fn( + mmap_obj: mmap.mmap, +) -> Callable[[mmap.mmap, int, int], None]: + """Select the pre-faulting method once for this mmap.""" + try: + _madvise_populate_write(mmap_obj, 0, mmap.PAGESIZE) + except OSError as e: + if e.errno != errno.EINVAL: + raise + logger.warning( + "MADV_POPULATE_WRITE is not supported; falling back to per-page " + "writes for EC mmap pre-population. Startup may be slower." + ) + return _fallback_populate_write + return _madvise_populate_write + def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0): """Spin-wait until the file reaches expected_size (creator truncated it).""" @@ -93,8 +128,8 @@ class ECSharedRegion: ) if self._is_creator: - _MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23) - self._mmap_obj.madvise(_MADV_POPULATE_WRITE, 0, total_size_bytes) + populate_write_fn = _get_populate_write_fn(self._mmap_obj) + populate_write_fn(self._mmap_obj, 0, total_size_bytes) # (num_blocks, block_size_bytes) int8 tensor over the mmap buffer. self.blocks: torch.Tensor = torch.frombuffer(