Order from Chaos: Replacing Race-Condition Partition Scheduling with FIFO Dispatch in CuZK

In any high-performance computing pipeline, scheduling discipline is the difference between throughput and thrashing. When partitions of a proof job are processed in random order, the entire system suffers from fragmentation: nearly-complete pipelines stall waiting for GPU proving while other pipelines with no synthesis work left consume budget, creating gaps that waste capacity and inflate end-to-end latency. This was precisely the problem facing the CuZK proving engine in message [msg 2763], where the assistant diagnosed a fundamental scheduling flaw and designed a clean fix using an ordered message-passing channel.

The Problem: A Thundering Herd on Budget

The root cause was architectural. CuZK's proof pipeline for both PoRep and SnapDeals proof types dispatched partitions using a straightforward tokio::spawn loop:

for partition_idx in 0..num_partitions {
    tokio::spawn(async move {
        let reservation = budget.acquire(PARTITION_BYTES).await;
        // synthesize...
    });
}

Every partition from every pipeline was spawned as an independent asynchronous task. These tasks then raced on budget.acquire(), a memory reservation mechanism backed by tokio's Notify primitive. When memory was released (e.g., a GPU prove completed), Notify woke all waiting tasks simultaneously, creating a thundering herd. Each task would attempt a fetch_add atomic operation, and only one would succeed. The others would fall back to sleep and wait for the next release.

The result was a random scheduling order. As the user reported in [msg 2747] with an annotated screenshot, pipeline 3 (ps-snap-3644168) had 6 of 16 partitions done while pipeline 4 (ps-snap-3644166) had 0 of 16 done but had partitions P3, P6, P7, and P15 scattered across the status display. Pipeline 5 had only 1 of 16 done. The scheduler was processing partitions in no discernible order — whichever task happened to win the race got to run next.

This was not merely an aesthetic concern. The lack of ordering meant that nearly-finished pipelines could stall waiting for GPU proving while other pipelines had no synthesis work left to do. A pipeline needing only its last two partitions to complete would have to wait for those partitions to win the race against partitions from entirely different jobs. The system had no concept of "finish what you started."

The Assistant's Reasoning Process

The assistant's thinking in [msg 2763] reflects a mature engineering approach: understand the constraint, evaluate alternatives, and choose the minimal structural change that solves the problem without over-engineering.

The message begins with a confident declaration: "OK I have a clear picture now." This statement is earned — it follows a multi-step investigation across several earlier messages. In [msg 2750], the assistant had examined the screenshot and identified the three visible symptoms. In [msg 2753], it traced the root cause to the tokio::spawn + budget.acquire() pattern and considered several approaches:

1. Priority-aware budget.acquire() — adding a priority parameter so that tasks with lower (job_index, partition_index) win. This was rejected because it required changing the MemoryBudget implementation, a deeper and more invasive change.
2. Sequential await within a job — awaiting each partition's budget acquisition sequentially within a job, allowing only synthesis_concurrency jobs to run their loops in parallel. The assistant quickly realized this would serialize partitions within a job, defeating the purpose of overlapping synthesis with GPU proving.
3. Ordered channel dispatch — replacing the spawn-all pattern with a shared mpsc channel where partitions are enqueued in FIFO order and a pool of synthesis workers pulls from the channel sequentially. This was chosen as the cleanest minimal fix. By [msg 2757], the assistant had settled on the channel approach and outlined the design. Message [msg 2763] formalizes this into a concrete four-point plan and begins the implementation by reading the current code at lines 1507–1511 to confirm the exact pattern that needs replacing.

The Four-Point Plan

The assistant's plan in [msg 2763] is concise and precise:

