fix(memory): honor cgroup memory limits when sizing expert banks

[Avicennasis]

Fixes #333.

Problem

The expert-bank loader (moe/expert_banks.py::_host_ram_fits_parallel), the CPU bandwidth bench (moe/benchbw.py) and the loader benchmark (benchmarks/bench_load_weight_generic.py) size host-RAM work from /proc/meminfo MemAvailable alone. Inside a container started with --memory=N (Docker, Kubernetes) that figure is the host's, so the parallel expert reader — whose transient is one extra whole shard — is admitted against tens of GiB the cgroup will not allow, and the process is OOM-killed mid-load with no diagnostic. benchbw.py already carried a private cgroup-v2-only clamp, so the bench and the loader disagreed about how much RAM a container has.

Change

One stdlib-only module, freetoken/memory.py, with effective_memory_available():

effective = min(host MemAvailable, tightest finite cgroup budget along the ancestry)

• cgroup v2: memory.max / memory.current (max = unlimited), walking /proc/self/cgroup ancestors with bounded reads; cgroup v1 memory.limit_in_bytes / memory.usage_in_bytes fallback with the unlimited sentinel.
• Absent procfs / no memory controller → None (unknown): non-Linux hosts keep today's behaviour and are never told they have 0 bytes. Present-but-malformed or unreadable control files fail closed to 0.
• The probe runs before any bank or buffer allocation; admission is the only decision point.

Every production MemAvailable reader now goes through it; benchbw's private _cgroup_mem_headroom parser is removed (one parser). The bench's clamping and synthetic-bank sizing are otherwise unchanged. Mount-namespace hardening of the lookup is deliberately out of scope.

Tests

• tests/test_memory.py: fake /proc + 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: the auto-admission path drops to the serial build when the fake cgroup budget (299 B) cannot hold banks + one shard transient (300 B) although MemAvailable is large; an explicit --expert-load override bypasses the probe; the bench receives a finite budget unchanged and keeps its 8 GiB fallback only for an unknown probe.

These tests patch the new resolver, so they cannot run against main as written; the behavioural difference is the reproduction in the linked issue — main's _host_ram_fits_parallel returns True for a 9 GiB load inside a 2 GiB container, this branch's admission drops to the serial build.

Validation

Environment: Ubuntu 24.04 (Linux 7.0.0, x86_64, 16 cores, 125 GiB), cgroup v2, Docker 29.1.3, Python 3.12.3, torch 2.11.0+cu130, transformers 5.16.1, triton 3.6.0. No NVIDIA GPU on this host — this path has no GPU dependency; GPU-marked tests skipped as usual and I did not reproduce an end-to-end ft serve OOM inside a container.

python -m pytest -q tests/test_memory.py tests/moe/test_expert_bank_memory_admission.py   # 47 passed
python -O -m pytest -q tests/test_memory.py tests/moe/test_expert_bank_memory_admission.py   # 47 passed
python -m pytest -q tests/        # same two pre-existing failures as main (tests/models/test_muse_glimmer.py), no new ones
ruff check / ruff format --check  # clean on new files; remaining E741 are pre-existing lines
git diff --check origin/main..HEAD

Live cgroup v2 probe, python:3.12-slim with the module loaded by path:

container	effective_memory_available()	memory.max − memory.current
--memory=2g	2 136 084 480	2 136 084 480
--memory=6g	6 432 010 240	6 432 010 240
no limit	36 877 078 528 (= host MemAvailable)	max

Performance: not measured / not changed — one extra bounded procfs read at admission time.

Prepared with AI assistance; I reviewed the change and ran the validation above.

[MT-z]

Thanks for this — it lands on something that bites here for a reason you did not have in front of
you. Your write-up is about containers (docker --memory=N, Kubernetes); the shape I run is a
systemd scope, and the same bug is there. So this is an independent confirmation from a
different deployment, plus a serving run on top of it.

Setup. RTX 4090 24 GB, 61 GB RAM, Ubuntu 26.04, cgroup v2 only (no v1 hierarchy mounted, so
I cannot exercise that path). This PR applied on main @ 03c28d2 plus PR #337 @ a6bd5c0.

The gap, measured

Everything on this box is launched through a small wrapper that runs the server in a transient
systemd scope with MemoryMax=46G. Same interpreter, same call, only the scope differs:

outside any scope                   55.9 GiB
inside MemoryMax=46G                46.0 GiB   == the scope's memory.max

Before this PR, _host_ram_fits_parallel sized the parallel expert reader's one-extra-shard
transient from /proc/meminfo MemAvailable alone. On this box that is the host's ~56 GiB, so
roughly 10 GiB the cgroup would never hand out was inside the admission decision. Same failure
you describe for containers, reached from a different direction — worth knowing that the fix is not
container-specific.

It does not cost the fast path here

The obvious worry with a tighter budget is that the parallel loader stops being admitted and every
boot goes serial. It does not, at least on this checkpoint:

