The Budget That Replaced the Semaphore: Reading the PoRep Partition Dispatch Loop

In the middle of a sweeping refactoring that replaced a static concurrency limiter with a memory-aware admission control system, the assistant paused to read a critical section of code. Message <msg id=2185> is deceptively simple on its surface: a single read operation targeting lines 1457–1464 of engine.rs, the central coordinator of the cuzk GPU proving daemon. But this read was the fulcrum on which an entire architectural transformation turned — the moment when the old world of fixed worker counts gave way to a new world of byte-level budget accounting.

The Context: Replacing the Semaphore

To understand why this message matters, one must understand what was being replaced. The cuzk proving engine had long used a partition_semaphore — a tokio semaphore initialized to a fixed count called partition_workers. This semaphore acted as a crude admission control mechanism: at most N partitions could be in-flight simultaneously, where N was a configuration value tuned by trial and error. If a machine had enough RAM for 3 simultaneous 32 GiB PoRep partitions, you set partition_workers = 3. If the machine changed, you changed the config. If a new proof type with different memory requirements was added, you crossed your fingers.

The new design, specified in the cuzk-memory-manager.md document and partially implemented across five files before this message, replaced that static semaphore with a MemoryBudget — a unified, byte-level budget auto-detected from system RAM. Instead of counting workers, the system would now track actual memory consumption: SRS (44 GiB pinned), PCE (26 GiB heap), and per-partition working memory (13.6 GiB for a 32 GiB PoRep partition). A partition could only proceed if the budget had enough remaining bytes to accommodate it.

By message <msg id=2185>, the assistant had already:

• Created memory.rs with MemoryBudget, MemoryReservation, and all estimation constants
• Updated config.rs with new budget fields and deprecation warnings for old ones
• Rewritten SrsManager for budget-aware loading with LRU eviction
• Replaced static OnceLock PCE caches with a PceCache struct in pipeline.rs
• Partially modified engine.rs to add reservation: Option<MemoryReservation> to SynthesizedJob, and wire the budget into the Engine struct But the core of the engine — the start() method, the dispatcher loop, the partition dispatch paths, the GPU worker loop — still used the old semaphore. That was the work remaining in this chunk.

The Specific Read: Lines 1457–1464

The message reads:

### Edit 10: Update PoRep partition dispatch loop — replace semaphore with budget.acquire()
[read] /tmp/czk/extern/cuzk/cuzk-core/src/engine.rs
<path>/tmp/czk/extern/cuzk/cuzk-core/src/engine.rs</path>
<type>file</type>
<content>1457:                                 partition_idx,
1458:                                 job_id: job_id.clone(),
1459:                                 request: req.clone(),
1460:                                 params: srs.clone(),
1461:                                 circuit_id: CircuitId::Porep32G,
1462:                                 num_partitions,
1463:                             };
1464:              ...

This is a read of the SynthesizedJob construction site inside the PoRep partition dispatch loop. The assistant is looking at the exact point where a per-partition job is created, just before it is spawned as a tokio task. The code shows the fields being set: partition_idx, job_id, request, params (the SRS parameters), circuit_id (set to CircuitId::Porep32G), and num_partitions.

