Taming the OOM Killer: How Cgroup-Aware Memory Detection, Empirical Probing, and Adaptive Recovery Made a GPU Proving Engine Survive in the Wild

Introduction

In the high-stakes world of GPU-accelerated zero-knowledge proving, memory is the most unforgiving resource. Run out, and the kernel's out-of-memory (OOM) killer terminates your process without warning, without a crash dump, and without recourse. For the CuZK proving engine—a high-performance GPU-accelerated proof generator for Filecoin—this was a recurring nightmare on cloud GPU instances. The daemon would be running, processing proofs, and then silently die. A broken pipe. A dead process. A benchmark that produced zero proofs before being terminated.

This article synthesizes a sustained engineering effort across two tightly coupled phases of work. In the first phase, the team identified and fixed a fundamental root cause: the Rust function that detected system memory was reading /proc/meminfo, which inside Docker containers reports 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. The fix—rewriting detect_system_memory() to be cgroup-aware—was deployed and validated on two production vast.ai instances.

But validation on a constrained 342 GiB instance revealed a deeper problem. Even with correct cgroup-aware budgeting, the system crashed. The investigation traced the root cause to a CUDA pinned memory pool operating partially outside the memory budget system, combined with invisible kernel/driver overhead that consumed the safety margin. This led to the second phase: building a memprobe C utility that empirically measures available memory up to the cgroup limit, and an OOM recovery loop that adaptively reduces the budget on failure. Running memprobe on a live 256 GiB instance confirmed the machine was operating at 99% of its cgroup limit with zero headroom—validating the desperate need for these adaptive measures.

The arc of this work is a masterclass in production debugging: from identifying a root cause, implementing a fix, 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.

Part I: From Host Blindness to Container Awareness

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 feeds into the MemoryBudget system that controls all memory allocations for GPU pinned buffers, SRS parameters, and proving operations. It 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.

Part II: Living on the Edge—Empirical Measurement and Adaptive Recovery

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 partially 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.

The assistant's initial hypothesis, developed in [msg 3984], pointed to the CUDA pinned memory pool. The reasoning was detailed and compelling: when a partition's synthesis completed and the a/b/c vectors were released, the MemoryBudget system freed their 13 GiB reservation. But the pinned buffers themselves were not freed—they were returned to the PinnedPool for reuse. The budget thought the memory was available; the kernel knew it was still allocated. With eight partitions in the GPU queue, this created a staggering 108 GiB of hidden, untracked memory.

The user corrected this assumption in [msg 3985] with a single, devastating sentence: "We track memory/pinned memory in memory manager." This twelve-word correction dismantled the entire hypothesis and forced a complete re-examination. If pinned memory was tracked, the accounting discrepancy the assistant had spent paragraphs constructing did not exist. The OOM must have a different cause.

This pivot—captured in [msg 3986]—is a masterclass in structured debugging under uncertainty. The assistant did not discard its analysis; it systematically re-traced the entire pipeline flow from first principles, enumerating each step: budget acquisition, vector allocation, GPU queue entry, budget release, pool return. The user's follow-up hint in [msg 3987]—"Maybe in SRS-mode it a bit weird"—redirected the investigation toward the Structured Reference String (SRS) loading path, where a transient double-allocation (simultaneous mmap of the 44 GiB file and cudaHostAlloc of the in-memory representation) created a temporary 88 GiB peak against a budget that only accounted for 44 GiB.

From Analysis to Action: The Memprobe and OOM Recovery

The investigation into SRS loading internals was thorough but ultimately inconclusive. The assistant traced the code through srs_manager.rs, supraseal_params.rs, and into the C++ groth16_srs.cuh file, discovering a conditional compilation flag (H_IS_STD__VECTOR) that affected allocation paths and speculating about point format size mismatches. But rather than getting lost in the internals, the assistant made a critical strategic decision, articulated in [msg 3999]: "Rather than getting lost in the internals, I should focus on the user's practical suggestions: building a memprobe tool to measure actual memory usage and adding OOM recovery to the benchmark."

This decision defines the entire second phase. The assistant chose measurement over speculation, resilience over perfect accounting.

The user's message at [msg 3998] had crystallized the way forward: "The hunch about 10gb being low also sounds right, maybe we should have a program before benchmark which tries to actually allocate memory, and in benchmark runs also recover from OOMs and retry with 50% overhead a few times." From this seed, the assistant designed two complementary mechanisms.