ornith-ai/Ornith-1.5-35B-A3B-NVFP4, --vision, --kv-reserve-tokens 262144, cap 46G
  expert banks: slow path (parallel build)
  Loading experts (parallel): 21.8G/21.8G [00:05<00:00, 4.62GB/s]
  API server is ready  (12 s from launch)

Then served an agent workload for ~55 minutes without incident — 994 prefill batches, 7.3 M prefill
tokens, peak context ~140 k, cgroup high-water 19.5 GiB of the 46 GiB cap.

Tests: tests/test_memory.py 43 passed, tests/moe/test_expert_bank_memory_admission.py 4 passed.
Across the 84 test files this PR could plausibly touch, taken before and after: 7 failures either
way, all pre-existing ULP-noise torch.equal comparisons. Nothing gained, nothing broken.

One thing I checked that turns out to be fine, in case it saves the next reader

The "max" handling looks asymmetric at first read — _cgroup_limit_remaining special-cases the
literal max only for v2 (if not v1 and limit_text in ("max", "max\n")), while v1 unlimited is
caught further down by the numeric _CGROUP_V1_UNLIMITED_MIN threshold. My first reaction was that
a v1 host writing max would fall through to the malformed cgroup memory limit raise.

It cannot: v1 and v2 use different sentinels. cgroup v1 memory.limit_in_bytes reports
unlimited as 9223372036854771712 (PAGE_COUNTER_MAX), and the string max is a v2 convention
only. 9223372036854771712 > 1 << 60, so the numeric branch catches it. The asymmetry is correct,
not an oversight. Flagging it only because it reads like one, and someone will raise it eventually.

Two observations, neither a defect

• The ancestor-walk termination guard is duplicated verbatim between _cgroup_control_directory and
  _cgroup_hierarchy_remaining. Both are right; it is a maintenance hazard if the walk ever changes.
• effective_memory_available() does a full /proc/meminfo read plus a full cgroup walk on each
  call, and the bench sampler calls it every 0.2 s. Invisible next to a model load, and the sampler
  only wants a low-water mark, so a cached bound would do — mentioning it as an observation, not a
  request.

Unrelated to the code: the branch is five commits behind

e05cff8 (Aug 31) is the base here, and main has since taken #311, #332, #336, #343 and #329.
GitHub still calls it cleanly mergeable and a local rebase onto 03c28d2 applies without conflict
with all 47 tests passing, so nothing is broken — just worth refreshing before it lands, since the
main it will actually merge into is not the one it was tested against.

Happy to test the v1 path if someone with a v1 or hybrid host wants a second pair of eyes; I only
have v2 here, so that half of the code is unexercised by anything I ran.

_{Written with AI assistance; every number above was measured on my hardware (RTX 4090, 61 GB
RAM) and I can reproduce it.}

[Avicennasis]

Thanks for the systemd-scope run — that is exactly the non-container shape the write-up could not cover, and 55.9 → 46.0 GiB is the same bug from the other direction.

Rebased onto current main (9535656). #418's quantization refactor rewrote expert_banks.py underneath this; the change itself is unchanged (a git range-diff against the previous head shows the import line moving and nothing else), and _host_ram_fits_parallel() still gates the parallel loader on the auto-pick path in load_expert_banks.

The only edit beyond that is in the two admission tests, which stub the builder to observe which loader was chosen: #418 renamed it and split it into _legacy_expert_banks / _method_expert_banks, so they now stub _legacy_expert_banks — the seam load_expert_banks reaches when no quant method is bound, at the same positional argument.

Re-ran on a CPU-only Ubuntu 24.04 x86_64 box (16 cores, 125 GiB, cgroup v2, Docker 29.1.3; Python 3.12, torch 2.11.0+cu130, transformers 5.17.0, pytest 9.1.1, run with PYTHONPATH=python since an editable install needs CUDA_HOME):

• tests/test_memory.py + tests/moe/test_expert_bank_memory_admission.py: 47 passed, normally and under python -O.
• Docker cgroup v2 probe against the rebased memory.py: --memory=2g → effective 2136227840, exactly memory.max − memory.current; --memory=6g → 6430826496, exact; unlimited → 57362182144 against a host MemAvailable of 57362636800 read milliseconds later.
• Full pytest -q tests from the repo root: this branch 4 failed / 1308 passed / 427 skipped, against main at the same commit in the same venv: 4 failed / 1261 passed / 427 skipped. Identical failure sets — tests/engine/test_cache_budget.py ×2, tests/moe/test_offload.py ×1, and tests/kernels/test_swiglu_clamp.py::test_cpu_extension_supports_swiglu_clamp (the C++ extension is not built in this venv). The +47 is this PR's tests.

Nothing in the change touches a kernel, which is why a CPU-only box is an honest environment for it; the loader path it gates is host-RAM admission.

Both of your observations are right. I have left them out of this PR so the reviewed core stays as it was, and will fold them into a follow-up (dedupe the ancestor-walk guard; cache the bound for the bench sampler) if you want them. The max asymmetry is as you say: v1 reports PAGE_COUNTER_MAX numerically and max is v2-only. Nobody on my side has a v1 host either, so that path stays covered by the fixture tests only.