fix: remove idle waits from checkpoint batching

[yaythomas]

Fixes #710

Problem

After collecting a synchronous checkpoint, the collector waited up to 100 ms on an empty queue. The caller waits until the batch persists, so it cannot add work. The wait only delayed it. Sequential steps paid it once per step.

Change

1. Once a batch holds a synchronous checkpoint, the collector waits 1 ms on an empty queue instead of 100 ms. A batch with no blocked caller keeps the full window.
2. A coordinator resuming timed waits needs a refresh, an empty checkpoint whose response shows the waits complete. It now requests the refresh when the branch suspends, with the resume time attached, through ExecutionState.schedule_refresh. The refresh goes into a heap under the checkpoint lock, and the collector sends every refresh due at one time in one request. Without this, change 1 would split refreshes from independent coordinators that arrive more than 1 ms apart.

The final due check seals a batch. A refresh scheduled before it joins that batch. Refreshes with different check times make one request per time. Sealing a batch never depends on how many ordinary operations the queue holds.

In the executor, a ResumeWave holds the branches suspended until one time and their refresh. Branches resume in index order.

Measurements

Before is main at 8742ad9. After is this branch. JavaScript and Java are the current SDKs, each with a handler of the same shape.

Local test runners, 30 ms simulated round trip, time measured inside the handler, three runs each.

	Python before	Python after	JavaScript	Java
200 sequential steps, checkpoint calls	200	200	200	398 to 399
200 sequential steps, handler time	26.2 s	6.34 s	6.32 s	12.2 s
200-branch fan-out, checkpoint calls	5	5	4	6
200-branch fan-out, handler time	413 ms	217 ms	136 ms	320 ms

Lambda, us-east-1, 1024 MB, warm, one run each.

	Python before	Python after	JavaScript	Java
1000 sequential steps, median interval	139 ms	41 ms	36 ms	38 ms
1000 sequential steps, p95 interval	161 ms	62 ms	58 ms	83 ms
1000 sequential steps, invocation	142.0 s	43.5 s	39.1 s	45.4 s
200-branch fan-out, invocation	855 ms	644 ms	458 ms	937 ms

Java has no batching window. It flushes each update at once on a separate thread, so a step's START and SUCCEED share a request only when the customer thread submits the SUCCEED before the flusher reads the queue. On Lambda at 1024 MB that happened for 841 of 1000 steps. On a many-core host locally it happened for none, which is the two requests per step in the table above. JavaScript flushes at the end of each event-loop turn, so a step whose body does not await makes one request.

Nested resume waves on Lambda, an outer parallel of 10 inner parallels, each with an 8 s step and a 2 s wait. Resume latency is the delay from a wait's scheduled end to its branch running again. Before, three runs, median 292 to 307 ms, 9 checkpoint calls. After, six runs, median 90 to 132 ms, 10 calls in four runs and 13 in two. The spread comes from the service. Some refresh round trips took 65 ms, and in two runs the service assigned two check times 36 to 69 ms apart to the ten waits, so two refresh requests.

The change does not alter operation counts, so it does not affect quota L-42D0A120.

API

schedule_refresh and ScheduledRefresh are new on ExecutionState, which the package does not export and whose docstring now says so. The executor is their only caller. Nothing else in the public surface changes.

Tests

Collector tests assert the timeout passed to each queue read, with no sleeps. Refresh tests cover deferral, grouping by check time, cancellation, wake coalescing, and settlement on failure, completion and shutdown. Four tests come from review findings. Cancelled refreshes do not hold a batch open. A refresh scheduled while the collector fills a batch joins that batch. A backlog of 5,000 queued operations does not change what one collection reads. Branches suspended until the same time resume in index order. Each of the four fails against the defect it guards.

The e2e test enqueues 300 refreshes before starting the batcher and asserts one request. It passed 10 of 10 runs under an 8-way CPU load. ci-checks.sh exits 0 with 1700 core tests and 98.17% coverage.

Gaps in the measurements

