Closing the Gap: Making Rust's Memory Detection Cgroup-Aware

Introduction

In the ongoing battle to stabilize the CuZK proving engine on memory-constrained vast.ai GPU instances, message <msg id=3886> marks a critical inflection point. After weeks of deploying shell-script workarounds, Docker wrappers, and empirical probing utilities, the assistant makes the decision to fix the root cause directly in Rust: rewriting detect_system_memory() to be natively cgroup-aware. This single edit transforms the memory budgeting system from one that relied on a fragile chain of external scripts into one that correctly understands its container-imposed limits from the ground up.

The message itself is deceptively brief — a few lines of reasoning followed by an edit command. But the decision it encodes represents a fundamental shift in architectural philosophy, one that resolves a class of OOM crashes that had plagued the project across multiple deployment cycles.

The Problem: A Blind Spot in Containerized Environments

The backstory is essential to understanding why this message matters. The CuZK proving engine had a detect_system_memory() function that read /proc/meminfo to determine available RAM. On bare metal, this works perfectly — MemTotal reflects the physical memory installed. But inside a Docker container, /proc/meminfo still reports the host's total RAM, not the container's cgroup-imposed limit. A vast.ai instance with 503 GiB of host RAM but a 342 GiB container limit would see the full 503 GiB, causing the memory budget to be set far too high. The result was predictable: the kernel's OOM killer would terminate the cuzk process when it exceeded the cgroup limit, producing exit code 137 and a failed benchmark.

The project had already developed an elaborate defense-in-depth strategy. A shell script called memcheck.sh (see <msg id=3879>) was written to detect cgroup limits and pass an explicit --budget flag to the cuzk binary. A memprobe utility was built to empirically measure kernel overhead. An OOM recovery loop was added to benchmark.sh to retry with reduced budgets. But all of these were external workarounds — bandages applied around a wound that the Rust code itself was creating every time it consulted /proc/meminfo.

The Reasoning: Why Fix It in Rust?

The assistant's reasoning in <msg id=3886> reveals a clear design philosophy. After reading the existing implementation and tracing the full budget resolution flow through config.rs (see <msg id=3883> and <msg id=3884>), the assistant concludes: "The fix is straightforward: make detect_system_memory() cgroup-aware."

This is not a casual observation. The assistant had just committed the shell-script workaround in <msg id=3879> — memcheck.sh, entrypoint.sh, and the entire Docker infrastructure. That commit was titled "cuzk: docker scripts, memcheck utility, vast-manager memcheck/SSH improvements" and represented hundreds of lines of defensive scripting. Yet the assistant immediately recognized that the proper fix belonged in Rust, not in shell. The reasoning is implicit but powerful: a workaround that passes an explicit --budget flag from a shell script is fragile, deployment-specific, and bypasses the elegant "auto" budget system that the Rust code already had. By making the Rust function itself cgroup-aware, every consumer of detect_system_memory() — whether invoked via total_budget = "auto" in config, or called programmatically — would automatically get the correct answer.

The Design: Three Sources, One Minimum

The assistant's plan, articulated in the message, is elegantly simple:

  1. Read /proc/meminfo for MemTotal (host RAM) — the existing behavior, preserved as the baseline.
  2. Check cgroup v2: /sys/fs/cgroup/memory.max — the modern cgroup interface, where the file contains either a byte count or the string "max" (meaning unlimited).
  3. Check cgroup v1: /sys/fs/cgroup/memory/memory.limit_in_bytes — the legacy interface, where values greater than approximately 2^60 indicate "no limit."
  4. Return the minimum of all available values. This design is deliberately conservative. By taking the minimum, the function ensures that even if one source reports an inflated value (host RAM from /proc/meminfo inside a container), the cgroup constraint will cap it. On bare metal where cgroup files don't exist or report "max"/"unlimited", the function falls back to /proc/meminfo — maintaining full backward compatibility. The assistant explicitly notes this backward-compatibility requirement: "I'll also keep the function backwards-compatible: if cgroup files don't exist (bare metal), it still falls back to /proc/meminfo." This is a production-minded decision. The code must work identically on a developer's workstation, a bare-metal server, and a Docker container on vast.ai, without any configuration changes.

Assumptions and Their Validity

The message rests on several assumptions, most of which are well-justified:

Cgroup filesystem paths are standard. The assistant assumes /sys/fs/cgroup/memory.max for v2 and /sys/fs/cgroup/memory/memory.limit_in_bytes for v1. These are the standard paths used by Linux cgroup controllers and are consistent across Docker, Kubernetes, and virtually all container runtimes. This assumption is safe.

The "max" sentinel in v2 and the ~2^60 threshold in v1 reliably indicate "unlimited." In cgroup v2, memory.max literally contains the string "max" when no limit is set. In cgroup v1, memory.limit_in_bytes contains a very large number (often the host's total RAM or a value like 9223372036854771712, which is close to 2^63). The assistant's plan to treat values above ~2^60 as "unlimited" is a well-known heuristic in the container ecosystem. This assumption is correct.

Reading these files is safe and fast. The assistant assumes that file I/O to /sys/fs/cgroup/ paths is reliable. In practice, these are pseudo-files exposed by the kernel's cgroup subsystem and are always readable (assuming the process is inside the cgroup). The operation is a tiny read, negligible in cost.

The minimum of host RAM and cgroup limit is the correct effective memory. This is the core semantic assumption. Is it always true that the effective memory for a process is min(host_ram, cgroup_limit)? In almost all practical cases, yes — the cgroup limit is the hard ceiling. However, there is a subtlety: the kernel itself reserves some memory for its own operations (page cache, slab allocator, etc.), and the cgroup limit applies to the sum of all processes in the cgroup. The function returns a value that is then reduced by a safety margin (default 10 GiB), so the assumption is conservative enough to be safe.

The Knowledge Flow: Input and Output

Input knowledge required to understand this message includes:

The Broader Context: A Shift from External to Internal

This message is particularly interesting when viewed as part of a larger pattern in the conversation. The project had been iterating on external solutions — shell scripts, Docker wrappers, empirical probing — for several segments (see segments 24–29). Each iteration added more sophistication: the pinned memory pool, the PI-controlled dispatch pacer, the memcheck utility, the memprobe binary, the OOM recovery loop. But all of these were external to the Rust core.

Message <msg id=3886> represents a conscious decision to pull the critical piece of logic — memory limit detection — back into the Rust code where it belongs. This is a classic refactoring pattern: workaround first, then root-cause fix. The shell scripts proved the concept (cgroup-aware detection works), validated the paths, and bought time. But the permanent solution had to live in the same layer as the code that uses the result.

The assistant's subsequent messages confirm this. In <msg id=3890>, the assistant explicitly notes: "With native cgroup detection, that becomes a safety net rather than the primary mechanism." The shell scripts are not removed — they remain as defense-in-depth — but their role shifts from essential to redundant. The Rust code now does the right thing on its own.

Conclusion

Message <msg id=3886> is a study in focused, principled engineering. In a few lines of reasoning and a single edit command, the assistant identifies the root cause of a persistent crash pattern, designs a minimal fix that respects backward compatibility, and executes it with surgical precision. The decision to implement cgroup awareness in Rust rather than continuing to rely on shell-script workarounds reflects a mature understanding of where architectural boundaries should be drawn. The fix is small in scope but large in impact — it transforms the memory budgeting system from one that was blind to container limits into one that sees them natively, closing the gap between what the code thought was available and what the container actually allowed.