The Warmup Gambit: How a Two-Phase Daemon Restart Solved OOM in GPU Proving

The Crisis in British Columbia

On a rented GPU instance in British Columbia, Canada, a machine with 125 GB of RAM and two RTX 3090s was being killed by the Linux kernel's Out-Of-Memory (OOM) killer. The instance, designated 32711932, had been tasked with running a PoRep (Proof of Replication) benchmark for the CuZK zero-knowledge proving engine. Instead of producing proofs, it was silently terminated, leaving behind a bench_rate of zero and a mystery for the engineer debugging the system remotely.

The proximate cause was clear enough: the machine ran out of memory. But the why was more subtle. The benchmark script started the CuZK daemon with a partition_workers setting of 10, and when the daemon processed its very first proof, it triggered a Pre-Compiled Constraint Evaluator (PCE) extraction. During PCE extraction, all partition workers ran simultaneously, each allocating substantial memory for constraint synthesis. On a 125 GB machine, this simultaneous allocation exceeded available RAM, and the kernel responded with a SIGKILL.

Meanwhile, a second instance in Norway—a single RTX 4090 with 500 GB of RAM—was running the same benchmark without issue. The contrast was instructive: the problem wasn't the proving workload itself, but the initialization phase, and it only manifested on machines with constrained memory. The fix would need to be surgical: reduce memory pressure during PCE extraction without compromising the performance of the actual benchmark.

The Message: A Strategy Takes Shape

It is in this context that message <msg id=1044> arrives. The assistant writes:

Now let me implement the fix. The strategy: when PCE file doesn't exist, start the daemon with partition_workers = 2 for warmup, then after warmup completes and PCE is cached, kill and restart the daemon with the full PARTITION_WORKERS for the benchmark.

[edit] /tmp/czk/docker/cuzk/benchmark.sh Edit applied successfully.

This is a deceptively simple message. In two sentences, it encapsulates a complete architectural response to a complex systems failure. The assistant has moved from diagnosis to implementation, and the edit it applies to benchmark.sh will become the cornerstone of the entire proving pipeline's reliability on memory-constrained hardware.

The Reasoning Behind the Strategy

The strategy rests on a key insight about CuZK's architecture: PCE extraction is a one-time cost. Once the PCE cache file exists on disk, subsequent daemon restarts load it instantly, bypassing the expensive synthesis step entirely. This means the memory spike is a singular event, limited to the very first proof ever run on a fresh instance.

The assistant's plan exploits this property with a two-phase approach:

Phase 1 (Warmup): If the PCE cache file is absent, start the daemon with partition_workers = 2. This is the minimum viable setting—enough workers to complete PCE extraction, but few enough that the per-worker memory allocation stays within safe bounds. On the BC Canada instance, this would mean roughly 2 workers × ~6 GB per worker ≈ 12 GB for synthesis, rather than 10 workers × ~6 GB ≈ 60 GB. The warmup proof runs slowly with only 2 workers, but that's acceptable because its sole purpose is to generate the PCE cache.

Phase 2 (Benchmark): Once the PCE file exists, kill the daemon and restart it with the full PARTITION_WORKERS count. Now the daemon loads the cached PCE instantly, avoiding the memory spike entirely. The actual benchmark—12 proofs at concurrency 5—runs at full speed with all partition workers available for synthesis.

The choice of partition_workers = 2 specifically is worth examining. Why 2 rather than 1 or 3? The assistant's reasoning, visible in earlier messages, treats 2 as "just enough to generate PCE." PCE extraction is parallelizable across partition workers, but with diminishing returns; a single worker would work but be very slow, while 2 workers provides a modest speedup without doubling memory pressure. The exact number is less critical than the principle: use the minimum number of workers needed to complete the task, then scale up once the memory-intensive initialization is done.

Assumptions and Their Risks

Every engineering decision rests on assumptions, and this strategy is no exception. The assistant implicitly assumes:

That PCE extraction is the sole source of the OOM. This is reasonable given the evidence—the BC Canada instance was killed during warmup, and the Norway instance with 4× the RAM had no issues. But it's possible that other operations (e.g., loading parameters, initializing GPU contexts) also contribute to memory pressure. The fix addresses the most likely cause but doesn't instrument memory usage to confirm.
That the daemon supports graceful kill/restart. The benchmark script kills the daemon by PID and waits for it to terminate before starting a new instance. This works in practice, but there's a risk of race conditions if the daemon doesn't release GPU resources promptly. The assistant mitigates this with a wait command and a cleanup handler.
That the PCE cache path is predictable. The script checks for a specific file path to determine whether PCE exists. If the path changes between daemon versions or configurations, the check would fail silently, potentially causing the daemon to run with reduced workers indefinitely.
That partition_workers = 2 is sufficient for PCE generation. While PCE extraction works with any positive number of workers, extremely constrained instances might still OOM if each worker's per-proof memory allocation is large. The 6 GB per worker estimate is empirical and may vary with proof parameters.

The Follow-Up Optimization

The assistant's first implementation had a subtle inefficiency, caught in message <msg id=1052>: when PCE didn't exist, the daemon was started twice—first with full PARTITION_WORKERS, then immediately killed and restarted with 2 workers. This was harmless (the first start would be killed before doing any real work), but wasteful. The assistant optimized by checking for PCE before the initial daemon start, skipping directly to the reduced-worker configuration when no cache existed. This refinement shows the iterative nature of the fix: the core strategy was sound from the start, but the implementation was tightened through self-review.

The Broader Significance

This message, though brief, represents a turning point in the session. Before the fix, low-RAM instances were unusable; after it, instances with as little as 125 GB of RAM could complete benchmarks reliably. The Czechia instance (2× RTX 3090, 251 GB RAM) and the Belgium instance (2× A40, 2 TB RAM) both deployed successfully with the hardened image, validating the approach.

More broadly, the strategy embodies a principle that recurs throughout systems engineering: separate initialization from steady-state operation, and allocate resources accordingly. The warmup phase uses conservative settings to survive the one-time cost; the benchmark phase uses aggressive settings to maximize throughput. This pattern—ramp up gently, then unleash—is applicable far beyond GPU proving, from database connection pooling to JVM heap sizing to Kubernetes pod startup sequences.

The assistant's decision to encode this logic in benchmark.sh rather than in the daemon itself is also noteworthy. It keeps the daemon simple (it doesn't need to know whether PCE exists) and places the orchestration responsibility in the script that manages the daemon's lifecycle. This separation of concerns is clean and maintainable.

Conclusion

Message <msg id=1044> is a masterclass in targeted intervention. Faced with an OOM crash that threatened to derail the entire proving pipeline, the assistant diagnosed the root cause (simultaneous PCE extraction across too many partition workers), devised a minimal fix (two-phase daemon restart with reduced warmup workers), and implemented it with a single edit. The fix was not a hack—it was a principled response that respected the system's architecture and exploited its properties (the one-time nature of PCE extraction) to achieve reliability without sacrificing performance.

In the end, the BC Canada instance's OOM crash became a catalyst for a more robust, hardware-aware proving pipeline. The warmup gambit worked.