1. Create a partition_work_tx/rx channel — an unbounded or large-bounded mpsc channel. The items (PartitionWorkItem structs) are lightweight, consisting of Arc pointers to parsed proof input, a partition index, a job ID, and a proof request clone. The channel is unbounded because the items are small and the real bottleneck is memory budget, not channel capacity.
2. Add a pool of synthesis worker tasks — these pull from partition_work_rx in FIFO order (guaranteed by mpsc), acquire budget, synthesize on a blocking thread, and push the result to synth_tx (the GPU channel). The worker count is set to max_partitions_in_budget, ensuring that budget — not worker availability — is the real bottleneck.
3. Replace the spawn loop in process_batch — instead of for partition_idx ... tokio::spawn, the code becomes for partition_idx ... partition_work_tx.send(item). This is a surgical change: the partition creation logic stays the same, but the dispatch mechanism changes from "fire and forget" to "enqueue in order."
4. Set worker count to max_partitions_in_budget — this ensures the worker pool is large enough that no partition is ever starved of a worker when budget is available, while keeping the budget as the true limiting factor.

Assumptions and Trade-offs

The assistant makes several assumptions in this design. First, it assumes that mpsc channel ordering is sufficient to guarantee FIFO processing. This is correct for tokio's mpsc — it is a single-producer, multi-consumer channel that preserves send order. However, the assistant is using a multi-producer variant (each pipeline's process_batch sends items), which means the ordering guarantee applies per-sender. If two pipelines send simultaneously, items may interleave at the receiver. The assistant addresses this by having each pipeline enqueue all its partitions in a single sequential loop before returning, which effectively batches each pipeline's partitions together in the channel.

Second, the assistant assumes that budget acquisition within the worker is the natural pacing mechanism. Workers pull a partition from the channel, acquire budget (blocking if insufficient), synthesize, and send to GPU. If budget is exhausted, workers block on acquire() rather than consuming channel items, which naturally stalls the channel without dropping items. This is sound.

Third, the assistant assumes that the worker pool size (max_partitions_in_budget) is sufficient. With 400 GiB total budget and 9 GiB per SnapDeals partition, this yields approximately 44 workers. This is large enough that budget, not worker count, will be the binding constraint — exactly the desired behavior.

Input Knowledge Required

To understand this message, the reader needs familiarity with:

• Tokio async Rust: the tokio::spawn pattern, mpsc channels, and Notify-based synchronization
• CuZK's proof pipeline: the two-stage architecture (CPU synthesis → GPU proving), the PartitionWorkItem struct, and the budget memory reservation system
• SnapDeals and PoRep proof types: both use the same partition dispatch pattern, and the fix must apply to both
• The memory budget manager: how acquire() works with Notify and the evictor callback

Output Knowledge Created

This message creates the design specification for a critical scheduling fix. It does not yet implement the code — that happens in subsequent messages — but it establishes:

• The architectural change (channel-based dispatch replacing spawn-all)
• The worker pool sizing rule (budget-derived, not config-derived)
• The integration points (both PoRep and SnapDeals dispatch loops in process_batch)
• The invariants to preserve (budget gating, GPU channel backpressure)

Why This Message Matters

Message [msg 2763] is the hinge point in a debugging session that had been chasing a confusing symptom — the /4 display bug in the status API — and discovered a deeper systemic issue. The random partition scheduling was not a display bug; it was a performance bug that could cause pipeline completion times to vary wildly depending on which partitions won the race. The assistant's decision to prioritize this fix over the display issue reflects sound judgment: the scheduling problem directly affects throughput, while the display problem is cosmetic.

The message also demonstrates a key skill in systems programming: recognizing when a "random" behavior is not actually random but is the deterministic result of an underspecified scheduling policy. The partitions were not being scheduled randomly — they were being scheduled by tokio's work-stealing scheduler, which has no concept of application-level priority. The fix does not fight the scheduler; it removes the scheduling decision from tokio entirely by using an ordered channel that preserves application-level ordering.

In the broader context of the CuZK proving engine, this fix transforms the pipeline from a chaotic free-for-all into a disciplined FIFO queue. Earlier pipelines get their partitions processed first. Within a pipeline, lower partition indices are processed first. This means a pipeline that is 14 of 16 partitions complete will have its last two partitions processed before a new pipeline's first partition, minimizing the gap between synthesis completion and GPU proving. The result is lower tail latency, more predictable completion times, and better utilization of GPU resources.