Repository navigation
Feat/token budget allocator - #189
Merged
Merged
Conversation
added 8 commits
September 14, 2026 16:19
max_calls and max_rollouts count invocations, and a reasoning model can spend 40k tokens on hidden thinking in a single one. A 20-round run with 4 workers is 80 rollouts and ~160 calls but, at 40k tokens each, 6.4M tokens — the count is not the cost. evolve(max_tokens=...) caps the total (prompt + completion): - sync path: checked at the round barrier alongside the other budgets - async path: checked per-rollout, for tighter control - stop_reason="max_tokens", reporting the spend actually incurred The token flow had a foot-gun that made the budget silently miss: the adapters that report tokens (claude, openai_compatible) take a Usage in their constructor, and the meter only saw those numbers when the caller passed the same object to evolve(usage=...). Forgetting the second thread made the run report tokens=0 while the bill went to an object only the caller held. evolve() now adopts the agent usage automatically when the caller did not pass one; an explicit usage= still wins. RoundInfo.tokens carries the cumulative spend to on_round (and so to rounds.jsonl), RunStatus.tokens lands it in status.json, and the CLI status line shows tok=N when the run has any. 13 new tests: budget fires on both paths, generous cap does not interfere, zero-token runs are honest, agent usage is adopted, explicit usage wins, RoundInfo/RunStatus serialization.
…ken wall
max_tokens was a brake: when the spend reaches it, the run stops. A brake is
the floor of cost control, not its ceiling. The last 10% of a token budget is
the most expensive part to waste — the search has paid for its exploration,
and a round dispatched there may never finish.
This adds a BudgetGovernor that degrades before the wall:
- at 75% of the budget (soft_floor), stop building fusion tournaments.
The tournament is an extra held-out sweep of every survivor plus the
fusion — pure ranking spend, and the ranking can survive without it for
the last few rounds. The aggregator acceptance gate runs unchanged.
- at 90% (hard_floor), stop the self-verify rollout. That rollout doubles
the cost of every proposal for a delta the acceptance test folds in as
a tie-breaker weight (observe_delta with weight min(1, 4*|delta|));
skipping it costs ranking precision on the advantage signal and nothing
on the commit gates, which read the full held-out set either way.
- the projection: at the round barrier, the governor projects the average
per-round spend and asks whether the next round fits the remaining
budget. If it does not, the run ends with a clean merge at the last
round it fully paid for, rather than being cut mid-round at the wall.
Reported as stop_reason="max_tokens" (the budget that fired).
All of this is inert without max_tokens: no budget, no governor, no change.
The governor mutates the fusion policy tournament flag at the round barrier,
which is safe because _tournament reads it in step() — which runs after the
barrier in the same round. The self-verify check is inline in the worker body,
gated by governor.allow_self_verify().
12 new governor tests + 13 token budget tests + 96 core regression tests,
121 passed total.
Three gaps closed: 1. EvolutionResult.budget: a new Optional[dict] field carrying the governor summary (max_tokens, spent, remaining, fusion_degraded, self_verify_degraded). None when no budget was set. A caller can now tell "the run finished because its budget ran out, degraded" from "it converged, at full quality" — and whether the tail of the search ran with fewer ranking passes. 2. CHANGELOG: the governor entry was missing entirely. Added a full description of the three-stage degradation (soft floor, hard floor, projection) and how it differs from the max_tokens brake. 3. Async path governor: was missing from the previous commit. The async_evolve path now constructs a BudgetGovernor, calls governor.spend() after each merger sweep, mutates fusion_policy.tournament at the soft floor, and gates self_verify in the worker body at the hard floor. Two new tests verify the governor fires on the async path (stop_reason=max_tokens, fusion degrades). 3 new tests: result.budget populated, result.budget None without cap, result.budget shows degradation.
…n hold cost fixed
The module argument (baselines.py) is that a budget in one unit is not a
comparison: matched on rollouts alone, an arm that asks for more proposals
per rollout spends more model and the table still says "equal budget".
That argument was missing the unit cost is measured in. Two arms at equal
rollouts AND equal calls can differ 10x in tokens when one uses a
reasoning model — and the table had no column that would show it.
Wiring:
- Budget gains tokens (split like rollouts/calls); Workload._evolve
passes it to evolve(max_tokens=), so the arm with a token budget
spends the same money as its siblings.
- compare(fixed="tokens") is a valid unit; the drift check runs over
all three units so a token confound appears in confounded when
rollouts is fixed.
- ArmResult.tokens is a measured property (prompt + completion), not
a requested number.
- to_markdown gains a token column that reads — (not 0) when the
actor reported none — an unreported unit is not zero spend.
Also fixed: compare([]) raised StatisticsError (a ValueError), so a test
validating the fixed argument was passing for the wrong reason. Empty
input now returns an empty Comparison.
7 new tests: budget split incl. tokens, ArmResult.tokens, fixed=tokens
flags unequal, tokens confounded under fixed=rollouts, markdown column,
empty compare.
The o1 test-time-scaling decision, as a search parameter. Until now every
token-budget mechanism this branch added was accounting: max_tokens caps,
the governor degrades, baselines can fix tokens. None of them changed *which
candidate the next expansion goes to*. The search still allocated compute by
score alone, blind to what a token buys.
Candidate gains cost (None = unknown, never zero — the same discipline
score and prior follow). SelectionContext gains budget_remaining (the
unspent fraction, 1.0 without a ceiling). CostEfficient is FlatPuct with
two additions:
- the exploration term is divided by cost: a speculative expansion of an
expensive candidate has to beat one of a cheap candidate at the same
rank. The exploitation term is NOT discounted, because charging cost
against a grounded pick biases the search against complex solutions
(longer artifacts cost more to expand) — a bias that is not a quality
judgment and must not be smuggled in as one.
- the exploration bonus anneals by budget_remaining: explore while there
is budget to exploit what you find, exploit once there is not. That is
what an anytime search must do, and the difference between spending a
budget and allocating it.
Safety properties, all tested:
- cost_exponent=0 is FlatPuct to the floating-point bit, so the mechanism
cannot change a run that did not ask for cost-awareness.
- uniform costs are unchanged at any exponent.
- an unknown cost is the pool mean, never "free" — reading a missing
number as 0 would make every unmeasured candidate the cheapest.
- a negative budget_remaining is clamped, so a caller bug cannot invert
the exploration sign.
PopulationAggregator fills Candidate.cost with a deterministic content-size
proxy (the prompt dominates input tokens; size is attributable to one
candidate without racing a shared meter) and passes budget_remaining into
the selection context. The engine feeds budget_remaining from the governor
each round, sync and async, and only to an aggregator that declares the
slot.
18 new tests.
Two things.
1. The economic stop. max_tokens is a ceiling; the question is whether to
*spend* it. evolve(stop_on_diminishing_returns=True) measures the run
own return per token (delta reward / delta tokens per round, from the
reward and meter the run already has) and stops when the recent rate has
fallen to efficiency_floor (default 0.25) of the run peak rate. Self-
calibrating: it compares the run against itself, so there is no absolute
rate to guess (reward is in [0,1], tokens in the millions, quotient has
no interpretable scale). Off by default -- a run whose reward only rises
late would be cut short -- and requires max_tokens. A run that never
improved has no peak to decline from, so this stays quiet and patience
handles it. Wired on sync and async; stops with
stop_reason="diminishing_returns".
2. tests/test_backward_compat.py -- the contract that every mechanism on
this branch is inert until asked for:
- no new options == the old path (no budget reported, no checkpoint dir,
default selection is still SingleHead).
- passing the new options at their defaults == passing none (identical
reward history).
- a never-binding max_tokens leaves the search identical, so a budget is
additive: it can end a run early and degrade optional spend when close,
and touches nothing with room to spare.
- the governor with max_tokens=None says nothing and changes nothing.
Also: documented the two new params on both entry points (test_api_docs
caught the omission), regenerated docs/api.md.
15 new tests (7 governor, 1 e2e plumbing, 6 backward-compat, 1 doc).
… search layer
The o1 insight, fully realized: a reasoning model per-call thinking
budget should be proportional to the expected value of that expansion,
not a fixed number for the whole run.
CallBudget (agentdescent/budget.py):
- base: the adapter configured max_tokens (the ceiling the engine
never exceeds — it never raises a bill the caller did not agree to)
- floor: base // 4 (a reasoning model below this returns empty content
— measured on deepseek-v4-flash at 1024)
- allocate(score, budget_remaining): sets per-call max_tokens as
(0.5 + 0.5*score) * (0.5 + 0.5*budget_remaining) * base, clamped
to [floor, base]. A promising parent at full budget gets ~base;
a dead-end parent with no budget gets ~floor.
- next: None by default (use base, never change) — inert until
the engine calls allocate() before each propose.
budgeted_completion(factory, budget):
- Wraps a max_tokens -> Completion factory into an adaptive Completion.
- Caches per max_tokens value, so a run that stays at the default pays
one construction and a run that adapts pays one per distinct step.
evolve(call_budget=...):
- Before each propose call, if call_budget is set, calls
budget.allocate(score, governor.remaining_fraction()).
- score is the parent held-out reward the worker just measured;
budget_remaining is the governor fraction the engine already tracks.
- Inert without call_budget: the engine never touches it and the
adapter uses its own fixed max_tokens.
Wired on both sync (evolution.py) and async (async_evolve.py) paths.
Async path delegates call_budget to async_evolve via the evolve()
dispatch.
13 new tests: allocate bounds, caching, end-to-end, inert-by-default.
163 total passed across the affected-area regression.
Birfy
approved these changes
Sep 16, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cost-aware search: spend a token budget by value, not by exhaustion
The problem
evolve()had four budgets (max_seconds,max_rollouts,max_calls,rounds), none of which map to cost. A reasoning model spends 30–50k tokens on a single hidden-thinking call, so a 20-round run with 4 workers is ~6.4M tokens — easily $50–100 on a real deployment — and the caller had no way to control, measure, or allocate that spend.The
Meteralready trackedprompt_tokensandcompletion_tokens. The infrastructure was there; the knob, the allocation, and the search-level intelligence were not.What this PR adds — four layers, all opt-in, all inert by default
Layer 1:
evolve(max_tokens=N)— a hard cap in the unit that maps to costChecked at the round barrier (sync) and per-rollout (async).
stop_reason="max_tokens".Agent usage auto-adoption: the adapters that report tokens (
claude,openai_compatible) take aUsagein their constructor. The caller who builtLLMAgent(claude(usage=u))had to passusage=uagain toevolve()for the meter to see the same numbers. Forgetting the second thread made the run reporttokens=0while the bill went tou.evolve()now adopts the agent's ownusagewhen the caller did not pass one.Layer 2:
BudgetGovernor— graceful degradation before the wallThe last 10% of a budget is the most expensive to waste. The governor degrades before the wall:
soft_floor)hard_floor)Design principle: every commit gate (Beta-posterior acceptance + regression guard) runs to the last round at full quality. Only optional ranking spend degrades.
stop_on_diminishing_returns=True: the economic stop — stop when the run's own return per token has fallen off its peak, even with budget left. Self-calibrating: compares the run against itself, so no absolute rate to guess. Stops withstop_reason="diminishing_returns".Layer 3:
CostEfficientselection — spend the budget, don't just count itThe o1 test-time-scaling decision, as a search parameter:
Cost discounts exploration, not exploitation. The exploration term is divided by the candidate's cost (a content-size proxy — the prompt dominates input tokens, and size is attributable to one candidate without racing a shared meter). A speculative expansion of an expensive candidate must beat one of a cheap candidate at the same rank. The exploitation term is left alone: charging cost against a grounded pick biases the search against complex solutions (longer artifacts cost more to expand) in favour of short ones, which is not a quality judgment.
Exploration anneals by
budget_remaining. Explore while there is budget to exploit what you find; exploit once there is not.anneal=Falseis the ablation that separates "cost-aware" from "budget-aware".Safety properties (all tested):
cost_exponent=0isFlatPuctto the floating-point bit; uniform costs are unchanged; unknown cost is the pool mean, never "free"; negativebudget_remainingis clamped.Candidategainscost(None= unknown, never zero).SelectionContextgainsbudget_remaining.PopulationAggregatorfills both.Layer 4:
CallBudget+budgeted_completion— adaptive per-call thinking budgetThe o1 insight fully realized: a reasoning model's per-call thinking budget should be proportional to the expected value of that expansion, not a fixed number for the whole run.
CallBudget.allocate(score, budget_remaining)sets the per-call ceiling:0.5 + 0.5 * score): a high-scoring parent gets more thinking; a zero-scoring one gets half.0.5 + 0.5 * budget_remaining): early in the run, full thinking; late, tightened.The product ranges from
floor(base // 4, the minimum below which a reasoning model returns empty content) tobase(the adapter's configured ceiling — the engine never raises a bill the caller did not agree to).budgeted_completioncaches permax_tokensvalue, so a run that stays at the default pays one construction.Layer 5:
baselines.Budget(tokens=...)— the A/B framework can hold cost fixedThe module's own argument is that a budget in one unit is not a comparison. Two arms at equal rollouts and equal calls can differ 10x in tokens when one uses a reasoning model — and the table had no column that would show it.
Budgetgainstokens(split like the others),Workload._evolvepasses it toevolve(max_tokens=),compare(fixed="tokens")is a valid unit,ArmResult.tokensis the measured sum, and the markdown table gains a token column that reads—(not0) when the actor reported none.Observability
RoundInfo.tokens: cumulative spend per round →on_round→rounds.jsonlRunStatus.tokens: instatus.json, CLI showstok=NEvolutionResult.budget: governor summary (max_tokens,spent,remaining,fusion_degraded,self_verify_degraded) —Nonewhen no budget was settokens=Nper round, logs when degradation firesBackward compatibility — guaranteed by tests
tests/test_backward_compat.py(6 behavioural assertions):SingleHead.max_tokensleaves the search identical: a budget is additive — it can end a run early and degrade optional spend when close, and touches nothing with room to spare.max_tokens=Nonesays nothing and changes nothing.CostEfficientis not the default: default selection isSingleHead.call_budget=Nonemeans the per-call ceiling never changes.Verification
Files
18 files changed, +2391, -35
New:
agentdescent/budget.py(+534) —BudgetGovernor,CallBudget,budgeted_completiontests/test_token_budget.py(+363) — end-to-end budget + observabilitytests/test_budget_governor.py(+220) — governor thresholds, projection, diminishing returnstests/test_cost_aware_selection.py(+318) — cost-aware selection + annealing + population wiringtests/test_call_budget.py(+206) — adaptive per-call budgettests/test_backward_compat.py(+110) — the guaranteeModified:
agentdescent/evolution.py(+160) —max_tokens,call_budget,stop_on_diminishing_returns,efficiency_floorparams; governor + budget_remaining + call_budget wiring;EvolutionResult.budgetagentdescent/async_evolve.py(+50) — governor on async path (spend, fusion degrade, self_verify gate, diminishing stop, call_budget)agentdescent/selection.py(+120) —Candidate.cost,SelectionContext.budget_remaining,CostEfficientagentdescent/population.py(+36) — fillCandidate.cost, passbudget_remainingagentdescent/baselines.py(+73) —Budget.tokens,compare(fixed="tokens"),ArmResult.tokens, markdown columnagentdescent/runstore.py(+8) —RunStatus.tokensagentdescent/cli.py(+3) —tok=Nin status lineagentdescent/__init__.py(+6) — exportsCHANGELOG.md(+86)docs/api.md(+43)