The Edit That Completed the Trilogy: Integrating memprobe into entrypoint.sh

The message is deceptively brief:

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

Seven words. A routine confirmation that a file was patched. But in the arc of this debugging session — a multi-hour investigation into why a GPU proving engine was being killed by the Linux out-of-memory (OOM) killer on memory-constrained vast.ai cloud instances — this single edit represents the culmination of a three-pronged strategy to solve a problem that had been plaguing the deployment for days. Message [msg 4012] is the moment the last piece clicked into place.

The Context: A Machine Living on the Edge

To understand why this edit matters, one must understand the crisis that preceded it. The CuZK proving engine, a high-performance GPU-based zero-knowledge proof system for Filecoin, had been crashing with OOM kills on vast.ai instances with tight memory budgets. The root cause was a cascade of problems: the CUDA pinned memory pool (PinnedPool) was operating outside the MemoryBudget accounting system, returning buffers to a pool after GPU work without ever freeing them from actual RSS. Combined with kernel/driver overhead — glibc malloc arenas consuming ~1.1 GiB across 18 threads, page tables for 94 GiB of pinned allocations, GPU driver allocations, and a transient spike during SRS parameter loading where a 44 GiB mmap coexisted with a 44 GiB cudaHostAlloc — the fixed 10 GiB safety margin was empirically insufficient.

The user had articulated the solution in [msg 3997] and [msg 3998]: "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." This was the seed of a three-part strategy.

The Three-Pronged Approach

The assistant had already executed two of the three prongs before [msg 4012]. In [msg 4000], it wrote memprobe.c — a 133-line C program that allocates 1 GiB chunks via mmap and memset until it nears the cgroup limit, then reports the actual allocatable memory and the kernel/driver overhead. In [msg 4003]-[msg 4004], it updated Dockerfile.cuzk to compile memprobe in the builder stage and copy it to the runtime image. In [msg 4009], it rewrote benchmark.sh with an OOM recovery loop: if the daemon is killed (exit code 137, the OOM signal), the budget is reduced by 10% and the benchmark retried up to three times.

What remained was the integration layer. The entrypoint.sh script is the container's initialization sequence — it runs memcheck.sh to detect cgroup-aware memory limits, calculates a budget, and then launches either a benchmark or the production daemon. The edit in [msg 4012] modified this script to run memprobe after memcheck and use the more conservative of the two budget estimates. This was the crucial wiring that connected the empirical measurement tool to the actual runtime decision.

The Reasoning Behind the Integration

The assistant's thinking, visible in the extended reasoning of [msg 3999], reveals a sophisticated understanding of the problem. It calculated the overheads precisely: "glibc malloc arenas: each thread gets a 64 MiB arena → 18 threads × 64 MiB = ~1.1 GiB. Thread stacks: 18 × 8 MiB = 144 MiB. Kernel slab for the container: varies but could be 1-2 GiB. GPU driver overhead: unclear, but could be significant. So total overhead might be 3-5 GiB beyond our 10 GiB safety margin."

But it also recognized the limitations of memprobe: "The memory probe approach is simpler in theory... but there's a timing issue: the probe runs before the daemon starts, so it only measures the empty container. Once the daemon is running, there's additional overhead from page tables for CUDA pinned memory and GPU driver allocations that weren't accounted for." This is why the integration used the more conservative estimate — acknowledging that the probe gives an optimistic baseline that will shrink under real workload.

The edit's logic was straightforward: memcheck.sh reads cgroup v2 (memory.max) and v1 (memory.limit_in_bytes) limits and returns the minimum of host RAM and the cgroup constraint, then subtracts a safety margin. memprobe empirically measures how much memory can actually be allocated before hitting the cgroup limit, accounting for kernel overhead that memcheck cannot see. By taking the more conservative of the two, the system would never use more than the cgroup allows, and would additionally account for hidden overhead that the cgroup limit doesn't directly expose.

What the Edit Actually Changed

The diff, visible in the commit at [msg 4016], shows 29 lines changed in entrypoint.sh. The edit inserted a memprobe invocation after the memcheck block, parsed its JSON output, extracted the overhead_gib field (which reports kernel/driver overhead invisible to cgroup accounting), and used it to compute a tighter safety margin. The budget passed to benchmark.sh and run.sh was then the minimum of the memcheck-derived budget and the memprobe-informed budget.

This was not a large change in terms of lines of code, but it was architecturally significant. Previously, the safety margin was a hardcoded guess — 10 GiB, which the user correctly suspected was too low. After the edit, the safety margin became a data-driven value, computed fresh on every container start by actually trying to allocate memory and observing where the system pushes back.

The Validation: Empirical Proof

The edit's value was validated almost immediately. In [msg 4022], the assistant deployed memprobe to a live 256 GiB instance (vast.ai instance 32897009) that was currently running a benchmark. The output was stark:

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

The machine was operating at 342 GiB out of a 342 GiB cgroup limit — 99% utilization — with only 14 GiB of additional allocatable space and 6 GiB of kernel/driver overhead. The instance was surviving the benchmark but with zero headroom. Any transient spike would trigger an OOM kill. The 10 GiB safety margin that had been the default was, as the user suspected, barely sufficient for this machine and would fail entirely on smaller instances.

Assumptions and Their Validity

The edit rested on several assumptions. First, that memprobe's measurement of an empty container is a useful proxy for the memory available under load. This is partially true — the kernel overhead (page tables, slab caches, driver allocations) is largely static and independent of the workload, so measuring it before the daemon starts gives a floor. But the pinned memory pool's accounting discrepancy — buffers returned to the pool but never freed from RSS — is a workload-dependent effect that memprobe cannot capture. The assistant acknowledged this and relied on the OOM recovery loop in benchmark.sh as the safety net for dynamic effects.

Second, that the more conservative estimate is always the safer choice. This is true for correctness but potentially wasteful on machines where memprobe's empty-container measurement is significantly lower than the actual usable memory under load. The assistant judged this acceptable — better to under-utilize memory than to OOM.

Third, that the memprobe binary would be available at runtime. The Dockerfile edits in [msg 4003]-[msg 4004] ensured this by compiling memprobe in the builder stage and copying it to the runtime image, but the edit in [msg 4012] did not include a fallback if memprobe is missing or fails. The entrypoint would silently fall through to the memcheck-only budget, which is safe but loses the benefit.

The Deeper Significance

Message [msg 4012] is a study in how small edits can have outsized impact when they complete a larger logical chain. The three parts — memprobe, OOM recovery, and entrypoint integration — were designed as a system. The memprobe provides a proactive, empirical safety margin. The OOM recovery provides a reactive, adaptive safety net. The entrypoint integration wires them together so that every container start benefits from both.

The edit also represents a philosophical shift: from guessing to measuring. The original 10 GiB safety margin was a guess, educated but ultimately arbitrary. The memprobe-based approach replaces guessing with empirical measurement, running a real allocation experiment on the actual hardware under the actual cgroup constraints. This is the engineering equivalent of "measure, don't guess" — a principle that this edit operationalized in a single, seven-word confirmation.

In the broader narrative of the CuZK deployment, [msg 4012] is the moment the memory management strategy became self-correcting. Before the edit, OOM kills required manual intervention — someone had to notice the crash, investigate, adjust the budget, and redeploy. After the edit, the system measures its own constraints, adapts its budget, and if it still gets it wrong, retries with a smaller budget automatically. The edit transformed a fragile, manually-tuned deployment into a resilient, self-tuning one.