From Host Blindness to Container Awareness: How Cgroup-Aware Memory Detection Rescued a GPU Proving Engine from OOM Kills

Introduction

In the world of GPU-accelerated zero-knowledge proving, memory is the most unforgiving resource. Run out, and the kernel's OOM killer terminates your process without warning, without a crash dump, and without recourse. The CuZK proving engine—a high-performance GPU-accelerated proof generator for Filecoin—had been suffering exactly this fate on vast.ai cloud instances. The root cause was subtle but devastating: inside Docker containers, the standard method of detecting system memory—reading /proc/meminfo—returns the host machine's total RAM, not the container's cgroup-imposed limit. A container limited to 342 GiB of RAM on a 503 GiB host would see 503 GiB available, budget 493 GiB, and promptly exceed the container's ceiling, triggering an OOM kill.

This article traces the arc of a multi-session engineering effort that solved this problem through a combination of Rust-level cgroup-aware memory detection, shell script debugging, empirical measurement tools, and adaptive recovery mechanisms. The work spanned implementing the core fix, fixing bugs discovered during live testing, deploying to production instances, validating the fix under real workloads, and ultimately building additional safety layers when the initial fix proved insufficient for the most constrained environments.

The Core Fix: Making detect_system_memory() Cgroup-Aware

The heart of the problem lay in the Rust function detect_system_memory() in extern/cuzk/cuzk-core/src/memory.rs. This function, which feeds into the MemoryBudget system that controls all memory allocations for GPU pinned buffers, SRS parameters, and proving operations, was reading /proc/meminfo directly. Inside a Docker container, this reports the host's total physical RAM, not the container's cgroup limit.

The fix, committed as dba80531 with the message "cuzk: make detect_system_memory() cgroup-aware for Docker containers" ([msg 3896]), rewrote the function to implement a three-source detection algorithm:

  1. Read /proc/meminfo to get MemTotal (host RAM in bytes) as a baseline.
  2. Attempt to read /sys/fs/cgroup/memory.max (cgroup v2). If the file exists and contains a numeric value (not the string "max"), parse it as a byte limit.
  3. 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.
  4. Return the minimum of all available values. If no cgroup limits exist, fall back to host RAM. This algorithm is simple, correct, and defensive. It handles bare metal (no cgroup files exist, falls back to /proc/meminfo), cgroup v1 containers, and cgroup v2 containers. The commit also updated the DEFAULT_SAFETY_MARGIN constant from 5 GiB to 10 GiB, matching the config default change that had been made earlier. The implementation was preceded by a thorough research phase. The assistant read the current implementation ([msg 3882]), traced the full call chain from config.rs to memory.rs ([msg 3883]), studied how total_budget = "auto" was resolved ([msg 3884]), and read the memcheck.sh script as a reference implementation for the cgroup file paths ([msg 3885]). This methodical approach—understand the architecture before writing code—is a hallmark of robust engineering.

Bugs Discovered in Live Testing: GPU JSON Parsing and Pinning Detection

When the updated Docker image was deployed to a live vast.ai instance and memcheck.sh was run, two bugs immediately surfaced ([msg 3911]). The first was a classic shell scripting pitfall: the GPU JSON parsing used IFS=', ' (Internal Field Separator set to comma and space) to split CSV fields from nvidia-smi. GPU names like "NVIDIA GeForce RTX 4090" contain spaces, so setting IFS=', ' caused bash to split on both commas and spaces, fragmenting the GPU name into separate tokens and producing malformed JSON. The fix was straightforward: change IFS=',' to split only on commas, preserving spaces within GPU names ([msg 3917]).

The second bug was more subtle and required experimental verification. The memcheck script checked whether the system could perform CUDA pinned memory allocations by examining ulimit -l—the maximum size of memory that a process can lock with mlock(). On the vast.ai instance, ulimit -l returned only 8192 kB (8 MB), far below the 50 GiB threshold the script considered necessary. This caused memcheck to report can_pin: false, which would prevent the system from using pinned memory for GPU transfers—a critical performance optimization.

