Closing the Gap: Making `detect_system_memory()` Cgroup-Aware in CuZK

The Message

On March 13, 2025, an assistant working on the CuZK zero-knowledge proving engine committed a critical fix to the detect_system_memory() function. The commit message reads:

cuzk: make detect_system_memory() cgroup-aware for Docker containers

detect_system_memory() now reads cgroup v2 (memory.max) and v1 (memory.limit_in_bytes) limits and returns min(host_ram, cgroup_limit).

Previously, /proc/meminfo MemTotal was used directly, which reports the HOST's full RAM inside Docker containers — causing massive over-allocation and OOM kills when the container has a cgroup limit smaller than host RAM (e.g. 256GiB container on a 512GiB host).

Also updates DEFAULT_SAFETY_MARGIN constant from 5GiB to 10GiB to match the config default change.

2 files changed, 100 insertions(+), 4 deletions(-)

This single commit, dba80531, represents the culmination of a multi-hour debugging session that had been chasing out-of-memory (OOM) crashes on GPU-provisioned cloud instances. It is a textbook example of how a seemingly simple systems programming fix — reading the right file from the right virtual filesystem — can resolve a class of failures that had been plaguing a distributed proving pipeline.

The Problem: When `/proc/meminfo` Lies

The CuZK proving engine uses a MemoryBudget system to auto-detect available system RAM and allocate memory for GPU workloads, SRS tables, and synthesis working sets. The detection was straightforward: read /proc/meminfo, parse MemTotal, subtract a safety margin, and use the result as the memory budget. On bare metal, this works perfectly.

But CuZK was being deployed on vast.ai, a cloud GPU marketplace where instances run inside Docker containers with cgroup memory limits. A typical vast.ai instance might have 503 GiB of host RAM but a cgroup limit of only 342 GiB. Inside such a container, /proc/meminfo still reports the host's full 503 GiB — the Linux kernel does not virtualize /proc/meminfo per container. The result was catastrophic: CuZK would compute a budget of ~493 GiB (503 GiB minus a 10 GiB safety margin) and attempt to allocate memory against a container that could only physically provide 342 GiB. The kernel's OOM killer would then terminate the process, causing benchmark failures and service interruptions.

This is the kind of bug that is easy to miss during local development (where containers typically have unlimited memory) and only manifests in production. The assistant had already built a shell-based workaround (memcheck.sh) that read cgroup limits and passed an explicit --budget flag to override auto-detection. But the proper fix was to make the Rust code itself cgroup-aware, so that total_budget = "auto" works correctly everywhere — bare metal, Docker, and any cgroup-constrained environment.

The Reasoning: Why This Message Was Written

The commit message at index 3896 was not written in isolation. It was the product of a deliberate, multi-step reasoning process that unfolded over the preceding messages ([msg 3875] through [msg 3895]). To understand why this message exists, we must trace the assistant's thinking.

The session began with the user asking the assistant to continue with next steps ([msg 3872]). The assistant surveyed the state of the repository and identified uncommitted changes: Docker scripts, a memcheck.sh utility, vast-manager improvements, and a safety margin increase in config.rs ([msg 3873]). The user then explicitly directed the assistant to "fix detect_system_memory() in Rust to be cgroup-aware natively" ([msg 3874]), recognizing that the shell-based workaround was insufficient.

The assistant committed the outstanding changes first ([msg 3879]), then turned to the Rust fix. It read the current implementation of detect_system_memory() ([msg 3882]), traced the full call chain from config.rs to memory.rs ([msg 3883]), and studied how total_budget = "auto" was resolved ([msg 3884]). It then read the memcheck.sh script to understand the cgroup file paths and logic that the shell workaround used ([msg 3885]).

This research phase reveals the assistant's methodical approach. Rather than blindly implementing cgroup detection, it first understood the existing architecture: detect_system_memory() returns an Option<u64>, called from resolve_total_budget() in config.rs with a fallback of 256 GiB. The function was simple — read MemTotal from /proc/meminfo — and needed to be extended without breaking backward compatibility.

The assistant then articulated the design in its reasoning ([msg 3886]):

The key paths: - cgroup v2: /sys/fs/cgroup/memory.max — contains bytes or "max" (unlimited) - cgroup v1: /sys/fs/cgroup/memory/memory.limit_in_bytes — contains bytes; values > ~2^60 mean unlimited

I'll also keep the function backwards-compatible: if cgroup files don't exist (bare metal), it still falls back to /proc/meminfo.

This design decision — to return the minimum of host RAM and any cgroup constraint, with graceful fallback — is the core intellectual contribution of this message. It is simple, correct, and defensive.

The Implementation: What Changed

The commit touched two files: extern/cuzk/cuzk-core/src/memory.rs (the main implementation and tests) and extern/cuzk/cuzk-core/src/config.rs (documentation update). The 100 lines added and 4 lines deleted represent a significant rewrite of the memory detection logic.

The new detect_system_memory() function follows this algorithm:

Read /proc/meminfo to get MemTotal (host RAM in bytes).
Attempt to read /sys/fs/cgroup/memory.max (cgroup v2). If the file exists and contains a numeric value (not "max"), parse it as a byte limit.
Attempt to read /sys/fs/cgroup/memory/memory.limit_in_bytes (cgroup v1). If the file exists and the value is below a threshold (values above ~2^60 indicate "unlimited" in cgroup v1), parse it as a byte limit.
Return the minimum of all available values. If no cgroup limits exist, return the host RAM value. This is a textbook implementation of cgroup-aware memory detection, and it mirrors the logic already present in memcheck.sh. The assistant also updated the DEFAULT_SAFETY_MARGIN constant from 5 GiB to 10 GiB to match the earlier config change, ensuring consistency between the Rust constant and the configuration default.

