diff --git a/benchmarks/bench_load_weight_generic.py b/benchmarks/bench_load_weight_generic.py index 89ddb30708..62807bc8fe 100644 --- a/benchmarks/bench_load_weight_generic.py +++ b/benchmarks/bench_load_weight_generic.py @@ -41,6 +41,8 @@ import torch +from freetoken.memory import effective_memory_available + GiB = float(1 << 30) ALL_MODES = ("parallel", "ftw") @@ -54,14 +56,6 @@ def _default_ftw_dir(model_path: str) -> str: # ---------------- memory sampling + checksum ---------------- -def _meminfo_available_bytes() -> int: - with open("/proc/meminfo") as f: - for line in f: - if line.startswith("MemAvailable:"): - return int(line.split()[1]) * 1024 - return 0 - - def _status_kb(key: str) -> int: with open("/proc/self/status") as f: for line in f: @@ -71,7 +65,7 @@ def _status_kb(key: str) -> int: class MemSampler(threading.Thread): - """Background sampler: peak process RSS + system MemAvailable low-water-mark.""" + """Background sampler: peak RSS + effective available-memory low-water-mark.""" def __init__(self, interval: float = 0.2): super().__init__(daemon=True) @@ -82,7 +76,9 @@ def __init__(self, interval: float = 0.2): def run(self) -> None: while not self._stop_evt.wait(self.interval): - self.min_avail = min(self.min_avail, _meminfo_available_bytes()) + available = effective_memory_available() + if available is not None: + self.min_avail = min(self.min_avail, available) self.max_rss = max(self.max_rss, _status_kb("VmRSS:")) def stop(self) -> None: diff --git a/python/freetoken/memory.py b/python/freetoken/memory.py new file mode 100644 index 0000000000..ef4294d1c2 --- /dev/null +++ b/python/freetoken/memory.py @@ -0,0 +1,334 @@ +"""Conservative host-memory admission for bare-metal and cgrouped processes. + +The kernel's host-wide ``MemAvailable`` figure is not an allocation budget for a +container. This module combines it with the process's tightest finite cgroup +budget and deliberately keeps the probing code independent of torch and the +rest of the runtime so startup checks and diagnostic tools can share it. +""" + +from __future__ import annotations + +import re +from pathlib import Path + + +# Kernel pseudo-files are small, but bound every read and hierarchy walk. A +# container may expose a synthetic or partially mounted cgroup tree; admission +# must not turn an unexpected file into an unbounded read or traversal. +_CGROUP_VALUE_MAX_BYTES = 128 +_CGROUP_STAT_MAX_BYTES = 64 * 1024 +_CGROUP_MEMBERSHIP_MAX_BYTES = 64 * 1024 +_CGROUP_MAX_ANCESTORS = 64 +_MEMINFO_MAX_BYTES = 256 * 1024 +_UINT64_MAX = (1 << 64) - 1 + +# Cgroup v1 represents "unlimited" with a page-aligned value just below the +# signed 64-bit maximum (commonly 9223372036854771712). No supported host can +# offer an exbibyte of usable RAM, so this threshold safely covers the kernel's +# sentinel variants without mistaking a practical finite limit for unlimited. +_CGROUP_V1_UNLIMITED_MIN = 1 << 60 + + +def _read_bounded_ascii(path: Path, max_bytes: int) -> tuple[bool, str | None]: + """Return ``(present, exact_text)`` for a small kernel pseudo-file. + + A missing path means that source/controller is unavailable. Other I/O + errors are authoritative probe failures and raise so the public resolver can + fail closed instead of silently treating the process as unconstrained. + ``None`` text means the present file was oversized or not ASCII. + """ + path = Path(path) + try: + with path.open("rb") as stream: + raw = stream.read(max_bytes + 1) + except (FileNotFoundError, NotADirectoryError): + return False, None + except OSError as exc: + raise ValueError(f"cannot read memory control file {path}: {exc}") from exc + if len(raw) > max_bytes: + return True, None + try: + return True, raw.decode("ascii") + except UnicodeDecodeError: + return True, None + + +def _parse_u64_counter(text: str | None) -> int | None: + """Parse one non-negative decimal counter with at most one final newline.""" + match = re.fullmatch(r"(\d+)\n?", text or "") + if match is None: + return None + value = int(match.group(1)) + return value if value <= _UINT64_MAX else None + + +def _host_memory_available(meminfo_path: Path) -> int | None: + """Read host-wide ``MemAvailable`` bytes, or ``None`` when /proc is absent.""" + present, text = _read_bounded_ascii(meminfo_path, _MEMINFO_MAX_BYTES) + if not present: + return None + if text is None: + raise ValueError(f"malformed or oversized memory information: {meminfo_path}") + match = re.search( + r"^MemAvailable:[ \t]+(\d+)[ \t]+kB[ \t]*(?:\n|$)", + text, + flags=re.MULTILINE, + ) + if match is None: + raise ValueError(f"MemAvailable is unavailable or malformed: {meminfo_path}") + kib = int(match.group(1)) + if kib > _UINT64_MAX // 1024: + raise ValueError(f"MemAvailable overflows a byte counter: {meminfo_path}") + return kib * 1024 + + +def _cgroup_file_cache(directory: Path, *, v1: bool = False) -> int: + """Return droppable file-LRU bytes, or zero if the statistics cannot be trusted. + + Dirty and writeback pages sit on the file lists too, but they cannot be + dropped until they are written, and a cgroup at its limit is OOM-killed + before that finishes, so they are taken back out of the credit. + """ + try: + _, text = _read_bounded_ascii( + Path(directory) / "memory.stat", _CGROUP_STAT_MAX_BYTES + ) + if text is None: + return 0 + counters: dict[str, int] = {} + for line in text.removesuffix("\n").split("\n"): + key, separator, counter = line.partition(" ") + if not separator or re.fullmatch(r"[A-Za-z0-9_]+", key) is None: + return 0 + value = _parse_u64_counter(counter) + if value is None or key in counters: + return 0 + counters[key] = value + except ValueError: + return 0 + + # v1 totals match hierarchical usage. Unlike `file`, file LRU lists exclude shmem. + prefix = "total_" if v1 else "" + cache = counters.get(prefix + "inactive_file", 0) + counters.get( + prefix + "active_file", 0 + ) + dirty, writeback = ( + ("total_dirty", "total_writeback") if v1 else ("file_dirty", "file_writeback") + ) + return max(0, cache - counters.get(dirty, 0) - counters.get(writeback, 0)) + + +def _cgroup_pair_remaining( + directory: Path, + limit_name: str, + current_name: str, + *, + v1: bool = False, +) -> tuple[bool, int | None]: + """Return ``(limit_present, finite_remaining_or_none)`` for one cgroup.""" + limit_path = Path(directory) / limit_name + current_path = Path(directory) / current_name + limit_present, limit_text = _read_bounded_ascii(limit_path, _CGROUP_VALUE_MAX_BYTES) + if not limit_present: + return False, None + if not v1 and limit_text in ("max", "max\n"): + return True, None + + limit = _parse_u64_counter(limit_text) + if limit is None: + raise ValueError(f"malformed cgroup memory limit: {limit_path}") + if v1 and limit >= _CGROUP_V1_UNLIMITED_MIN: + return True, None + + current_present, current_text = _read_bounded_ascii( + current_path, _CGROUP_VALUE_MAX_BYTES + ) + if not current_present: + raise ValueError(f"cgroup memory usage is unavailable: {current_path}") + current = _parse_u64_counter(current_text) + if current is None: + raise ValueError(f"malformed cgroup memory usage: {current_path}") + working_set = max(0, current - _cgroup_file_cache(directory, v1=v1)) + return True, max(0, limit - working_set) + + +def _read_cgroup_memberships(proc_cgroup_path: Path) -> tuple[str | None, str | None]: + """Return the v2 and v1-memory paths from bounded ``/proc/self/cgroup``. + + Hybrid hosts list leftover v1 hierarchies (``1:net_cls:/``) above the unified + ``0::/...`` line; non-memory v1 controllers are simply skipped. + """ + present, text = _read_bounded_ascii(proc_cgroup_path, _CGROUP_MEMBERSHIP_MAX_BYTES) + if not present: + return None, None + if text is None: + raise ValueError( + f"malformed or oversized cgroup membership: {proc_cgroup_path}" + ) + + v2_path = None + v1_path = None + for line in text.splitlines(): + fields = line.split(":", 2) + if len(fields) != 3: + continue + hierarchy, controllers, member_path = fields + if not member_path.startswith("/"): + continue + # Real cgroup membership paths are absolute and contain no dot + # components. Reject lexical escapes in synthetic/malformed proc data. + parts = Path(member_path).parts[1:] + if any(part in ("", ".", "..") for part in parts): + continue + if hierarchy == "0" and not controllers: + v2_path = member_path + elif "memory" in controllers.split(","): + v1_path = member_path + return v2_path, v1_path + + +def _control_file_exists(path: Path) -> bool: + """Existence probe that fails closed on errors other than a missing path.""" + path = Path(path) + try: + path.stat() + return True + except (FileNotFoundError, NotADirectoryError): + return False + except OSError as exc: + raise ValueError(f"cannot inspect cgroup control file {path}: {exc}") from exc + + +def _cgroup_control_directory( + root: Path, member_path: str | None, limit_name: str +) -> Path | None: + """Map a proc membership path into a mounted, possibly namespaced tree.""" + root = Path(root) + if member_path: + parts = Path(member_path).parts[1:] + if len(parts) >= _CGROUP_MAX_ANCESTORS: + raise ValueError( + f"cgroup membership nesting exceeds {_CGROUP_MAX_ANCESTORS - 1} levels" + ) + candidate = root.joinpath(*parts) + # A leaf can omit controller files when only an ancestor has the memory + # controller enabled. Find the nearest visible ancestor without a + # filesystem scan or a walk outside the mounted hierarchy. + for _ in range(_CGROUP_MAX_ANCESTORS): + if _control_file_exists(candidate / limit_name): + return candidate + if candidate == root: + break + parent = candidate.parent + if parent == candidate or (parent != root and root not in parent.parents): + break + candidate = parent + + # Container runtimes may mount the process's own subgroup as the hierarchy + # root while /proc still reports a host-side membership path. + if _control_file_exists(root / limit_name): + return root + return None + + +def _cgroup_hierarchy_remaining( + root: Path, + member_path: str | None, + limit_name: str, + current_name: str, + *, + v1: bool = False, +) -> tuple[bool, int | None]: + """Return ``(hierarchy_seen, tightest finite ancestor headroom)``.""" + root = Path(root) + current = _cgroup_control_directory(root, member_path, limit_name) + if current is None: + return False, None + + seen = False + remaining = None + for _ in range(_CGROUP_MAX_ANCESTORS): + present, candidate = _cgroup_pair_remaining( + current, limit_name, current_name, v1=v1 + ) + seen = seen or present + if candidate is not None: + remaining = candidate if remaining is None else min(remaining, candidate) + if current == root: + break + parent = current.parent + if parent == current or (parent != root and root not in parent.parents): + break + current = parent + return seen, remaining + + +def _cgroup_memory_remaining( + cgroup_root: Path = Path("/sys/fs/cgroup"), + proc_cgroup_path: Path = Path("/proc/self/cgroup"), +) -> int | None: + """Return this process's tightest finite cgroup memory headroom. + + Cgroup v2 is authoritative whenever its memory controller is visible. + ``max`` is a known-unlimited value. Only when v2 is absent do we try the + conventional, bounded v1 controller roots; no filesystem scan is used. + Malformed, overflowing, or unreadable limits or usage raise ``ValueError``. + Optional file-cache statistics that cannot be parsed subtract nothing. + """ + root = Path(cgroup_root) + v2_path, v1_path = _read_cgroup_memberships(proc_cgroup_path) + v2_seen, remaining = _cgroup_hierarchy_remaining( + root, v2_path, "memory.max", "memory.current" + ) + if v2_seen: + return remaining + + for v1_root in (root / "memory", root): + v1_seen, remaining = _cgroup_hierarchy_remaining( + v1_root, + v1_path, + "memory.limit_in_bytes", + "memory.usage_in_bytes", + v1=True, + ) + if v1_seen: + return remaining + return None + + +def effective_memory_available( + *, + meminfo_path: Path = Path("/proc/meminfo"), + cgroup_root: Path = Path("/sys/fs/cgroup"), + proc_cgroup_path: Path = Path("/proc/self/cgroup"), +) -> int | None: + """Return the effective available host-memory budget in bytes. + + The result is ``min(MemAvailable, tightest finite cgroup headroom)`` when + both measurements exist, either known bound when only one exists, and + ``None`` only when neither source is available -- no procfs, a non-Linux + host, or a cgroup tree without a visible memory controller. ``None`` means + *unknown* and is never conflated with ``0``; callers keep their historical + best-effort behaviour on it. A known zero remains zero. + + A *present* but malformed, overflowing, or permission-denied control file + is an authoritative signal that the budget cannot be trusted, so it fails + closed to ``0`` rather than being mistaken for an unlimited controller. + Optional ``memory.stat`` failures only disable the file-cache subtraction. + """ + try: + host_available = _host_memory_available(Path(meminfo_path)) + cgroup_remaining = _cgroup_memory_remaining( + Path(cgroup_root), Path(proc_cgroup_path) + ) + except ValueError: + return 0 + + if host_available is None: + return cgroup_remaining + if cgroup_remaining is None: + return host_available + return min(host_available, cgroup_remaining) + + +__all__ = ["effective_memory_available"] diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index 911ca664b3..79807d8a1d 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -59,6 +59,7 @@ single_gpu_arg, ) from freetoken.kernel.pinned import alloc_pinned_tensor +from freetoken.memory import effective_memory_available from freetoken.moe.bench_profile import PROFILE_VERSION from freetoken.moe.cpu_executor import physical_core_cpus, resolve_threads_and_affinity from freetoken.utils import init_logger @@ -164,29 +165,10 @@ def default_out_path(gpu_uuid: str | None = None) -> str: return default_profile_path(gpu_uuid) -def _cgroup_mem_headroom() -> int | None: - """Free bytes under this process's cgroup v2 memory limit, or None if unlimited.""" - try: - with open("/sys/fs/cgroup/memory.max") as f: - raw = f.read().strip() - if raw == "max": - return None - limit = int(raw) - with open("/sys/fs/cgroup/memory.current") as f: - used = int(f.read().strip()) - return max(0, limit - used) - except (OSError, ValueError): - return None - - def _available_ram_bytes() -> int: - """Free host RAM, clamped to the cgroup memory limit so a container isn't over-estimated.""" - try: - host = os.sysconf("SC_AVPHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") - except (ValueError, OSError, AttributeError): - host = 8 << 30 - cg = _cgroup_mem_headroom() - return min(host, cg) if cg is not None else host + """Effective free RAM, retaining the legacy fallback only when probing is absent.""" + available = effective_memory_available() + return (8 << 30) if available is None else available # ============================== ceilings (hardware) ============================== diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 197dbb9b04..c2c5eebad2 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -15,6 +15,7 @@ import torch from freetoken.layers.quantization import QuantKind +from freetoken.memory import effective_memory_available from freetoken.utils import init_logger from .host_banks import HostBank @@ -255,16 +256,10 @@ def _host_ram_fits_parallel(model_path: str) -> bool: """Best-effort: can free host RAM hold the expert banks plus the parallel reader's one extra (non-reclaimable) whole-shard buffer? Unknown (non-local path / no /proc) -> True, i.e. keep the fast path. Banks ~= checkpoint size (experts dominate); transient ~= the - largest shard. Uses MemAvailable (counts reclaimable cache) -- the OOM-relevant figure.""" - avail = None - try: - with open("/proc/meminfo") as f: - for line in f: - if line.startswith("MemAvailable:"): - avail = int(line.split()[1]) * 1024 - break - except OSError: - pass + largest shard. Uses the effective available memory (host ``MemAvailable``, which counts + reclaimable cache, clamped by the process's tightest finite cgroup budget) -- the + OOM-relevant figure on bare metal and inside a ``--memory``-limited container alike.""" + avail = effective_memory_available() if avail is None: return True try: # the loader's index-first shard list also covers subdirectory layouts (a root glob does not) diff --git a/tests/moe/test_expert_bank_memory_admission.py b/tests/moe/test_expert_bank_memory_admission.py new file mode 100644 index 0000000000..3d368cce73 --- /dev/null +++ b/tests/moe/test_expert_bank_memory_admission.py @@ -0,0 +1,123 @@ +import inspect +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.layers.quantization import QuantKind +from freetoken.memory import effective_memory_available + + +def _record_read_mode(monkeypatch: pytest.MonkeyPatch, expert_banks) -> list: + """Stand in for the raw-checkpoint bank builder and record the read mode it is handed. + + The mode is looked up by parameter name on the real builder's signature, so a + reordered signature cannot make this read a different argument. + """ + signature = inspect.signature(expert_banks._method_expert_banks) + chosen = [] + + def build(*args, **kwargs): + chosen.append(signature.bind(*args, **kwargs).arguments["parallel"]) + return expert_banks.ExpertBanks("bf16", {}) + + monkeypatch.setattr(expert_banks, "_method_expert_banks", build) + return chosen + + +@pytest.mark.parametrize( + ("cgroup_current", "parallel"), + [ + pytest.param(701, False, id="299-bytes-left"), + pytest.param(699, True, id="301-bytes-left"), + ], +) +def test_loader_auto_admission_consumes_effective_memory( + tmp_path, monkeypatch: pytest.MonkeyPatch, cgroup_current: int, parallel: bool +): + """Regression: production auto-load must use the shared effective budget.""" + import freetoken.checkpoint.ftw as ftw + import freetoken.models.weight as model_weight + import freetoken.moe.expert_banks as expert_banks + + # The bank estimate is 200 bytes plus a 100-byte parallel-reader transient. + # Host MemAvailable is large; the fake cgroup has 299 bytes left (one short, + # so the serial build) or 301 (enough, so the parallel reader is kept). + (tmp_path / "a.safetensors").write_bytes(b"a" * 100) + (tmp_path / "b.safetensors").write_bytes(b"b" * 100) + meminfo = tmp_path / "meminfo" + meminfo.write_text("MemAvailable: 1024 kB\n", encoding="ascii") + proc_cgroup = tmp_path / "proc-self-cgroup" + proc_cgroup.write_text("0::/\n", encoding="ascii") + cgroup_root = tmp_path / "cgroup" + cgroup_root.mkdir() + (cgroup_root / "memory.max").write_text("1000\n", encoding="ascii") + (cgroup_root / "memory.current").write_text(f"{cgroup_current}\n", encoding="ascii") + + def effective_fixture(): + return effective_memory_available( + meminfo_path=meminfo, + cgroup_root=cgroup_root, + proc_cgroup_path=proc_cgroup, + ) + + monkeypatch.setattr(expert_banks, "_PARALLEL_READER_SUPPORTED", True) + monkeypatch.setattr(expert_banks, "effective_memory_available", effective_fixture) + monkeypatch.setattr(model_weight, "experts_scattered", lambda _path: True) + monkeypatch.setattr(ftw, "is_ftw_checkpoint", lambda _path: False) + chosen = _record_read_mode(monkeypatch, expert_banks) + config = SimpleNamespace(num_moe_layers=1, expert_quant="none") + + expert_banks.load_expert_banks( + str(tmp_path), + config, + method=SimpleNamespace(kind=QuantKind.NONE), + device=torch.device("cpu"), + ) + + assert chosen == [parallel] + + +def test_explicit_parallel_loader_override_bypasses_auto_admission( + tmp_path, monkeypatch: pytest.MonkeyPatch +): + import freetoken.checkpoint.ftw as ftw + import freetoken.moe.expert_banks as expert_banks + + monkeypatch.setattr(expert_banks, "_PARALLEL_READER_SUPPORTED", True) + monkeypatch.setattr(ftw, "is_ftw_checkpoint", lambda _path: False) + monkeypatch.setattr( + expert_banks, + "effective_memory_available", + lambda: (_ for _ in ()).throw(AssertionError("auto admission should not run")), + ) + chosen = _record_read_mode(monkeypatch, expert_banks) + config = SimpleNamespace(num_moe_layers=1, expert_quant="none") + + expert_banks.load_expert_banks( + str(tmp_path), + config, + method=SimpleNamespace(kind=QuantKind.NONE), + device=torch.device("cpu"), + parallel=True, + ) + + assert chosen == [True] + + +def test_benchmark_ram_estimate_uses_effective_headroom(monkeypatch): + import freetoken.moe.benchbw as benchbw + + # A finite cgroup budget (here 3 GiB) must reach the benchmark unchanged + # instead of the host-wide figure the old sysconf probe reported. + monkeypatch.setattr(benchbw, "effective_memory_available", lambda: 3 << 30) + + assert benchbw._available_ram_bytes() == 3 << 30 + + +def test_benchmark_legacy_fallback_is_only_for_unknown_probe(monkeypatch): + import freetoken.moe.benchbw as benchbw + + monkeypatch.setattr(benchbw, "effective_memory_available", lambda: None) + + assert benchbw._available_ram_bytes() == 8 << 30 diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 0572678721..6198769ba3 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -1229,11 +1229,9 @@ def boom(addr, nbytes): def test_host_ram_fits_parallel_sizes_subdirectory_shards(tmp_path, monkeypatch): """The RAM probe must size the shards the reader will actually open, indexed subdirectory layouts included.""" - import builtins - import io import json - from freetoken.moe.expert_banks import _host_ram_fits_parallel + import freetoken.moe.expert_banks as expert_banks sub = tmp_path / "tensors" sub.mkdir() @@ -1249,12 +1247,6 @@ def test_host_ram_fits_parallel_sizes_subdirectory_shards(tmp_path, monkeypatch) } ) ) - real_open = builtins.open - - def fake_open(file, *args, **kwargs): - if str(file) == "/proc/meminfo": - return io.StringIO("MemAvailable: 512 kB\n") - return real_open(file, *args, **kwargs) - - monkeypatch.setattr(builtins, "open", fake_open) - assert not _host_ram_fits_parallel(str(tmp_path)) + # 512 kB available against 2 MiB of shards; the probe reads memory through the shared resolver + monkeypatch.setattr(expert_banks, "effective_memory_available", lambda: 512 * 1024) + assert not expert_banks._host_ram_fits_parallel(str(tmp_path)) diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 0000000000..219a9058cd --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,609 @@ +from pathlib import Path + +import pytest + +from freetoken.memory import ( + _CGROUP_MAX_ANCESTORS, + _CGROUP_V1_UNLIMITED_MIN, + _UINT64_MAX, + _cgroup_memory_remaining, + effective_memory_available, +) + + +class FakeMemoryFiles: + """Pure fake /proc and cgroup trees; never touches the runner's hierarchy.""" + + def __init__(self, base: Path): + self.root = base / "cgroup" + self.root.mkdir() + self.proc_cgroup = base / "proc-self-cgroup" + self.meminfo = base / "meminfo" + self.meminfo.write_text( + "MemTotal: 32768 kB\nMemAvailable: 16384 kB\n", encoding="ascii" + ) + + def membership(self, *lines: str) -> None: + self.proc_cgroup.write_text("\n".join(lines) + "\n", encoding="ascii") + + @staticmethod + def pair( + directory: Path, + limit_name: str, + current_name: str, + limit: object, + current: object, + ) -> None: + directory.mkdir(parents=True, exist_ok=True) + (directory / limit_name).write_text(f"{limit}\n", encoding="ascii") + (directory / current_name).write_text(f"{current}\n", encoding="ascii") + + def v2(self, directory: Path, limit: object, current: object) -> None: + self.pair(directory, "memory.max", "memory.current", limit, current) + + def v1(self, directory: Path, limit: object, current: object) -> None: + self.pair( + directory, + "memory.limit_in_bytes", + "memory.usage_in_bytes", + limit, + current, + ) + + def remaining(self) -> int | None: + return _cgroup_memory_remaining(self.root, self.proc_cgroup) + + def effective(self) -> int | None: + return effective_memory_available( + meminfo_path=self.meminfo, + cgroup_root=self.root, + proc_cgroup_path=self.proc_cgroup, + ) + + +@pytest.fixture +def memory_files(tmp_path: Path) -> FakeMemoryFiles: + return FakeMemoryFiles(tmp_path) + + +@pytest.fixture(params=[False, True], ids=["v2", "v1"]) +def file_cache_cgroup( + memory_files: FakeMemoryFiles, request: pytest.FixtureRequest +) -> tuple[Path, str]: + v1 = request.param + memory_files.membership("7:memory:/" if v1 else "0::/") + directory = memory_files.root / "memory" if v1 else memory_files.root + write = memory_files.v1 if v1 else memory_files.v2 + write(directory, 24 << 30, 20 << 30) + memory_files.meminfo.write_text("MemAvailable: 33554432 kB\n", encoding="ascii") + return directory, "total_" if v1 else "" + + +@pytest.mark.parametrize( + "current,inactive,active,remaining", + [ + (20, 8, 6, 18), + (20, 14, 0, 18), + (20, 0, 14, 18), + (20, 10, 11, 24), + (26, 1, 0, 0), + (26, 4, 0, 2), + (24, 0, 0, 0), + ], +) +def test_file_cache_counts_as_room( + memory_files: FakeMemoryFiles, + file_cache_cgroup, + current: int, + inactive: int, + active: int, + remaining: int, +): + directory, prefix = file_cache_cgroup + current_name = "memory.usage_in_bytes" if prefix else "memory.current" + (directory / current_name).write_text(f"{current << 30}\n", encoding="ascii") + (directory / "memory.stat").write_text( + f"{prefix}inactive_file {inactive << 30}\n" + f"{prefix}active_file {active << 30}\n", + encoding="ascii", + ) + + assert memory_files.remaining() == remaining << 30 + assert memory_files.effective() == remaining << 30 + + +@pytest.mark.parametrize("key", ["inactive_file", "active_file"]) +def test_file_cache_missing_one_key_uses_the_other( + memory_files: FakeMemoryFiles, file_cache_cgroup, key: str +): + directory, prefix = file_cache_cgroup + (directory / "memory.stat").write_text( + f"{prefix}{key} {14 << 30}\n", encoding="ascii" + ) + + assert memory_files.remaining() == 18 << 30 + + +def test_file_cache_uses_only_matching_file_lists( + memory_files: FakeMemoryFiles, file_cache_cgroup +): + directory, prefix = file_cache_cgroup + other_prefix = "" if prefix else "total_" + (directory / "memory.stat").write_text( + f"file {20 << 30}\nshmem {10 << 30}\ncache {20 << 30}\n" + f"{other_prefix}inactive_file {8 << 30}\n" + f"{other_prefix}active_file {6 << 30}\n", + encoding="ascii", + ) + + assert memory_files.remaining() == 4 << 30 + + +@pytest.mark.parametrize( + "dirty,writeback,remaining", + [(2, 1, 15), (2, 0, 16), (0, 1, 17), (16, 0, 4), (8, 8, 4)], +) +def test_dirty_and_writeback_pages_are_not_room( + memory_files: FakeMemoryFiles, + file_cache_cgroup, + dirty: int, + writeback: int, + remaining: int, +): + directory, prefix = file_cache_cgroup + dirty_key, writeback_key = ( + ("total_dirty", "total_writeback") + if prefix + else ("file_dirty", "file_writeback") + ) + (directory / "memory.stat").write_text( + f"{prefix}inactive_file {8 << 30}\n{prefix}active_file {6 << 30}\n" + f"{dirty_key} {dirty << 30}\n{writeback_key} {writeback << 30}\n", + encoding="ascii", + ) + + assert memory_files.remaining() == remaining << 30 + assert memory_files.effective() == remaining << 30 + + +def test_other_hierarchy_dirty_counters_are_ignored( + memory_files: FakeMemoryFiles, file_cache_cgroup +): + directory, prefix = file_cache_cgroup + other_dirty, other_writeback = ( + ("file_dirty", "file_writeback") + if prefix + else ("total_dirty", "total_writeback") + ) + (directory / "memory.stat").write_text( + f"{prefix}inactive_file {8 << 30}\n{prefix}active_file {6 << 30}\n" + f"{other_dirty} {2 << 30}\n{other_writeback} {1 << 30}\n", + encoding="ascii", + ) + + assert memory_files.remaining() == 18 << 30 + + +@pytest.mark.parametrize( + "stat_text", + [ + None, + "", + "anon 1\n", + "{prefix}active_file x\n", + "{prefix}active_file -1\n", + "{prefix}active_file +1\n", + "{prefix}active_file 1.5\n", + "{prefix}active_file 1 2\n", + "{prefix}active_file 1 \n", + "{prefix}active_file 1\r\n", + "{prefix}active_file {overflow}\n", + "{prefix}inactive_file 1\n{prefix}inactive_file 2\n", + "{prefix}active_file 1\n{prefix}active_file 2\n", + "{prefix}inactive_file 8589934592\n{prefix}active_file x\n", + "{prefix}inactive_file 8589934592\nbad line here\n", + "{prefix}inactive_file 8589934592\nanon {overflow}\n", + "{prefix}inactive_file 8589934592\nanon 1\nanon 2\n", + "{prefix}inactive_file 8589934592\n{prefix}active_file " + "9" * 5000 + "\n", + "{prefix}inactive_file 8589934592\n\n", + ], + ids=[ + "absent", + "empty", + "missing-keys", + "non-numeric", + "negative", + "signed", + "fractional", + "extra-field", + "trailing-space", + "carriage-return", + "overflow", + "duplicate-inactive", + "duplicate-active", + "partial-cache", + "malformed-line", + "unrelated-overflow", + "unrelated-duplicate", + "oversized-counter", + "blank-line", + ], +) +def test_unusable_file_cache_subtracts_nothing( + memory_files: FakeMemoryFiles, file_cache_cgroup, stat_text: str | None +): + directory, prefix = file_cache_cgroup + if stat_text is not None: + (directory / "memory.stat").write_text( + stat_text.format(prefix=prefix, overflow=_UINT64_MAX + 1), encoding="ascii" + ) + + assert memory_files.remaining() == 4 << 30 + assert memory_files.effective() == 4 << 30 + + +# The oversized case is well-formed line by line (70,000 bytes of distinct +# counters), so only the read bound can reject it. +@pytest.mark.parametrize( + "suffix", + [b"\xff", b"".join(b"pad%04d 0\n" % n for n in range(7000))], + ids=["non-ascii", "oversized"], +) +def test_non_ascii_or_oversized_file_cache_subtracts_nothing( + memory_files: FakeMemoryFiles, file_cache_cgroup, suffix: bytes +): + directory, prefix = file_cache_cgroup + (directory / "memory.stat").write_bytes( + f"{prefix}inactive_file {14 << 30}\n".encode("ascii") + suffix + ) + + assert memory_files.remaining() == 4 << 30 + + +def test_unreadable_file_cache_subtracts_nothing( + memory_files: FakeMemoryFiles, file_cache_cgroup, monkeypatch: pytest.MonkeyPatch +): + directory, _ = file_cache_cgroup + denied = directory / "memory.stat" + original_open = Path.open + + def deny_stat(path: Path, *args, **kwargs): + if path == denied: + raise PermissionError("fixture denied memory.stat") + return original_open(path, *args, **kwargs) + + monkeypatch.setattr(Path, "open", deny_stat) + assert memory_files.remaining() == 4 << 30 + assert memory_files.effective() == 4 << 30 + + +@pytest.mark.parametrize("leaf_current,remaining", [(20, 6), (4, 18)]) +def test_file_cache_is_adjusted_at_each_ancestor( + memory_files: FakeMemoryFiles, file_cache_cgroup, leaf_current: int, remaining: int +): + directory, prefix = file_cache_cgroup + memory_files.membership("7:memory:/job" if prefix else "0::/job") + write = memory_files.v1 if prefix else memory_files.v2 + write(directory / "job", 24 << 30, leaf_current << 30) + for level, inactive, active in [(directory, 8, 6), (directory / "job", 1, 1)]: + (level / "memory.stat").write_text( + f"{prefix}inactive_file {inactive << 30}\n" + f"{prefix}active_file {active << 30}\n", + encoding="ascii", + ) + + assert memory_files.remaining() == remaining << 30 + + +def test_file_cache_room_is_still_clamped_by_host( + memory_files: FakeMemoryFiles, file_cache_cgroup +): + directory, prefix = file_cache_cgroup + memory_files.meminfo.write_text("MemAvailable: 8388608 kB\n", encoding="ascii") + (directory / "memory.stat").write_text( + f"{prefix}inactive_file {14 << 30}\n", encoding="ascii" + ) + + assert memory_files.remaining() == 18 << 30 + assert memory_files.effective() == 8 << 30 + + +@pytest.mark.parametrize("control", ["limit", "usage"]) +def test_file_cache_does_not_mask_malformed_controls( + memory_files: FakeMemoryFiles, file_cache_cgroup, control: str +): + directory, prefix = file_cache_cgroup + limit_name, current_name = ( + ("memory.limit_in_bytes", "memory.usage_in_bytes") + if prefix + else ("memory.max", "memory.current") + ) + name = current_name if control == "usage" else limit_name + (directory / name).write_text("bad\n", encoding="ascii") + (directory / "memory.stat").write_text( + f"{prefix}inactive_file {14 << 30}\n", encoding="ascii" + ) + + with pytest.raises(ValueError, match=f"malformed cgroup memory {control}"): + memory_files.remaining() + assert memory_files.effective() == 0 + + +def test_v2_finite_budget_clamps_host_memavailable(memory_files: FakeMemoryFiles): + memory_files.membership("0::/tenant/job") + memory_files.v2(memory_files.root / "tenant" / "job", 12_000_000, 3_000_000) + + assert memory_files.remaining() == 9_000_000 + assert memory_files.effective() == 9_000_000 + + +def test_host_memavailable_wins_when_lower(memory_files: FakeMemoryFiles): + memory_files.membership("0::/tenant/job") + memory_files.v2(memory_files.root / "tenant" / "job", 100_000_000, 1) + + assert memory_files.effective() == 16_384 * 1024 + + +def test_nested_v2_uses_tightest_finite_ancestor(memory_files: FakeMemoryFiles): + memory_files.membership("0::/tenant/job") + memory_files.v2(memory_files.root, "max", "not-read-for-unlimited") + memory_files.v2(memory_files.root / "tenant", 20_000_000, 14_000_000) + memory_files.v2(memory_files.root / "tenant" / "job", 12_000_000, 3_000_000) + + assert memory_files.remaining() == 6_000_000 + + +def test_unlimited_v2_leaf_still_honors_finite_parent(memory_files: FakeMemoryFiles): + memory_files.membership("0::/tenant/job") + memory_files.v2(memory_files.root / "tenant", 9_000_000, 4_000_000) + memory_files.v2( + memory_files.root / "tenant" / "job", "max", "not-read-for-unlimited" + ) + + assert memory_files.remaining() == 5_000_000 + + +def test_namespaced_v2_mount_uses_root_controls(memory_files: FakeMemoryFiles): + memory_files.membership("0::/host/tenant/job") + memory_files.v2(memory_files.root, 10_000_000, 4_000_000) + + assert memory_files.remaining() == 6_000_000 + + +def test_v2_max_is_known_unlimited(memory_files: FakeMemoryFiles): + memory_files.membership("0::/") + memory_files.v2(memory_files.root, "max", "malformed-but-irrelevant") + + assert memory_files.remaining() is None + assert memory_files.effective() == 16_384 * 1024 + + +def test_counter_accepts_no_trailing_newline(memory_files: FakeMemoryFiles): + memory_files.membership("0::/") + (memory_files.root / "memory.max").write_text("10000", encoding="ascii") + (memory_files.root / "memory.current").write_text("4000", encoding="ascii") + + assert memory_files.remaining() == 6_000 + + +@pytest.mark.parametrize("current", [10_000, 10_001]) +def test_v2_current_at_or_above_limit_is_zero( + memory_files: FakeMemoryFiles, current: int +): + memory_files.membership("0::/") + memory_files.v2(memory_files.root, 10_000, current) + + assert memory_files.remaining() == 0 + assert memory_files.effective() == 0 + + +@pytest.mark.parametrize( + "value", + ["", "-1", "+1", "1.5", "garbage", " 1", "1 ", "\t1", "1\n", "1\r", "max "], +) +def test_malformed_v2_limit_fails_closed(memory_files: FakeMemoryFiles, value: str): + memory_files.membership("0::/") + memory_files.v2(memory_files.root, value, 1) + + with pytest.raises(ValueError, match="malformed cgroup memory limit"): + memory_files.remaining() + assert memory_files.effective() == 0 + + +@pytest.mark.parametrize("value", ["-1", "+1", "1.5", "unknown", "1 "]) +def test_malformed_v2_usage_fails_closed(memory_files: FakeMemoryFiles, value: str): + memory_files.membership("0::/") + memory_files.v2(memory_files.root, 10_000, value) + + with pytest.raises(ValueError, match="malformed cgroup memory usage"): + memory_files.remaining() + assert memory_files.effective() == 0 + + +def test_overflow_and_oversized_controls_fail_closed(memory_files: FakeMemoryFiles): + memory_files.membership("0::/") + memory_files.v2(memory_files.root, _UINT64_MAX + 1, 1) + assert memory_files.effective() == 0 + + memory_files.v2(memory_files.root, "9" * 256, 1) + with pytest.raises(ValueError, match="malformed cgroup memory limit"): + memory_files.remaining() + assert memory_files.effective() == 0 + + +def test_finite_limit_without_usage_fails_closed(memory_files: FakeMemoryFiles): + memory_files.membership("0::/") + (memory_files.root / "memory.max").write_text("10000\n", encoding="ascii") + + with pytest.raises(ValueError, match="usage is unavailable"): + memory_files.remaining() + assert memory_files.effective() == 0 + + +def test_permission_error_on_visible_counter_fails_closed( + memory_files: FakeMemoryFiles, monkeypatch: pytest.MonkeyPatch +): + memory_files.membership("0::/") + memory_files.v2(memory_files.root, 10_000, 1_000) + denied = memory_files.root / "memory.current" + original_open = Path.open + + def deny_current(path: Path, *args, **kwargs): + if path == denied: + raise PermissionError("fixture denied") + return original_open(path, *args, **kwargs) + + monkeypatch.setattr(Path, "open", deny_current) + with pytest.raises(ValueError, match="cannot read memory control file"): + memory_files.remaining() + assert memory_files.effective() == 0 + + +def test_excessive_membership_nesting_is_bounded(memory_files: FakeMemoryFiles): + deep = "/" + "/".join(f"level-{i}" for i in range(_CGROUP_MAX_ANCESTORS)) + memory_files.membership(f"0::{deep}") + memory_files.v2(memory_files.root, "max", 1) + + with pytest.raises(ValueError, match="membership nesting exceeds"): + memory_files.remaining() + assert memory_files.effective() == 0 + + +def test_oversized_membership_read_fails_closed(memory_files: FakeMemoryFiles): + memory_files.proc_cgroup.write_text("0::/" + "x" * (65 * 1024), encoding="ascii") + + with pytest.raises(ValueError, match="oversized cgroup membership"): + memory_files.remaining() + assert memory_files.effective() == 0 + + +def test_malformed_leaf_is_not_masked_by_valid_parent(memory_files: FakeMemoryFiles): + memory_files.membership("0::/tenant/job") + memory_files.v2(memory_files.root / "tenant", 8_000_000, 3_000_000) + memory_files.v2(memory_files.root / "tenant" / "job", "bad", 1) + + assert memory_files.effective() == 0 + + +def test_v2_is_authoritative_over_stale_v1(memory_files: FakeMemoryFiles): + memory_files.membership("0::/", "7:memory:/legacy") + memory_files.v2(memory_files.root, "max", 1) + memory_files.v1(memory_files.root / "memory" / "legacy", 4_000_000, 3_000_000) + + assert memory_files.remaining() is None + assert memory_files.effective() == 16_384 * 1024 + + +def test_missing_v2_controller_uses_bounded_v1(memory_files: FakeMemoryFiles): + memory_files.membership("0::/unified", "7:memory:/legacy") + memory_files.v1(memory_files.root / "memory" / "legacy", 8_000_000, 3_000_000) + + assert memory_files.remaining() == 5_000_000 + assert memory_files.effective() == 5_000_000 + + +def test_v1_honors_tighter_parent(memory_files: FakeMemoryFiles): + memory_files.membership("7:memory:/tenant/job") + v1_root = memory_files.root / "memory" + memory_files.v1(v1_root / "tenant", 10_000_000, 7_000_000) + memory_files.v1(v1_root / "tenant" / "job", 8_000_000, 2_000_000) + + assert memory_files.remaining() == 3_000_000 + + +def test_v1_unlimited_sentinel_is_known_unlimited(memory_files: FakeMemoryFiles): + memory_files.membership("7:memory:/") + memory_files.v1( + memory_files.root / "memory", + _CGROUP_V1_UNLIMITED_MIN, + "malformed-but-irrelevant", + ) + + assert memory_files.remaining() is None + assert memory_files.effective() == 16_384 * 1024 + + +def test_v1_current_above_limit_is_zero(memory_files: FakeMemoryFiles): + memory_files.membership("7:memory:/") + memory_files.v1(memory_files.root / "memory", 10_000, 20_000) + + assert memory_files.remaining() == 0 + assert memory_files.effective() == 0 + + +def test_malformed_membership_can_use_namespaced_root(memory_files: FakeMemoryFiles): + memory_files.v2(memory_files.root, 9_000_000, 2_000_000) + for proc_text in ("", "not:a:valid:line\n"): + memory_files.proc_cgroup.write_text(proc_text, encoding="ascii") + assert memory_files.remaining() == 7_000_000 + + +def test_absent_cgroup_files_leave_host_measurement(memory_files: FakeMemoryFiles): + memory_files.membership("0::/missing") + + assert memory_files.remaining() is None + assert memory_files.effective() == 16_384 * 1024 + + +def test_finite_cgroup_is_bound_when_meminfo_is_absent(memory_files: FakeMemoryFiles): + memory_files.membership("0::/") + memory_files.v2(memory_files.root, 7_000_000, 2_000_000) + + assert ( + effective_memory_available( + meminfo_path=memory_files.meminfo.with_name("missing"), + cgroup_root=memory_files.root, + proc_cgroup_path=memory_files.proc_cgroup, + ) + == 5_000_000 + ) + + +def test_malformed_or_overflowing_meminfo_fails_closed(memory_files: FakeMemoryFiles): + memory_files.membership("0::/missing") + for available in ("not-a-number", str(_UINT64_MAX // 1024 + 1)): + memory_files.meminfo.write_text( + f"MemTotal: 32768 kB\nMemAvailable: {available} kB\n", encoding="ascii" + ) + assert memory_files.effective() == 0 + + +def test_permission_error_on_meminfo_fails_closed( + memory_files: FakeMemoryFiles, monkeypatch: pytest.MonkeyPatch +): + memory_files.membership("0::/missing") + denied = memory_files.meminfo + original_open = Path.open + + def deny_meminfo(path: Path, *args, **kwargs): + if path == denied: + raise PermissionError("fixture denied") + return original_open(path, *args, **kwargs) + + monkeypatch.setattr(Path, "open", deny_meminfo) + assert memory_files.effective() == 0 + + +def test_completely_absent_probe_is_unknown(tmp_path: Path): + assert ( + effective_memory_available( + meminfo_path=tmp_path / "missing-meminfo", + cgroup_root=tmp_path / "missing-cgroup", + proc_cgroup_path=tmp_path / "missing-proc-cgroup", + ) + is None + ) + + +def test_hybrid_v1_and_v2_membership_is_not_a_failure(memory_files: FakeMemoryFiles): + """systemd hosts keep leftover v1 hierarchies (``1:net_cls:/``) above the unified + ``0::/...`` line; those lines must be skipped, not treated as malformed or as a + memory controller, and the real unified root carries no ``memory.max`` of its own.""" + memory_files.membership("2:cpu,cpuacct:/", "1:net_cls:/", "0::/tenant/job") + memory_files.v2(memory_files.root / "tenant", 20_000_000, 14_000_000) + memory_files.v2( + memory_files.root / "tenant" / "job", "max", "not-read-for-unlimited" + ) + + assert memory_files.remaining() == 6_000_000 + assert memory_files.effective() == 6_000_000