The assistant hypothesized that CUDA's cudaHostAlloc function, which allocates pinned (page-locked) host memory for GPU transfers, goes through the NVIDIA kernel driver (/dev/nvidia*), not through the standard mlock() system call. To test this, the assistant SSH'd into the live instance and ran a Python script using ctypes to call cuInit and cudaHostAlloc directly ([msg 3926]). The result: cudaHostAlloc(1GiB): 0 (0=success). Pinned memory allocation worked perfectly despite the 8 MB ulimit. The NVIDIA kernel driver bypasses RLIMIT_MEMLOCK for its own pinned memory allocations.

The fix was to check for CUDA capability via nvidia-smi presence instead of relying on ulimit -l. If an NVIDIA GPU is detected, the system can pin memory—regardless of the ulimit setting ([msg 3930]). The ulimit value is still reported for informational purposes, but it no longer determines the can_pin flag.

A third, latent bug was also fixed: the entrypoint script used set -e (exit on error) combined with raw jq calls that had no fallback behavior. If jq encountered a parse error—for example, because the GPU JSON from memcheck was malformed—the entire entrypoint would crash. The fix wrapped jq calls in a helper function with fallback defaults ([msg 3918]).

These three fixes were committed together at [msg 3947] with the message "cuzk: fix memcheck GPU JSON parsing and pinning detection." The commit touched two files with 33 insertions and 18 deletions.

Deployment and Validation on Production Instances

With the fixes committed and a new Docker image built and pushed, the assistant deployed to two real vast.ai instances to validate the cgroup-aware memory detection under production conditions.

The 961 GiB Instance

The first instance (C.32874928) was an RTX 4090 machine with 2003 GiB of host RAM and a cgroup limit of 961 GiB. The assistant deployed the updated image, waited for the entrypoint to complete, and checked the status ([msg 3958]). The results were clear:

cat /sys/fs/cgroup/memory.current  → 584199278592 (~544 GiB RSS)
cat /sys/fs/cgroup/memory.max      → 1032768716800 (~961 GiB limit)

The system was operating at roughly 56% of its allowed memory with ample headroom. The fix was holding.

The 342 GiB Instance: Deployment Goes Awry

The user then directed the assistant to test on a more constrained machine: "Try on a 256G machine" ([msg 3962]). This was the real test—the smaller instances were the ones that had been OOM-killing before.

The assistant searched for suitable offers on vast.ai ([msg 3963]), filtering for machines with 200–300 GiB of CPU RAM, fast GPUs, and low cost. It found offer 28256188: an RTX 5090 with 251 GiB advertised RAM at $0.33/hr in Canada. The assistant attempted to deploy via the vast-manager API ([msg 3964]), but the command returned empty output.

When the assistant checked the dashboard ([msg 3965]), a new instance had appeared—but it was not the one requested. Instance C.32896749 was an RTX PRO 4000 with 503 GiB RAM in Norway. The assistant's reasoning captures the confusion: "Hmm, the newly deployed instance has cpu_ram_mb: 515604 which is ~503 GiB, not 256 GiB. And the GPU is 'RTX PRO 4000' not RTX 5090" ([msg 3966]).

The investigation deepened when the assistant checked the manager logs via journalctl ([msg 3967]). The log revealed that the manager had deployed offer 31639239, not offer 28256188. The deploy endpoint had silently substituted a different machine. The assistant's reasoning shows the dawning realization: "So the manager actually deployed offer 31639239, not 28256188. It seems the deploy endpoint might have auto-selected a different offer, or there's a mismatch in how I called it" ([msg 3968]).

The assistant made a critical decision: bypass the unreliable manager API and deploy directly via the vastai CLI ([msg 3969]). The command was:

vastai create instance 28256188 --image theuser/curio-cuzk:latest --disk 250 ...

This succeeded immediately, returning instance ID 32897009—an RTX 5090 machine in Canada. When the assistant checked the instance details ([msg 3970]), a new surprise awaited: the instance showed cpu_ram: 503 GiB, not the 251 GiB advertised in the offer.