Assumptions Made

Several assumptions underpin this commit:

Assumption 1: Cgroup limits are always more restrictive than host RAM. The function returns min(host_ram, cgroup_limit), which assumes the cgroup limit is the binding constraint. This is correct for Docker containers, but in theory a cgroup could have a limit larger than host RAM (e.g., a swap-backed limit). In practice, this is vanishingly rare, and taking the minimum is the safe choice.

Assumption 2: The cgroup v2 "max" token means unlimited. The cgroup v2 interface uses the string "max" to indicate no limit. The assistant correctly handles this by skipping the v2 value if it equals "max".

Assumption 3: Cgroup v1 values above ~2^60 mean unlimited. This is a well-known convention in the cgroup v1 interface: memory.limit_in_bytes is set to a very large value (often 9223372036854771712 or similar) when no limit is configured. The assistant's threshold-based check is the standard approach.

Assumption 4: The function remains backwards-compatible. If no cgroup files exist (bare metal, or containers without memory limits), the function falls back to /proc/meminfo as before. This ensures the change is safe to deploy everywhere.

Assumption 5: The shell-based workaround (--budget override) can remain as defense-in-depth. The assistant explicitly noted ([msg 3890]) that the entrypoint's explicit --budget flag "is still useful as a double-safety" and that "both mechanisms are complementary." This is a pragmatic assumption: having two layers of protection is better than one, especially when dealing with memory budgeting in production.

Input Knowledge Required

To understand this message, a reader needs:

Linux cgroup internals: Knowledge of cgroup v1 and v2 memory limit files, the "max" token convention in v2, and the large-value convention in v1.
Docker memory semantics: Understanding that Docker containers use cgroups to enforce memory limits, and that /proc/meminfo is not virtualized — it always reports host RAM.
The CuZK memory architecture: Awareness that detect_system_memory() feeds into MemoryBudget, which is used to allocate SRS tables, PCE working sets, and GPU pinned memory pools.
The deployment context: Understanding that CuZK runs on vast.ai cloud instances where cgroup limits are the norm, not the exception.
Rust file I/O patterns: Familiarity with reading files in Rust, parsing integers from strings, and handling Option return types.

Output Knowledge Created

This commit produces several forms of knowledge:

A correct, deployable memory detection function that works in Docker containers, bare metal, and any cgroup-constrained environment.
A documented design pattern for cgroup-aware resource detection in Rust, which can be reused for other resource limits (CPU, swap, I/O).
A validated fix that was subsequently tested on real vast.ai instances (a 961 GiB cgroup-limited machine and a 342 GiB cgroup-limited machine), confirming that the budget was correctly derived from the cgroup limit rather than host RAM.
A defense-in-depth strategy combining native Rust detection with shell-based override, providing two layers of protection against OOM kills.

The Thinking Process

The assistant's thinking process is visible in the sequence of tool calls leading up to this commit. Let us examine the key reasoning steps.

Step 1: Problem identification. The assistant recognized that detect_system_memory() reading /proc/meminfo was the root cause of OOM kills in Docker containers. This was not a new insight — the shell workaround already existed — but the assistant understood that the proper fix was to make the Rust code itself cgroup-aware.

Step 2: Architecture tracing. The assistant traced the call chain from config.rs to memory.rs, reading both files and understanding how total_budget = "auto" flows through the system. This ensured the fix would be applied at the right level.

Step 3: Reference implementation study. The assistant read memcheck.sh to understand the cgroup file paths and parsing logic. This is a pragmatic choice: rather than reinventing the wheel, it reused the proven logic from the shell script.

Step 4: Design articulation. The assistant explicitly stated the three cgroup paths to check and the fallback behavior, demonstrating a clear mental model before writing code.

Step 5: Implementation and testing. The assistant edited the Rust file, updated the tests, and verified compilation (though the local Cargo was too old for the full dependency tree, rustfmt --check confirmed the file was syntactically valid).

Step 6: Complementary strategy. The assistant noted that the shell-based override could remain as defense-in-depth, showing an understanding of production deployment realities.

Step 7: Commit and documentation. The commit message itself is well-structured, explaining the problem, the fix, and the impact. It uses a concrete example ("256GiB container on a 512GiB host") to illustrate the severity.

Broader Significance

This commit is a microcosm of a larger pattern in systems engineering: the gap between development environments and production environments. In development, containers often have unlimited memory, and /proc/meminfo happens to report the correct value. In production, cgroup limits are the norm, and the same code path produces catastrophic over-allocation. The fix — reading the right virtual file — is simple in retrospect, but it required understanding the Linux kernel's container isolation model, the Docker runtime's cgroup configuration, and the CuZK memory architecture.

The commit also illustrates the value of defense-in-depth. The assistant kept the shell-based override even after implementing the native fix, recognizing that two layers of protection are better than one. This is a mature engineering judgment that prioritizes reliability over elegance.

Finally, the commit demonstrates the importance of tracing the full call chain before making changes. The assistant did not just edit detect_system_memory() in isolation; it read config.rs, understood how resolve_total_budget() called the function, and ensured the fix would propagate correctly through the entire memory budgeting system. This systems thinking — understanding the whole before changing the part — is what separates a robust fix from a fragile one.

Closing the Gap: Making detect_system_memory() Cgroup-Aware in CuZK