feat: run-budget policy and call accounting (closes #106)

[amarzullo24]

Summary

• Adds benchmaxxing/budget.py: the policy/accounting layer on top of the existing gateway retry/rate-limit handling (Gateway retry, rate-limit handling, and error surfacing #35).
• RunBudget (max_cases, max_calls, seed) + subsample_cases: deterministic subsampling seeded by seed, so a smaller max_cases is always a strict prefix of a larger run over the same pool (a pilot is a subset of the full run, not an independent sample).
• CallAccountant: tallies calls (and tokens, when a backend reports them) per model, exposed via .summary() for the run report.
• BudgetedBackend + BudgetExceeded: wraps a backend so calls are tallied and max_calls is enforced; raises BudgetExceeded instead of making the call once the shared accountant's budget is hit, so a caller can stop cleanly and flush partial results.
• truncation_note(planned, completed): an explicit, non-silent note for the run summary when a run stops short.

Closes #106.

Known scope boundaries

• No backend today returns token usage (Backend.complete() returns a bare string), so tokens_per_model degrades gracefully to empty rather than inventing a wider Backend interface change — out of scope for this issue.
• No --max-calls CLI flag yet: there's no benchmaxxing run entrypoint on main to attach it to (tracked separately in Experiment entrypoint: a benchmaxxing run CLI that drives the stage runners from a config #101 / PR feat: benchmaxxing run entry point for the stage runners #251). The reusable primitives here are ready to wire in once that lands.

Test plan

• ruff check . passes
• pytest -q — 653 passed, 6 skipped (unrelated, pre-existing)
• New tests cover: deterministic subsampling, prefix/superset property across different max_cases, per-model/total call tallying, BudgetExceeded at the right call, and a partial-result-flush pattern in a mock run loop

[sebasmos]

Clean budget/accounting layer, and the deterministic-prefix subsampling (a pilot being a strict subset of the full run, not an independent sample) is the right property for reproducibility. Enforcing the cap by raising before the call rather than after is correct. Infra, so mock tests suffice under the real-data rule. Approving; it is independent of the run-CLI stack, so it can land whenever.

[felipeocampoos]

Reviewed. I'd independently written the same feature (RunBudget/subsample/CallAccountant/BudgetedBackend, near-identical shape) before noticing this PR was already open -- sorry for the duplicate effort, withdrawing mine and reviewing this one instead since it landed first.

Two correctness issues, both unexercised by the current tests:

1. BudgetedBackend.complete checks the accountant's global total, not this model's own count (benchmaxxing/budget.py):

if self.max_calls is not None and self.accountant.total_calls >= self.max_calls:

CallAccountant is explicitly designed to be shared across models (it tallies calls_per_model), and max_calls is a per-BudgetedBackend (i.e. per-model) parameter. But the check uses accountant.total_calls, the sum across every model sharing that accountant. Concretely: two backends, model-a (max_calls=10) and model-b (max_calls=10), sharing one accountant. If model-b makes its 10 calls first, model-a is refused on its first call (total_calls == 10 >= 10) even though it hasn't spent any of its own budget yet. Each model's budget should be independent unless that's an intentional shared-pool design (worth calling out explicitly in the docstring if so). Suggested fix: self.accountant.calls_per_model.get(self.model, 0) >= self.max_calls.

2. No lock around the check-then-record: a race under concurrent calls. The check and the accountant.record() that follows aren't atomic, so two threads can both pass the total_calls >= max_calls check before either records, and both proceed -- overspending the budget. This isn't theoretical here: every real experiment script in this repo drives its backend calls through ThreadPoolExecutor (4-6 workers is typical), so a shared CallAccountant/BudgetedBackend used the way this module is meant to be used will race in practice, not just in principle.

Also worth double-checking (not a bug, just flagging): subsample_cases uses np.random.default_rng(seed).permutation(len(cases)), which is index-based over whatever order cases arrives in. The nesting/prefix property (smaller max_cases is a subset of larger) only holds if the same cases list, in the same order, is passed to both calls. If an upstream step ever reorders or re-filters the pool between a pilot run and the full run (e.g. a different query, a reloaded manifest with different row order), the two subsamples would diverge silently. Ranking by a stable hash of (seed, case_id) instead of by list-index permutation would make the nesting property hold regardless of input order -- might be worth a test asserting invariance to input order either way, since that's the property the issue actually needs ("a cheap pilot is a strict subset of the larger run").

Happy to help fix these if useful, or feel free to take it from here since it's your branch.

[amarzullo24]

Thanks @felipeocampoos, and no worries about the duplicate — all three addressed in 389b87b.

1. Per-model starvation. Root cause was two disconnected max_calls: a dead one on RunBudget and a live per-BudgetedBackend one. Rather than the per-model check, I collapsed them into a single run-wide cap owned by CallAccountant — that matches the issue's run_budget block (one scalar max_calls + a per-model breakdown for the summary), and CallAccountant.from_budget(RunBudget) now wires the previously-dead field to it. BudgetedBackend no longer takes max_calls. Models that genuinely need independent budgets get their own accountant — proven by test_independent_budgets_do_not_starve_each_other (your exact model-a/model-b scenario). Called out as intentional shared-pool in the docstring.

2. Race. Check + increment are now one atomic try_record() under a threading.Lock. test_cap_is_atomic_under_concurrent_calls drives 500 tasks through a 16-worker ThreadPoolExecutor against a cap of 100 and asserts exactly 100 succeed. One deliberate consequence worth flagging: I went claim-before-spend (reserve under the lock, then call the backend), because holding the lock across the network call would serialize all API traffic and defeat the fan-out. Trade-off is that an errored mid-flight call still counts against the budget — the conservative direction for a spend cap (a flapping backend can't retry-storm past it), pinned by test_a_failed_backend_call_still_counts_against_the_budget.

3. Order-dependent subsample. Switched from the numpy index permutation to ranking by a stable blake2b(seed, case_id) hash (index as tiebreak), exactly your stable-hash suggestion. The nested-subset property now holds regardless of input order — test_subsample_is_invariant_to_input_order shuffles the pool between calls and asserts the same subset. Dropped the numpy dependency here since index-based RNG was the root of the fragility.

ruff check . clean, full suite green (658 passed / 6 skipped).

[sebasmos]

The two fixes (per-model starvation, check-record race) are verified in the head. Before merge though, please run it for real once: a real budgeted run with max_calls below the natural count, showing the cap actually stops calls and the per-model tally matches. Green tests were not enough here (that is exactly what the two bugs showed), so one real run is the bar.

[maximinl]

Peer review after the Felipe fixup (389b87b).

The three corrections look right:

1. Run-wide max_calls on CallAccountant matches the issue’s one-scalar budget, and independent accountants are the documented escape hatch (pinned by test_independent_budgets_do_not_starve_each_other).
2. try_record under a lock + claim-before-spend is the right concurrent shape for ThreadPoolExecutor fan-out; counting failed backend calls is the conservative direction for a spend cap.
3. blake2b(seed, case_id) ranking makes the nested-subset property order-invariant.

Remaining merge gate is Seb’s ask: one real budgeted run with max_calls below the natural count, showing the cap stops calls and the per-model tally matches. Happy to approve the code; that live check should still land before merge.

[sebasmos]

Formalizing the ask above: one real budgeted run (cap actually stops calls, per-model tally shown) before merge.

[Agastya191]

Nice work on budget.py, especially the atomic try_record cap and the ThreadPoolExecutor stress test. One thing on subsample_cases: the early return for max_cases None or >= len(cases) hands cases back in pool order while the subsampling path returns them in rank order, so a bounded pilot is a subset of an unbounded full run but not a prefix of it, and the "strict prefix of a larger run" from the summary only holds when the larger run is also bounded (the prefix test compares max_cases=5 against 20, both bounded). In practice that means diffing a pilot against the real, unbounded full run won't line up positionally even though the case sets nest correctly. You've got options; I'd lean toward ranking on the unbounded path too and slicing at len(cases) so the full run and every pilot share one order, or if subset is all you really need, softening the wording from prefix to subset.

[sebasmos]

Ran the real budgeted check (real Gemini API): cap holds at 4 (per-model tally flash 2 / flash-lite 2), calls refused before spending, and 200 concurrent threads admit exactly 10 at a cap of 10, so the per-model-starvation and check-then-record fixes both hold on real hardware. Also pushed the fix for Agastya's point: the unbounded path now ranks too, so a bounded pilot is a strict prefix of the full run (test updated + prefix test added, 16/16 pass, ruff clean). Merge-ready. Thanks both for the careful review.