From de9df4a7c98d54bd381b6429851be932a1819a9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Tue, 1 Sep 2026 15:36:05 -0400 Subject: [PATCH] fix(memory): honor cgroup memory limits when sizing expert banks The expert-bank loader, the CPU bandwidth bench and the loader benchmark all sized host-RAM work from /proc/meminfo MemAvailable alone. Inside a container started with --memory=N that figure is the host's, not the cgroup's: the parallel expert reader was admitted against tens of GiB of "available" RAM while the cgroup could hold a fraction of it, and the process was OOM-killed mid-load. benchbw.py carried its own cgroup-v2-only parser; the loader had none. Rule: effective available memory = min(host MemAvailable, remaining finite cgroup budget). The cgroup budget is the tightest finite limit-minus-usage along the process's cgroup ancestry (v2 memory.max / memory.current with the `max` sentinel; v1 memory.limit_in_bytes / memory.usage_in_bytes with the unlimited sentinel as a fallback), resolved from /proc/self/cgroup with bounded reads and a bounded ancestor walk. Absent procfs or no visible memory controller yields None (unknown), so non-Linux hosts keep their previous best-effort behaviour and are never told they have 0 bytes; a present but malformed or unreadable control file fails closed to 0. The probe runs before any bank or buffer allocation: admission is the only decision point, because once the parallel reader's whole-shard anonymous buffers exist there is nothing reclaimable left to give back. freetoken/memory.py is stdlib-only (no torch import) so it can be loaded by path inside a slim container for diagnosis. Every production MemAvailable reader now goes through it: moe/expert_banks ._host_ram_fits_parallel, moe/benchbw._available_ram_bytes (the private _cgroup_mem_headroom parser is removed so there is one cgroup parser; the bench's clamping and synthetic-bank sizing are otherwise unchanged), and benchmarks/bench_load_weight_generic.MemSampler. Tests: tests/test_memory.py drives the resolver against fake /proc and cgroup trees under tmp_path -- unlimited, finite, nested ancestor-min, namespaced root, usage above limit, malformed, oversized, permission denied, v1 fallback and sentinel, absent files, and the hybrid `1:net_cls:/` + `0::/...` membership layout systemd and Docker both produce. tests/moe/test_expert_bank_memory_admission.py checks that the auto-admission path drops to the serial build when the fake cgroup budget (299 B) cannot hold the banks plus one shard transient (300 B) although MemAvailable is large, that an explicit --expert-load override bypasses the probe, and that the bench receives a finite budget unchanged while keeping its legacy 8 GiB fallback only for an unknown (None) probe. Validated on CPU-only Linux (Ubuntu 24.04, Docker 29.1, cgroup v2; no GPU is involved in this path): 47 tests pass under python and python -O; the full suite has the same two pre-existing failures as origin/main (tests/models/test_muse_glimmer.py) and no new ones; python:3.12-slim with --memory=2g / --memory=6g / no limit reports effective = memory.max - memory.current (2136084480 / 6432010240) and host MemAvailable respectively. Mount-namespace hardening of the cgroup lookup is deliberately out of scope for this change. --- benchmarks/bench_load_weight_generic.py | 16 +- python/freetoken/memory.py | 294 +++++++++++++++ python/freetoken/moe/benchbw.py | 26 +- python/freetoken/moe/expert_banks.py | 15 +- .../moe/test_expert_bank_memory_admission.py | 108 ++++++ tests/test_memory.py | 346 ++++++++++++++++++ 6 files changed, 763 insertions(+), 42 deletions(-) create mode 100644 python/freetoken/memory.py create mode 100644 tests/moe/test_expert_bank_memory_admission.py create mode 100644 tests/test_memory.py diff --git a/benchmarks/bench_load_weight_generic.py b/benchmarks/bench_load_weight_generic.py index 0e1905a0a..b368c9194 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 000000000..32832aa9a --- /dev/null +++ b/python/freetoken/memory.py @@ -0,0 +1,294 @@ +"""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_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_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}") + return True, max(0, limit - current) + + +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 visible controls raise ``ValueError``. + """ + 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. + """ + 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 f3e5359a4..a7c45e407 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -56,6 +56,7 @@ single_gpu_arg, ) from freetoken.kernel.pinned import alloc_pinned_tensor +from freetoken.memory import effective_memory_available from freetoken.moe.cpu_executor import physical_core_cpus, resolve_threads_and_affinity from freetoken.utils import init_logger @@ -149,29 +150,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 8b6116ba8..cbea81945 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -22,6 +22,7 @@ import torch +from freetoken.memory import effective_memory_available from freetoken.utils import init_logger from .offload_cache import _BANK_BYTES_PER_EXPERT, _BANK_SCHEMAS @@ -349,16 +350,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: # resolve a hub id to its local cache dir (no-op for a local path) so glob sees the shards 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 000000000..4638f61b9 --- /dev/null +++ b/tests/moe/test_expert_bank_memory_admission.py @@ -0,0 +1,108 @@ +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.memory import effective_memory_available + + +def test_loader_auto_admission_consumes_effective_memory( + tmp_path, monkeypatch: pytest.MonkeyPatch +): + """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 + import freetoken.utils.hf as hf + + # The bank estimate is 200 bytes plus a 100-byte parallel-reader transient. + # Host MemAvailable is large, but the fake cgroup has only 299 bytes left. + (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("701\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(hf, "download_hf_weight", lambda path: path) + monkeypatch.setattr(ftw, "is_ftw_checkpoint", lambda _path: False) + + chosen = [] + + def build(*args, **kwargs): + chosen.append(args[5]) # _build_expert_banks(..., parallel, ...) + return expert_banks.ExpertBanks("bf16", {}) + + monkeypatch.setattr(expert_banks, "_build_expert_banks", build) + config = SimpleNamespace(num_moe_layers=1, expert_quant="none") + + expert_banks.load_expert_banks( + str(tmp_path), config, device=torch.device("cpu"), dtype=torch.bfloat16 + ) + + assert chosen == [False] + + +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 = [] + + def build(*args, **kwargs): + chosen.append(args[5]) + return expert_banks.ExpertBanks("bf16", {}) + + monkeypatch.setattr(expert_banks, "_build_expert_banks", build) + config = SimpleNamespace(num_moe_layers=1, expert_quant="none") + + expert_banks.load_expert_banks( + str(tmp_path), + config, + device=torch.device("cpu"), + dtype=torch.bfloat16, + 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/test_memory.py b/tests/test_memory.py new file mode 100644 index 000000000..c136b88e8 --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,346 @@ +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) + + +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