From 2aa1c6362f999525aafcc03b907cc95945be5b35 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 1/6] 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 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..32832aa9a6 --- /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 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..221814212d --- /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]) # _legacy_expert_banks(..., dummy, parallel, ...) + return expert_banks.ExpertBanks("bf16", {}) + + monkeypatch.setattr(expert_banks, "_legacy_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, "_legacy_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 0000000000..c136b88e88 --- /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 From 6f1ce1de4a0d855c7cad7bad1cde5882b888f3ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Fri, 13 Nov 2026 01:41:43 -0500 Subject: [PATCH 2/6] test(moe): drive the admission tests through the current bank builder The two loader tests replaced expert_banks._legacy_expert_banks and passed dtype= to load_expert_banks. Both went away with the q4_0 provider (#690), so on current main the tests stop at the monkeypatch (AttributeError) before any admission code runs. They now stand in for _method_expert_banks, the builder a raw checkpoint reaches, pass a bound-method stand-in (kind NONE), and read the chosen mode by parameter name from the real builder's signature instead of by position. The download_hf_weight patch is dropped: the shard list now comes from the loader's iter_weight_files, which resolves a local directory itself. The auto-admission test is two-sided: 299 bytes of cgroup room against the 300 the banks and one shard need gives the serial build, 301 keeps the parallel reader. With only the serial half, a loader that never proposed the parallel reader would pass without reaching the probe. --- .../moe/test_expert_bank_memory_admission.py | 61 ++++++++++++------- 1 file changed, 38 insertions(+), 23 deletions(-) diff --git a/tests/moe/test_expert_bank_memory_admission.py b/tests/moe/test_expert_bank_memory_admission.py index 221814212d..3d368cce73 100644 --- a/tests/moe/test_expert_bank_memory_admission.py +++ b/tests/moe/test_expert_bank_memory_admission.py @@ -1,22 +1,48 @@ +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 + 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 - 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. + # 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" @@ -26,7 +52,7 @@ def test_loader_auto_admission_consumes_effective_memory( 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") + (cgroup_root / "memory.current").write_text(f"{cgroup_current}\n", encoding="ascii") def effective_fixture(): return effective_memory_available( @@ -38,23 +64,18 @@ def effective_fixture(): 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]) # _legacy_expert_banks(..., dummy, parallel, ...) - return expert_banks.ExpertBanks("bf16", {}) - - monkeypatch.setattr(expert_banks, "_legacy_expert_banks", build) + 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, device=torch.device("cpu"), dtype=torch.bfloat16 + str(tmp_path), + config, + method=SimpleNamespace(kind=QuantKind.NONE), + device=torch.device("cpu"), ) - assert chosen == [False] + assert chosen == [parallel] def test_explicit_parallel_loader_override_bypasses_auto_admission( @@ -70,20 +91,14 @@ def test_explicit_parallel_loader_override_bypasses_auto_admission( "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, "_legacy_expert_banks", build) + 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"), - dtype=torch.bfloat16, parallel=True, ) From a51a212ae3c25f4638433c5748896e77811097b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Fri, 13 Nov 2026 01:42:24 -0500 Subject: [PATCH 3/6] test(moe): give the shard-sizing probe test its memory figure at the resolver test_host_ram_fits_parallel_sizes_subdirectory_shards (#819) fakes /proc/meminfo by replacing builtins.open. _host_ram_fits_parallel no longer opens that file itself: it asks effective_memory_available(), which reads through pathlib, so the fake is never consulted, the probe sees the machine's real figure and the test fails on any host with more than 3 MiB available. The test's subject is which shards get sized, not how the figure is read, so it now sets the same 512 kB at the name the probe calls. The fixture and the assertion are unchanged. --- tests/moe/test_offload.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) 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)) From f9b2797cec26ba7e2e3ed4246564de061f657a9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Wed, 4 Nov 2026 22:45:13 -0500 Subject: [PATCH 4/6] fix(memory): count cgroup file cache as available room --- python/freetoken/memory.py | 35 +++++- tests/test_memory.py | 214 +++++++++++++++++++++++++++++++++++++ 2 files changed, 247 insertions(+), 2 deletions(-) diff --git a/python/freetoken/memory.py b/python/freetoken/memory.py index 32832aa9a6..40054d7bf4 100644 --- a/python/freetoken/memory.py +++ b/python/freetoken/memory.py @@ -16,6 +16,7 @@ # 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 @@ -81,6 +82,33 @@ def _host_memory_available(meminfo_path: Path) -> int | None: return kib * 1024 +def _cgroup_file_cache(directory: Path, *, v1: bool = False) -> int: + """Return file-LRU bytes, or zero if optional statistics cannot be trusted.""" + 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 "" + return counters.get(prefix + "inactive_file", 0) + counters.get( + prefix + "active_file", 0 + ) + + def _cgroup_pair_remaining( directory: Path, limit_name: str, @@ -111,7 +139,8 @@ def _cgroup_pair_remaining( 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) + 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]: @@ -234,7 +263,8 @@ def _cgroup_memory_remaining( 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``. + 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) @@ -275,6 +305,7 @@ def effective_memory_available( 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)) diff --git a/tests/test_memory.py b/tests/test_memory.py index c136b88e88..539f3b2d0d 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -66,6 +66,220 @@ 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( + "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 + + +@pytest.mark.parametrize( + "suffix", [b"\xff", b"x" * (65 * 1024)], 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) From d87b5086df515c1f1cceea3663b2093d5fd6a94a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Tue, 10 Nov 2026 02:39:02 -0500 Subject: [PATCH 5/6] test(memory): make the oversized memory.stat case need the read bound The oversized fixture was 65 KiB of "x", which the line parser rejects on its own, so the case passed with the 64 KiB bound removed. Use 70,000 bytes of distinct well-formed counters instead. --- tests/test_memory.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_memory.py b/tests/test_memory.py index 539f3b2d0d..83717e08b9 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -197,8 +197,12 @@ def test_unusable_file_cache_subtracts_nothing( 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"x" * (65 * 1024)], ids=["non-ascii", "oversized"] + "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 From 323fc93504a49f2a41e4129236ed4ab9c5a072b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Tue, 10 Nov 2026 02:57:23 -0500 Subject: [PATCH 6/6] fix(memory): do not count dirty or writeback pages as cgroup room The file LRU lists include pages that are dirty or under writeback. They cannot be dropped until written, and a cgroup at its limit is OOM-killed before that finishes. Take file_dirty + file_writeback (v1: total_dirty + total_writeback) back out of the file-cache credit, clamped at zero. --- python/freetoken/memory.py | 13 +++++++++-- tests/test_memory.py | 45 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/python/freetoken/memory.py b/python/freetoken/memory.py index 40054d7bf4..ef4294d1c2 100644 --- a/python/freetoken/memory.py +++ b/python/freetoken/memory.py @@ -83,7 +83,12 @@ def _host_memory_available(meminfo_path: Path) -> int | None: def _cgroup_file_cache(directory: Path, *, v1: bool = False) -> int: - """Return file-LRU bytes, or zero if optional statistics cannot be trusted.""" + """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 @@ -104,9 +109,13 @@ def _cgroup_file_cache(directory: Path, *, v1: bool = False) -> int: # v1 totals match hierarchical usage. Unlike `file`, file LRU lists exclude shmem. prefix = "total_" if v1 else "" - return counters.get(prefix + "inactive_file", 0) + counters.get( + 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( diff --git a/tests/test_memory.py b/tests/test_memory.py index 83717e08b9..219a9058cd 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -139,6 +139,51 @@ def test_file_cache_uses_only_matching_file_lists( 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", [