The Memory Probe (memprobe)

The memprobe utility, written as a 133-line C program in [msg 4000], is elegant in its simplicity. Rather than trying to parse kernel memory accounting structures or interrogate /proc filesystem entries, it simply probes the available memory by allocating 1 GiB chunks via mmap with MAP_ANONYMOUS and touching each page via memset until it nears the cgroup limit. It stops 2 GiB before the limit to avoid triggering the OOM killer on itself.

The assistant considered several design alternatives. A bash-based approach using tmpfs was rejected because it would test tmpfs limits, not RSS. The stress and stress-ng tools were absent from the runtime Docker image. Python was unavailable. C, compiled in the Docker builder stage and copied to the runtime image, was the only viable option.

The key limitation was explicitly acknowledged: the probe runs before the daemon starts, measuring only the empty container's memory capacity. Once the daemon is running, additional overhead from page tables for CUDA pinned memory, GPU driver allocations, thread stacks, and glibc arenas would shrink the available headroom. The memprobe provides an optimistic baseline—valuable but not sufficient on its own.

The OOM Recovery Loop

The OOM recovery mechanism, implemented in benchmark.sh ([msg 4009]), is a retry loop wrapped around the benchmark execution. When the cuzk daemon is killed with exit code 137 (128 + SIGKILL, the signal delivered by the OOM killer), the script detects this, reduces the memory budget by 10%, waits 30 seconds for kernel memory reclaim, and retries the benchmark up to three times.

The assistant recognized that this recovery loop is "actually the more reliable safety mechanism" because it tests under real workload conditions. The memprobe provides a baseline; the OOM recovery provides a safety net for everything the baseline misses. Together, they form a defense-in-depth strategy for memory management.

Integration and Deployment

The integration of these two mechanisms required changes across four files, captured in the commit at [msg 4015] and [msg 4016]. The Dockerfile.cuzk was updated to compile memprobe.c in the builder stage and copy the binary to the runtime image (6 insertions). The entrypoint.sh was updated to run memprobe after memcheck and use the more conservative of the two budget estimates (29 insertions). The benchmark.sh was restructured with the OOM recovery loop (319 insertions, significant deletions). And memprobe.c itself was added as a new file (133 insertions).

The commit message, quoted in [msg 4016], reads:

cuzk: add memprobe utility and benchmark OOM recovery

>

- memprobe.c: tiny C program that allocates 1 GiB chunks via mmap+memset to find actual usable memory. Accounts for kernel overhead (page tables, slab, driver) that consumes cgroup budget invisibly. Stops 2 GiB before cgroup limit. - benchmark.sh: wrap benchmark in OOM recovery loop. If daemon is OOM-killed (exit 137), reduce budget by 10% and retry up to 3 times. Waits 30s between retries for kernel memory reclaim. - entrypoint.sh: run memprobe after memcheck, use the more conservative of the two budget estimates. - Dockerfile.cuzk: compile memprobe in builder stage, copy to runtime.

The Docker image was built and pushed to the registry in [msg 4017] and [msg 4018]. But the team faced a practical challenge: the 256 GiB instance (32897009) was already running a benchmark with the old scripts. Restarting the container would interrupt the benchmark and waste hours of parameter download time.

The solution, executed in [msg 4020] and [msg 4021], was a surgical hot-patch. The assistant extracted the memprobe binary from the newly built Docker image using docker create and docker cp, then deployed it to the running instance via scp. This avoided restarting the container while still getting the diagnostic tool onto the live system.

The 99% Machine: Empirical Validation

The moment of truth arrived in [msg 4022]. The assistant deployed memprobe to the live 256 GiB instance and captured the output:

{
  "cgroup_limit_bytes": 367003697152,
  "usage_before_bytes": 342817820672,
  "allocated_bytes": 15032385536,
  "allocated_gib": 14,
  "usage_peak_bytes": 364346327040,
  "usage_after_bytes": 350478913536,
  "overhead_bytes": 6496120832,
  "overhead_gib": 6.0,
  "probe_chunks": 14
}