The assistant's reasoning in [msg 3971] reveals a sophisticated inference: "This makes sense if the offer is for a slice of the machine with a cgroup limit around 251 GiB while the host itself has 503 GiB total—that's actually ideal for testing whether my detection logic properly respects cgroup constraints." Rather than treating the discrepancy as a failure, the assistant recognized it as the perfect test case for the cgroup-aware fix.

Validation on the 342 GiB Instance

The assistant SSH'd into the instance and read the raw cgroup files ([msg 3972]):

cgroup limit: 367003697152 bytes = 341.8 GiB
Host /proc/meminfo: 527977840 kB = 503.4 GiB

The analysis was stark: "Without the fix: would budget 503 GiB - 10 GiB = 493 GiB → OOM kill on a 342 GiB container. With the fix: will budget 342 GiB - 10 GiB = 332 GiB → safe."

The assistant then checked the entrypoint logs ([msg 3973]) and found the memcheck output confirming the fix was working:

[entrypoint] Effective RAM: 341GiB, Budget: 331GiB, Bench concurrency: 4
[entrypoint] GPUs: 1, Can pin: true, cgroup limit: 367003697152
[entrypoint] WARNINGS: cgroup limit 341GiB is 161GiB less...

The summary message at [msg 3974] presented a before-and-after comparison table:

| Metric | Without fix | With fix | |---|---|---| | Detected RAM | 503 GiB (host) | 341 GiB (cgroup) | | Budget | 493 GiB | 331 GiB | | Max PoRep conc. | would try ~30 | 6 | | Result | OOM kill | safe |

The cgroup-aware memory detection was validated on two distinct production instances. The fix worked.

The Crash That Revealed Deeper Problems

Despite the successful validation, the 342 GiB instance later crashed with a broken pipe error during GPU proving ([msg 3983]). The investigation revealed that the CUDA pinned memory pool (PinnedPool) was operating outside the MemoryBudget system—pinned buffers returned to the pool after GPU work were never freed from actual RSS, creating a massive accounting discrepancy. Combined with kernel/driver overhead (glibc arenas, page tables, GPU driver allocations) and the transient SRS loading spike (simultaneous mmap + cudaHostAlloc), the 10 GiB safety margin was empirically insufficient for constrained instances.

This led to the development of two additional tools:

  1. memprobe: A C utility that empirically measures available memory up to the cgroup limit by allocating 1 GiB chunks via mmap/memset until it nears the limit, providing a data-driven safety margin that accounts for hidden kernel overhead.
  2. OOM recovery loop: An enhancement to benchmark.sh that detects when the cuzk daemon is killed by OOM (exit code 137), reduces the budget by 10%, and retries up to three times. Running memprobe on the live 342 GiB instance confirmed the machine was operating at 99% of its cgroup limit with only 14 GiB of additional allocatable space and 6 GiB of kernel/driver overhead—validating the need for these adaptive measures.

Conclusion

The work in this chunk represents a complete arc of production debugging: from identifying a root cause (cgroup-unaware memory detection), implementing a fix in Rust, discovering and fixing additional bugs during live testing, deploying and validating on production instances, and finally building adaptive safety layers when the initial fix proved insufficient for the most constrained environments.

The cgroup-aware memory detection transformed the CuZK proving engine from a system that would inevitably OOM-kill on constrained vast.ai instances to one that correctly respects container limits and operates safely within them. The three bug fixes (GPU JSON parsing, pinning detection, jq resilience) hardened the deployment pipeline against real-world failures. The deployment confusion with the vast-manager API revealed the importance of verifying infrastructure assumptions. And the memprobe utility and OOM recovery loop provided the final layer of adaptive safety that made the system robust on the tightest memory budgets.

The lesson is clear: in containerized environments, the difference between host resources and container-available resources is a constant source of bugs. Fixing the detection layer is necessary but not sufficient—you must also account for kernel overhead, driver allocations, and transient spikes that traditional accounting misses. The combination of correct detection, empirical measurement, and adaptive recovery creates a system that can survive the unpredictable realities of cloud GPU computing.