What's notable is what's missing from this struct literal: there is no reservation field. The assistant had earlier added reservation: Option<MemoryReservation> to the SynthesizedJob struct definition, but the construction sites throughout the file still built the struct without it. Every single place that created a SynthesizedJob needed to be updated to either pass reservation: Some(reservation) (for budget-gated paths) or reservation: None (for paths that didn't acquire budget).

This read was the reconnaissance step before the surgical edit. The assistant needed to see the exact structure of the code around the spawn point to understand:

1. Where the semaphore acquisition happened (the line that would be replaced)
2. How the SynthesizedJob was constructed (to add the reservation field)
3. What variables were in scope (budget, pce_cache, srs_mgr)
4. How the spawn and error handling were structured (to preserve the control flow)

The Reasoning Behind the Read

The assistant's thinking, visible in the sequence of edits across this chunk, reveals a methodical, top-down approach to the refactoring. Earlier edits had already:

• Edit 1: Replaced the SRS preload block and PCE preload block in start() with evictor wiring
• Edit 2: Updated channel capacity sizing to remove partition_workers dependency
• Edit 3: Removed partition_workers and partition_semaphore setup
• Edit 4-7: Updated the dispatcher spawn block and all dispatch_batch signatures and call sites
• Edit 8-9: Updated process_batch signature and the PoRep PCE background extraction Now, with Edit 10, the assistant was at the innermost loop — the per-partition dispatch that actually spawns the synthesis+prove tasks. This was the heart of the change: instead of acquiring a semaphore permit (which only limited count of concurrent partitions), the code would acquire a memory reservation (which limited bytes of concurrent memory usage). The assistant's decision to read before editing reflects a disciplined approach to large-scale refactoring. The PoRep partition dispatch loop (spanning roughly lines 1317–1568 in the original) was a complex section with multiple concerns: SRS loading, PCE extraction, partition spawning, error handling, and status reporting. Changing the admission control mechanism required understanding how all these pieces fit together, because the budget acquisition needed to happen at the right point — early enough to prevent memory oversubscription, but late enough that the budget could be accurately estimated.

Assumptions Made

The assistant made several assumptions in this read-and-edit sequence:

1. That the SynthesizedJob struct already had the reservation field. This was a safe assumption because the assistant had added it in an earlier edit (in a previous chunk). The read confirmed that the construction site didn't include it yet, which was expected — the assistant was systematically visiting each construction site.
2. That the budget acquisition pattern would mirror the semaphore acquisition pattern. The old code used partition_sem.acquire_owned().await at the top of the spawn block. The new code would use budget.acquire(POREP_PARTITION_FULL_BYTES).await in the same position. This assumption proved correct because both mechanisms served the same purpose — gating admission — even though they measured different resources.
3. That POREP_PARTITION_FULL_BYTES was the correct constant. This constant (defined in memory.rs as approximately 13.6 GiB for a 32 GiB PoRep partition) represented the full working memory needed for one partition: the a/b/c vectors (~12.5 GiB), auxiliary data (~0.74 GiB), and density wire (~48 MB). The assistant assumed this was the right amount to reserve, which it was — the spec had carefully calculated these values.
4. That the reservation would be attached to the SynthesizedJob and released in two phases. This was the key architectural insight from the memory manager spec: the a/b/c portion (~12.5 GiB) could be released immediately after gpu_prove_start, while the remaining ~1.1 GiB would be held until gpu_prove_finish. The reservation needed to travel with the job through the pipeline.

Input Knowledge Required

To understand this message, one needs knowledge of:

• The cuzk proving pipeline: PoRep (Proof-of-Replication) proofs are computed in multiple partitions, each of which goes through CPU synthesis followed by GPU proving. Partitions can be dispatched in parallel, but memory is the limiting factor.
• The old semaphore-based architecture: partition_workers was a config value that set the capacity of a tokio semaphore. Each partition would acquire_owned() a permit before spawning, and release it when done. This limited the number of concurrent partitions but not their memory usage.
• The new budget-based architecture: MemoryBudget tracks total memory consumption in bytes. budget.acquire(bytes) returns a MemoryReservation that holds the budget quota. The reservation can be partially released (two-phase release) or converted to permanent ownership (into_permanent()).
• The SynthesizedJob struct: A per-partition job container that carries the request, SRS parameters, circuit ID, and (after this refactoring) a memory reservation through the synthesis and proving pipeline.
• The memory constants: POREP_PARTITION_FULL_BYTES (~13.6 GiB), POREP_ABC_BYTES (~12.5 GiB), and the two-phase release pattern.

Output Knowledge Created

This message, combined with the edit that immediately followed it ([msg 2186]), created:

1. A budget-gated PoRep partition dispatch: Instead of acquiring a semaphore permit, the code now calls budget.acquire(POREP_PARTITION_FULL_BYTES) before spawning each partition task. If the budget is exhausted, the acquire() future waits until memory is freed (by another partition completing its GPU prove phase).
2. A reservation-attached SynthesizedJob: The reservation returned by budget.acquire() is stored in the SynthesizedJob.reservation field, carrying it through the pipeline so it can be released in two phases.
3. A pattern for the remaining dispatch paths: The SnapDeals partition dispatch (updated in Edit 11, <msg id=2188-2191>) follows the same pattern, using SNAP_PARTITION_FULL_BYTES instead. The monolithic synthesis path (updated later) also acquires budget, though with a different byte estimate.

Broader Significance

Message [msg 2185] represents a microcosm of the entire memory manager refactoring. The read operation — examining a few lines of struct construction — was the moment where the abstract design (budgets, reservations, two-phase release) met the concrete reality of the codebase. The assistant wasn't just replacing one function call with another; it was changing the fundamental admission model of the proving engine.

The old semaphore answered the question "how many partitions can run concurrently?" with a static integer. The new budget answers the question "how much memory can we use?" with a dynamic, system-aware calculation. This is the difference between counting tokens and measuring resources — a shift from coarse-grained to fine-grained control, from configuration-dependent to environment-adaptive.

The read in message [msg 2185] was the necessary precondition for that transformation. Without seeing the exact structure of the spawn point — the variables in scope, the fields of SynthesizedJob, the error handling around the spawn — the assistant could not have written the correct edit. The read is invisible in the final code (the edit is what persists), but it was the essential act of understanding that made the edit possible.