I ran the Lambda sequential and fan-out scenarios once per function, so those figures carry run-to-run noise. I ran the nested waves six times after and three before. I measured only at 1024 MB, at most 500 branches, and at most 50 nested coordinators. The JavaScript and Java Lambda functions report no checkpoint call counts, so those cells are local only.

[ParidelPooya]

Three pre-existing wall-clock assertions in state_test.py lost their slack. Each test sleeps 0.15 s in a processor thread, then collects a batch, and asserts the caller blocked at least 0.15 s. The caller records its own start time after the sleep began, so the measured time is 0.15 s minus the thread stagger plus the collection time. On main the collection
added ~100 ms, which covered the stagger. Now it adds ~1 ms.

The three tests are at lines 3575, 3647 and 4184. You already fixed a fourth test of this exact shape (test_create_checkpoint_multiple_concurrent_callers) by replacing the wall-clock assertion with event-based ones. Do the same here.

[yaythomas]

thanks! Done. The three tests now use the event-based shape from test_create_checkpoint_multiple_sync_calls_all_block, through one shared helper. A fourth test of the same shape, test_synchronous_checkpoint_blocks_until_complete, had the same assertion and is converted too. The helper fails when the call is made with is_sync=False, so it still tests blocking.

[github-actions]

Codex AI review · Finding arf_v1_vxz4y425f47cdgl6c7z2tvac2i

[P2] Preserve the coalescing window for synchronous empty checkpoints. Every due ScheduledRefresh becomes QueuedOperation(None, completion_event), so this marks the batch as blocked and _read_timeout seals it after 1 ms. With the collector running, same-time immediate/past refreshes requested by independent coordinators more than 1 ms apart are therefore split into multiple empty API calls, potentially producing O(coordinators) latency and throttling. The updated e2e masks this by enqueueing every handle before starting the batcher. Track synchronous empty checkpoints separately and retain the normal coalescing window whenever one is present, including mixed batches; add live-batcher coverage with same-time due refreshes arriving more than 1 ms apart.

[yaythomas]

By design. Restoring the full window for every synchronous empty checkpoint added about 100 ms to every normal refresh wave, undermining this PR’s goal. The PR also explicitly tests that a refresh requested after its time may require another request.

The worker collects a batch, then sends it, then collects again. While a request is in flight, every arrival pools into the next batch. So a burst of N past-due refreshes spread over less than one round trip, about 35 ms on Lambda, makes two requests, not N. The first refresh goes alone, the other N minus 1 share the next request, and every coordinator resumes sooner than under main's 100 ms wait. Arrivals spaced more than a round trip apart make more requests, but main then makes one request per 100 ms, so the ratio is bounded by about 100 ms over one round trip, roughly three. That is the trade rule A already makes for synchronous step checkpoints from parallel branches. The fan-out measurement shows it is a good trade. 200 branches made 5 requests before and 5 after, in 855 ms and 644 ms.

Keeping the 100 ms window whenever a synchronous empty checkpoint is present puts the idle wait back on the common path. Every wave resume is a batch of exactly such checkpoints. The heap already groups them by check time without any window, so the window would buy nothing and cost about 100 ms per wave. Nested resume latency went from about 290 ms to about 90 ms. This would move it back to about 190 ms.

Propose to reconsider as an addition later in separate PR if necessary.

[github-actions]

Codex AI review

One P2 batching regression remains: live concurrent due refreshes can still fragment into multiple API calls. The updated e2e covers only prequeued refreshes.

Reviewed commit 6a22723140df024f16b0e4afbe8eccd0a0033b7e. Workflow run

[ParidelPooya]

the three wall-clock assertions. test_create_checkpoint_blocks_until_completion_default (line 3575), ..._explicit_true (3647), and test_create_checkpoint_sync_with_empty_checkpoint (4184). The PR takes their margin from +100 ms to ~0: measured min elapsed 0.2523 s on main versus 0.1499 s on the PR against a 0.15 s threshold, with 1 of 30 rounds
already below. It failed twice in the real suite under CPU load. GitHub's runners are noisier than my machine, so this will produce intermittent red builds that someone else will have to bisect.

[yaythomas]

the three wall-clock assertions
I think this might've been against an older version of this PR, already fixed in previous commit?