The data was stark. The instance was operating at 99% of its cgroup limit (340 GiB / 342 GiB) before the probe even started. Only 14 GiB of additional memory could be allocated. Six gigabytes of kernel/driver overhead—page tables, slab allocator caches, GPU driver allocations, glibc arena metadata—were consuming cgroup budget invisibly. The 10 GiB safety margin was being eaten alive by overhead that standard tools could not detect.

The assistant's commentary in [msg 4023] captured the tension: "The machine is running at 99% of its cgroup limit during benchmark." This was not a system that had failed—it was a system that was surviving, but on a knife's edge.

The follow-up check in [msg 4024] confirmed the daemon was still alive, with RSS at 340 GiB against a 342 GiB limit and a budget of 331 GiB. The daemon log showed a clean initialization: "memory budget initialized total_budget_gib=331," "CUDA pinned memory pool initialized," "starting pipeline." No errors, no warnings. The system was running, but the 9 GiB gap between the budget (331 GiB) and actual RSS (340 GiB) was the PinnedPool accounting discrepancy in action—visible, measurable, and dangerous.

In [msg 4025], the assistant drilled deeper, querying both the cgroup memory controller and the daemon's internal status API. The cgroup reported memory.current at 320.6 GiB against a memory.max of 342 GiB. The daemon's status API showed used_bytes at 321.8 GiB with only 9.2 GiB available. Both GPU workers were in "proving" state—the pipeline was fully saturated. The assistant's assessment: "The machine is currently surviving! RSS is 340/342 GiB which is extremely tight but hasn't OOM'd yet."

Discussion: Assumptions, Limitations, and Lessons Learned

The memprobe and OOM recovery approach rests on several assumptions that deserve examination.

The memprobe measures anonymous mmap, not CUDA pinned memory. The probe allocates via mmap with MAP_ANONYMOUS, which is a reasonable proxy for general memory pressure but does not capture the specific overhead of CUDA pinned allocations or GPU driver state. The 6 GiB overhead figure is a lower bound—the real overhead during GPU proving could be higher.

Exit code 137 reliably indicates OOM. In practice, exit code 137 is 128 + 9 (SIGKILL), which the OOM killer delivers. But any SIGKILL produces this exit code, including manual kill -9 or systemd unit timeouts. The recovery loop could retry on non-OOM failures, potentially masking other bugs.

The 10% budget reduction per retry is a heuristic. There is no theoretical justification for 10% rather than 5% or 20%. On a 342 GiB machine, three 10% reductions give a maximum adjustment of ~27%. This is gentle enough to avoid collapsing performance on the first retry but might be insufficient if the overshoot is large.

The 30-second wait between retries is a guess. Kernel memory reclaim after an OOM kill can complete in seconds on some systems and take minutes on others, especially with large pinned allocations. A fixed 30-second wait is a compromise.

The approach addresses symptoms, not the root cause. The memprobe and OOM recovery loop do not fix the PinnedPool accounting issue. They build a safety net around it. The pool still operates partially outside the budget system; the RSS still diverges from the budget. A future workload change could invalidate the memprobe's measurements.

Conclusion

The work in this segment represents a fundamental shift in how the CuZK proving engine approaches memory management. The team moved from a model of perfect accounting—where the budget should exactly match reality—to a model of empirical measurement and adaptive resilience.

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.

The empirical validation on the live 256 GiB instance confirmed the severity of the problem: 6 GiB of invisible kernel overhead, 99% cgroup utilization, zero headroom. The system was surviving, but barely. The memprobe and OOM recovery loop gave it a fighting chance.

In the broader narrative of the CuZK project, this segment represents the moment when the system learned to survive in the wild—on memory-constrained GPU instances where every gigabyte counts and the kernel is always watching. The commit that codified this survival ([msg 4016]) is not just a set of code changes; it is the crystallized artifact of a debugging odyssey, encoding hard-won knowledge about kernel memory accounting, cgroup limits, and the gap between what a budget thinks is allocated and what the kernel knows is real.

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.

References

[1] "From Host Blindness to Container Awareness: How Cgroup-Aware Memory Detection Rescued a GPU Proving Engine from OOM Kills" — Detailed analysis of the detect_system_memory() rewrite, memcheck bug fixes, and production deployment validation.

[2] "Living on the Edge: How Empirical Measurement and OOM Recovery Tamed Memory Death in a GPU Proving Pipeline" — Analysis of the OOM crash investigation, memprobe utility, and OOM recovery loop implementation and validation.