diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c039e78..9c20641c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,107 @@ All notable changes to AgentDescent are documented here. The format follows ### Added +- **`evolve(max_tokens=...)`: a budget in the unit that maps to the bill.** + `max_calls` and `max_rollouts` count invocations, and a reasoning model can + spend 40k tokens on hidden thinking in a single one -- so a 20-round run with + 4 workers is 80 rollouts and ~160 calls but, at 40k tokens each, 6.4M tokens. + The cap is checked at the round barrier on the sync path (alongside the other + budgets) and per-rollout on the async path, stops with + `stop_reason="max_tokens"`, and reports the spend it actually incurred. + + Tokens only reach the meter when the model adapter reports them, and the + adapters that do (`claude`, `openai_compatible`) take a `Usage` in their + constructor. The caller who built `LLMAgent(claude(usage=u))` had to pass + `usage=u` **again** to `evolve()` for the meter — and with it `max_tokens`, + the `tokens=` column on every round, and `status` — to see the same numbers. + Forgetting the second thread made the run report `tokens=0` while the bill + went to `u`, an object only the caller held, and a token budget silently + never fired. `evolve()` now adopts the agent's own `usage` when the caller + did not pass one; an explicit `usage=` still wins, and a bare + `run`/`propose` pair still reports zero (the honest reading, not an error). + + `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. + +- **`BudgetGovernor`: graceful degradation before the token wall.** + `max_tokens` is 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. The governor degrades before + the wall on both paths (sync and async): + + - 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; skipping it costs ranking precision on the advantage + signal and nothing on the commit gates. + - the projection at the round barrier (sync) 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. + + All of this is inert without `max_tokens`. `EvolutionResult.budget` carries + the governor's summary (spent, remaining, whether degradation fired) so a + caller can tell "the run finished because its budget ran out, degraded" from + "it converged, at full quality". + +- **`SelectionContext.budget_remaining` + `CostEfficient`: spend the budget, don't just count it.** + The o1 test-time-scaling decision, as a search parameter. `SelectionContext` + gains `budget_remaining` (the unspent fraction, `1.0` with no ceiling), and + `Candidate` gains `cost` (`None` = unknown, never zero — the same distinction + `score` and `prior` make). `CostEfficient` is `FlatPuct` with two additions: + + - the **exploration** term is divided by cost, so 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 + would bias the search against complex solutions (longer artifacts cost more + to expand) in favour of short ones, which is not a quality judgment and must + not be smuggled in as one; + - the exploration bonus **anneals by `budget_remaining`**, so a run explores + while it has budget to exploit what it finds and exploits once it does not — + what an anytime search has to do, and the difference between spending a + budget and allocating it. + + `cost_exponent=0` is `FlatPuct` to the floating-point bit (tested), so the + mechanism cannot change a cost-blind run; uniform costs are unchanged too. + `anneal=False` is the ablation that separates "cost-aware" from "budget-aware". + The population layer fills `Candidate.cost` with a deterministic content-size + proxy (the prompt dominates the input tokens, and size is attributable to one + candidate without racing a shared meter) and the engine feeds + `budget_remaining` from the governor each round, on both sync and async paths. + +- **`evolve(stop_on_diminishing_returns=True)`: stop when the search stops + paying for itself.** `max_tokens` is a ceiling; the economic question is + whether to *spend* it, and the answer is to keep buying compute while it + still returns quality and stop when it does not — even with budget left. The + governor measures the run's own return per token (`Δreward / Δtokens` per + round, from the same reward and meter the run already has) and stops when the + recent rate has fallen to `efficiency_floor` (default a quarter) of the run's + *peak* rate. Self-calibrating: it compares the run against itself, so there + is no absolute rate to guess — a reward is in `[0, 1]`, a token count is in + the millions, and their quotient has no interpretable scale. Off by default, + because a run whose reward only rises late would be cut short by it, and + because it needs `max_tokens`. Stops with + `stop_reason="diminishing_returns"`. A run that never improved has no peak to + decline *from*, so this stays quiet and the patience counter handles it. + +- **`baselines.Budget(tokens=...)`: the A/B framework can hold cost fixed.** + The module's own argument 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. `Budget` gains `tokens` (split like the + others), `Workload._evolve` passes it to `evolve(max_tokens=)`, + `compare(fixed="tokens")` is a valid unit, `ArmResult.tokens` is the measured + sum, and the markdown table gains a token column that reads `—` (not `0`) + when the actor reported none. A comparison that fixes tokens is the only one + that has shown an arm to be more *efficient* rather than merely given more + model. + - **`evolve(checkpointing=True)`: the search survives a process restart.** `repo_path` already resumed the *artifact* -- the ledger's whole job -- but the aggregator's search state is only ever in memory, so a resume re-seeded diff --git a/agentdescent/__init__.py b/agentdescent/__init__.py index 732777c9..b636aa48 100644 --- a/agentdescent/__init__.py +++ b/agentdescent/__init__.py @@ -50,6 +50,7 @@ Archive, Beam, Candidate, + CostEfficient, MCTS, MultiHeadUnsupported, ParetoFrontier, @@ -120,6 +121,7 @@ clear_checkpoints, list_checkpoints, ) +from .budget import BudgetGovernor, CallBudget, budgeted_completion from .executor import Executor, Result, ThreadExecutor from .sandbox import LocalWorkspaceSandbox, SandboxPool, WorkspaceProvider from .sandbox_shared import SharedSandboxPool @@ -256,6 +258,7 @@ "Beam", "ParetoFrontier", "Archive", + "CostEfficient", "MCTS", "MultiHeadUnsupported", "pareto_front", @@ -440,4 +443,8 @@ "resolve_from_mapping", "verifier_fingerprint", "residual_bias", + # the cost-aware search layer + "BudgetGovernor", + "CallBudget", + "budgeted_completion", ] diff --git a/agentdescent/async_evolve.py b/agentdescent/async_evolve.py index 1cb19e8a..93c680ea 100644 --- a/agentdescent/async_evolve.py +++ b/agentdescent/async_evolve.py @@ -39,6 +39,7 @@ from .advantage import GroupAdvantage from .agents import Usage +from .budget import BudgetGovernor, CallBudget from .policies import Policies from .evolution import ( _publish_stable, _safe_log, @@ -73,6 +74,14 @@ def async_evolve( max_seconds: float = 20.0, max_iters: Optional[int] = None, max_calls: Optional[int] = None, + max_tokens: Optional[int] = None, + #: Stop when the run's measured return per token falls off its own peak. + #: See :func:`evolve`; off by default, needs ``max_tokens``. + stop_on_diminishing_returns: bool = False, + #: How far below the peak counts as diminishing. See :func:`evolve`. + efficiency_floor: float = 0.25, + #: An adaptive per-call thinking budget. See :func:`evolve`. + call_budget: Optional[CallBudget] = None, target_reward: Optional[float] = None, patience: Optional[int] = None, max_worker_errors: int = 3, @@ -183,6 +192,28 @@ def async_evolve( more proposals per rollout, and the cheaper unit is the one a reader assumes was held fixed. Both bounds are checked as each rollout lands, so a run overshoots only by what was already in flight. + max_tokens: + Stop after this many tokens in total (``prompt + completion``), as the + meter measured them. The third budget unit, and the one cost is measured + in: ``max_calls`` and ``max_iters`` count invocations, and a reasoning + model can spend 40k tokens on hidden thinking in a single one, so neither + bounds the bill. Checked as each rollout lands (tighter than the + synchronous path's round barrier), and a :class:`~agentdescent.budget.BudgetGovernor` + degrades optional spend -- fusion tournaments at 75% of the budget, + self-verify at 90% -- before the wall. ``None`` (default) means + unbounded, and no governor is constructed. + stop_on_diminishing_returns: + Stop when the run's measured return per token falls off its own peak by + more than ``efficiency_floor``. Off by default; needs ``max_tokens``. + See :func:`evolve`. + efficiency_floor: + How far below the peak counts as diminishing, when the stop above is on. + A self-calibrating ratio, not an absolute rate. See :func:`evolve`. + call_budget: + An adaptive per-call thinking budget (o1-style test-time scaling). + See :func:`evolve`; the allocation happens before each worker's + ``propose`` call, gated by the same ``score`` and + ``governor.remaining_fraction()`` the synchronous path uses. eval_concurrency: How many held-out tasks the merger scores at once. ``1`` restores the old sequential behaviour. @@ -352,6 +383,17 @@ def async_evolve( usage=usage, verifier=_pol.verifier, ledger_impl=_pol.ledger, policies_bundle=_pol, checkpointing=checkpointing) eng.meter.start() + # The cost-aware governor for the async path. The merger loop has no round + # barrier, so the governor checks per sweep (after each merge, before the + # next worker batch's self-verify fires) rather than at the barrier. The + # soft floor turns off fusion tournaments between sweeps; the hard floor + # turns off self-verify in the worker body — read through eng.checkpointing + # would be wrong (that is checkpoint, not budget), so the governor lives on + # eng alongside it. + governor = BudgetGovernor(max_tokens=max_tokens, + stop_on_diminishing=stop_on_diminishing_returns, + efficiency_floor=efficiency_floor) + eng.governor = governor if n_workers < 1: raise ValueError(f"n_workers must be >= 1, got {n_workers}") policy = staleness_policy or get_policy("guarded") @@ -611,6 +653,9 @@ def _worker(wid: int, shard: List[Task]) -> None: adv = advantage.observe( advantage.key(base_v, str(task.meta.get("cluster", ""))), score) if score < solved_threshold: + # o1-style test-time scaling: see evolve() for the reasoning. + if call_budget is not None: + call_budget.allocate(score, governor.remaining_fraction()) proposal = _checked_proposal( eng.propose(artifact.render(), task, output, score), task) if proposal: @@ -621,7 +666,7 @@ def _worker(wid: int, shard: List[Task]) -> None: # diff applied for a before/after signal. Faithful repos that # only score the candidate on held-out (e.g. EvoSkill) pass # self_verify=False to skip this extra rollout. - if self_verify: + if self_verify and governor.allow_self_verify(): after = _checked_reward( eng.reward(task, eng.run(artifact.apply(diff).render(), task)), task) delta = after - score @@ -723,6 +768,16 @@ def _worker(wid: int, shard: List[Task]) -> None: # ratio and the second budget cannot be derived from the first. stop_reason[0] = "max_calls" stop.set() + elif max_tokens is not None and eng.meter.usage.total_tokens >= max_tokens: + # Same reasoning as max_calls: read from the meter, because + # a reasoning model can spend 40k tokens on one hidden-thinking + # call and `calls` says nothing about the bill. The async + # path checks per-rollout rather than per-round-barrier, so a + # token budget fires sooner here than on the sync path — which + # is correct: the barrier-free loop is the one that can run + # away with spend if the gate keeps accepting. + stop_reason[0] = "max_tokens" + stop.set() def _drain_and_merge() -> None: with intake_lock: @@ -794,6 +849,32 @@ def _discarded() -> None: _discarded() # DISCARD -> drop the card with eng.meter.timed("merge_gate_seconds"): reports = check_reports(_gated_step(), eng.aggregator) + # Governor: record the spend after each merge and degrade fusion before + # the next sweep. The merger is the single-threaded decision point, so + # the mutation is safe — no worker reads ``fusion_policy.tournament`` + # until the next ``_gated_step`` which is the next sweep. + if governor.active: + gov_tokens = eng.meter.usage.total_tokens + governor.spend(gov_tokens) + # The previous sweep's reward, for the diminishing-returns measure + # (this runs before `record_round` for the current sweep, so + # `history[-1]` is the last *completed* one; `None` before any). + governor.observe(history[-1].held_out_reward if history else None) + if hasattr(eng.aggregator, "budget_remaining"): + eng.aggregator.budget_remaining = governor.remaining_fraction() + fp = getattr(eng.aggregator, "fusion_policy", None) + if fp is not None and hasattr(fp, "tournament"): + should = governor.allow_fusion_tournament() + if fp.tournament != should: + fp.tournament = should + if verbose: + tag = "on" if should else "off" + print(f" fusion tournament {tag} " + f"(token budget: {gov_tokens:,}/" + f"{governor.max_tokens:,})") + if governor.diminishing_returns() and not stop.is_set(): + stop_reason[0] = "diminishing_returns" + stop.set() if not reports and not batch: # A pipelined poll whose candidate is still being measured has # nothing to report yet. Returning keeps the merger draining -- @@ -859,8 +940,10 @@ def _discarded() -> None: epoch[0] += 1 # every worker resyncs on its next loop stall.force() if verbose: + u = eng.meter.usage print(f"sweep {len(history):>3} reward={r:.3f} merged={len(batch)} " - f"+{committed} pending={len(intake)}") + f"+{committed} pending={len(intake)}" + + (f" tokens={u.total_tokens:,}" if u.total_tokens else "")) if early_stop is not None: stop_reason[0] = early_stop stop.set() diff --git a/agentdescent/baselines.py b/agentdescent/baselines.py index fae821be..411ae8c4 100644 --- a/agentdescent/baselines.py +++ b/agentdescent/baselines.py @@ -114,6 +114,16 @@ class Budget: budget"; matched on calls alone, an arm with a cheaper proposer buys more rollouts. Neither is derivable from the other -- a rollout that solves its task never proposes at all. + + ``tokens`` is the third unit, and the one that maps to the bill. It is + ``None`` by default because it is only measurable when the actor reports + token counts (an adapter built with ``claude(usage=u)``, and the same + ``u`` threaded to the run); a comparison between two API models should set + it, and a comparison between two wrappers around one opaque CLI cannot. + When it is set, ``evolve``'s ``max_tokens`` enforces it and the governor's + degradation applies -- so an arm at a token budget spends the *same* money + as its siblings, which neither ``rollouts`` nor ``calls`` can promise + across arms that differ in how much the model thinks per call. """ rollouts: int @@ -121,13 +131,19 @@ class Budget: #: arm shares one actor and the ratio is stable. Set it as soon as arms #: differ in how often they propose. calls: Optional[int] = None + #: ``None`` means "do not bound tokens". Set it when the actor reports them + #: -- it is the unit cost is measured in, and an arm that wins at equal + #: tokens is the only thing that has been shown to be more *efficient* + #: rather than merely given more model. + tokens: Optional[int] = None def split(self, ways: int) -> "Budget": """The share of this budget one of ``ways`` independent runs may spend.""" if ways < 1: raise ValueError(f"cannot split a budget {ways} ways") return Budget(rollouts=self.rollouts // ways, - calls=None if self.calls is None else self.calls // ways) + calls=None if self.calls is None else self.calls // ways, + tokens=None if self.tokens is None else self.tokens // ways) @dataclass(frozen=True) @@ -172,6 +188,7 @@ def _evolve(self, *, seed: int, n_workers: int, budget: Budget, self.tasks, self.reward, seed=seed, n_workers=n_workers, max_concurrency=kwargs.pop("max_concurrency", n_workers), max_rollouts=budget.rollouts, max_calls=budget.calls, + max_tokens=budget.tokens, usage=usage, **kwargs) @@ -230,8 +247,7 @@ class ArmResult: rollouts: int calls: int prompt_tokens: int - completion_tokens: int - #: Sum of every run in the arm. For fork this is what it cost, not what it + completion_tokens: int #: Sum of every run in the arm. For fork this is what it cost, not what it #: would take: N forks are independent and run concurrently given N workers, #: which is why `wallclock_parallel` exists beside it. wallclock: float @@ -262,6 +278,17 @@ class ArmResult: #: which is itself the right answer for them. fusion: Optional["FusionStats"] = None + @property + def tokens(self) -> int: + """Total tokens the arm spent (prompt + completion), as measured. + + A property rather than a field, so it is always the sum of the two + counts beside it and cannot drift from them. Zero when the actor + reported none -- which is why a comparison on ``fixed="tokens"`` must + check that the arms actually reported before reading the column (see + :func:`compare`).""" + return self.prompt_tokens + self.completion_tokens + def _tally(arm: str, seed: int, width: int, runs: Sequence[EvolutionResult], usages: Sequence[Usage], *, scoring_error: str = "", @@ -406,7 +433,10 @@ class Comparison: """Several seeds of several arms, and whether they are comparable at all.""" arms: Dict[str, List[ArmResult]] - #: The unit the comparison holds fixed: ``"rollouts"`` or ``"calls"``. + #: The unit the comparison holds fixed: ``"rollouts"``, ``"calls"`` or + #: ``"tokens"``. Tokens is the unit cost is measured in, so it is the one to + #: fix when the actor reports them; it is not the default only because it + #: needs reporting the other two do not. fixed: str = "rollouts" #: Arms whose spend in the **fixed** unit drifted from the median by more #: than the tolerance, as ``(arm, unit, spend, median)``. Non-empty means the @@ -478,24 +508,33 @@ def compare(results: Sequence[ArmResult], *, fixed: str = "rollouts", tolerance: float = 0.1) -> Comparison: """Group arm results by arm and check what the comparison actually held fixed. - ``fixed`` names the unit the budget was set in. Drift there goes to - ``unequal`` and invalidates the comparison; divergence in the other unit goes - to ``confounded``, which is a caveat rather than a defect -- see the module - docstring for why it cannot be designed away. + ``fixed`` names the unit the budget was set in -- ``"rollouts"``, + ``"calls"``, or ``"tokens"``. Drift in the fixed unit goes to ``unequal`` + and invalidates the comparison; divergence in the other units goes to + ``confounded``, which is a caveat rather than a defect -- see the module + docstring for why they cannot all be equalised by design. ``tolerance`` is a fraction of the median spend. The default allows the barrier overshoot -- an arm can end up to one round past its bound, and a wide arm's round is wide -- while still catching an arm that ran on twice the model. """ - if fixed not in ("rollouts", "calls"): - raise ValueError(f"fixed must be 'rollouts' or 'calls', not {fixed!r}") + if fixed not in ("rollouts", "calls", "tokens"): + raise ValueError( + f"fixed must be 'rollouts', 'calls' or 'tokens', not {fixed!r}") by_arm: Dict[str, List[ArmResult]] = {} for r in results: by_arm.setdefault(r.arm, []).append(r) comparison = Comparison(arms=by_arm, fixed=fixed, tolerance=tolerance) - for unit in ("rollouts", "calls"): + if not by_arm: + # No arms, nothing to check. Returning an empty comparison rather than + # raising keeps a sweep that produced nothing from crashing the caller + # that is about to report it. (`statistics.median([])` raised + # `StatisticsError`, a `ValueError`, which a caller validating the + # `fixed` argument caught by accident.) + return comparison + for unit in ("rollouts", "calls", "tokens"): spends = {arm: statistics.median([getattr(r, unit) for r in group]) for arm, group in by_arm.items()} median = statistics.median(list(spends.values())) @@ -510,9 +549,9 @@ def compare(results: Sequence[ArmResult], *, fixed: str = "rollouts", def to_markdown(comparison: Comparison) -> str: """A table whose caption cannot claim more than the numbers support.""" - lines = ["| arm | seeds | rollouts | calls | dev | test (min/med/max) | " - "fork oracle |", - "|---|---|---|---|---|---|---|"] + lines = ["| arm | seeds | rollouts | calls | tokens | dev | " + "test (min/med/max) | fork oracle |", + "|---|---|---|---|---|---|---|---|"] for arm in sorted(comparison.arms): group = comparison.arms[arm] spread = comparison.spread(arm) @@ -525,10 +564,16 @@ def to_markdown(comparison: Comparison) -> str: f"{spread[0]:.3f} / {spread[1]:.3f} / {spread[2]:.3f}") seeds = (f"{len(group)}" if comparison.scored(arm) == len(group) else f"{comparison.scored(arm)}/{len(group)}") + # Tokens read "—" when the actor reported none. That is not zero spend, + # it is an unreported unit -- and a `0` in the column would read as + # "spent nothing", which is the one thing it cannot be. + toks = [r.tokens for r in group] + token_col = f"{statistics.median(toks):.0f}" if any(toks) else "—" lines.append( f"| {arm} | {seeds} | " f"{statistics.median([r.rollouts for r in group]):.0f} | " f"{statistics.median([r.calls for r in group]):.0f} | " + f"{token_col} | " f"{statistics.median([r.dev_reward for r in group]):.3f} | " f"{quality} | {oracle} |") diff --git a/agentdescent/budget.py b/agentdescent/budget.py new file mode 100644 index 00000000..dc3eb284 --- /dev/null +++ b/agentdescent/budget.py @@ -0,0 +1,385 @@ +"""Budget governor: graceful degradation before the hard token stop. + +``max_tokens`` is a brake: when the spend reaches it the run stops, full stop. +A brake is the floor of cost control, not its ceiling. The last 10% of a token +budget is the most expensive part of the run to waste — the search has already +paid for its exploration, and a round dispatched there may never finish. What a +cost-aware search does instead, and what this module provides, is **degrade +before the wall**: + +* at ``soft_floor`` (default 75%) of the budget, stop building fusion + tournaments. The tournament is an extra held-out sweep of the surviving + candidates and their fused candidate — pure ranking spend, and the ranking + can survive without it for the last few rounds; +* at ``hard_floor`` (default 90%), 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 a weight of ``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** tells the round barrier whether the *next* round fits the + remaining budget, so a run that cannot afford round N+1 ends at N with a + clean merge instead of being cut mid-round at the wall. + +The three thresholds are one design: every stage of a merge that is optional +ranking spend degrades before anything that decides a commit degrades, and +nothing that decides a commit ever degrades. The gate, the conflict resolution +and the Beta-posterior acceptance run to the last round exactly as they ran in +the first. + +All of this is inert without ``max_tokens``: no budget, no governor, no change +to the run. With a budget, the governor changes *how* the tail of the run is +spent, never whether the run's numbers mean anything. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Optional + + +#: Where the fusion tournament stops being built: 75% of the budget spent. +#: The tournament is the most expensive piece of optional ranking (a held-out +#: sweep of every survivor plus the fusion), and it is the piece whose absence +#: costs the least: candidates still commit through the same acceptance gate. +SOFT_FLOOR = 0.75 + +#: Where the self-verify rollout stops: 90% of the budget spent. This is the +#: last thing to go because it is the cheapest of the optional spends and the +#: one closest to the merge: its delta feeds the acceptance test's +#: ``observe_delta`` tie-breaker. Skipping it at 90% trades that tie-breaker +#: for the run actually finishing its round. +HARD_FLOOR = 0.90 + + +@dataclass +class BudgetGovernor: + """Tracks token spend against a cap and decides what to degrade. + + Constructed once per run (or not at all, when ``max_tokens`` is ``None``), + read at the points where optional spend is about to happen. All reads are + from the merger's thread or between rounds — the same single-threaded + contract the round barrier already has — so no lock is needed. + + ``spend`` is called by the engine with the meter's cumulative + ``prompt_tokens + completion_tokens`` once per round, before the barrier + checks; ``affords`` answers for the *next* round using the per-round + average of everything seen so far. + """ + + #: The hard cap, in tokens. ``None`` disables the governor entirely — + #: every method returns the "no budget" answer, which is the run's + #: behaviour without this class at all. + max_tokens: Optional[int] = None + #: Fraction of the budget at which fusion tournaments stop (see SOFT_FLOOR). + soft_floor: float = SOFT_FLOOR + #: Fraction of the budget at which self-verify stops (see HARD_FLOOR). + hard_floor: float = HARD_FLOOR + #: Stop when the run's *measured* return per token has fallen off its own + #: peak by more than :attr:`efficiency_floor`. Off by default: it is an + #: economic rule, not a safety one, and a run whose reward only ever rises + #: late would be cut short by it. Turn it on to spend a budget by value + #: rather than by exhaustion -- it is what lets a run finish *under* its + #: budget when the search has stopped paying for itself. + stop_on_diminishing: bool = False + #: How far below the run's own peak efficiency counts as diminishing. The + #: default ``0.25`` means "the last rounds returned less than a quarter of + #: the best rate this run ever achieved". Self-calibrating: it compares the + #: run against itself, so there is no magic absolute rate to set (a reward + #: is in ``[0, 1]`` and a token count is in the millions; their ratio has no + #: interpretable scale, which is exactly why an absolute floor would be a + #: guess). + efficiency_floor: float = 0.25 + #: How many efficiency samples before the rule may fire. Four rounds is the + #: floor at which "the peak" and "the recent rate" are different claims. + min_efficiency_samples: int = 4 + #: Cumulative spend, set by ``spend``. Kept as a field so ``affords`` can + #: be asked without re-passing the number the caller just passed. + _spent: int = 0 + #: Per-round deltas of the spend, for the projection in ``affords``. + _round_deltas: list = field(default_factory=list) + #: The spend at the last ``spend`` call, so a delta can be computed without + #: the caller tracking the previous value. + _last: int = 0 + #: Return-per-token for each completed round, from ``observe``. + _efficiencies: list = field(default_factory=list) + #: The previous round's reward and spend, so ``observe`` can difference. + _prev_reward: Optional[float] = None + _prev_spent: int = 0 + + # -- what the engine reports -------------------------------------------- + + def spend(self, tokens: int) -> None: + """Record the cumulative token spend as of this round's barrier. + + Called once per round, before the budget checks; the deltas between + calls are what ``affords`` projects on. A first call with ``0`` (a run + whose meter reports no tokens) is fine and leaves the projection + uninformative rather than wrong. + """ + self._round_deltas.append(max(0, tokens - self._last)) + self._last = self._spent = tokens + + # -- what the engine asks ------------------------------------------------ + + @property + def active(self) -> bool: + """Whether the governor has anything to say (a budget was set).""" + return self.max_tokens is not None + + def _remaining(self) -> Optional[float]: + if self.max_tokens is None: + return None + return self.max_tokens - self._spent + + def over_budget(self) -> bool: + """The hard stop: the spend has reached the cap.""" + if self.max_tokens is None: + return False + return self._spent >= self.max_tokens + + def allow_fusion_tournament(self) -> bool: + """Whether the extra held-out sweep of a tournament is affordable. + + The first thing to degrade: ranking precision, not a commit gate. + """ + if self.max_tokens is None: + return True + return self._spent < self.max_tokens * self.soft_floor + + def allow_self_verify(self) -> bool: + """Whether the double-cost verification rollout is affordable. + + The last thing to degrade: its delta feeds the acceptance test's + tie-breaker, so it survives the soft floor and only dies at the hard + one. + """ + if self.max_tokens is None: + return True + return self._spent < self.max_tokens * self.hard_floor + + def affords_next_round(self) -> bool: + """Whether the projection says one more round fits the budget. + + Projects the *average* per-round spend of the rounds seen so far. A + run with no history (nothing spent yet) always affords a round — a + budget that cannot afford round one is a ``ValueError`` the caller + should have raised, not a silent no-op — and a run whose meter reports + no tokens (``_spent == 0`` after a round) also always affords, because + the projection has nothing to project with and the token brake will + never fire anyway. + """ + if self.max_tokens is None: + return True + remaining = self._remaining() + if remaining <= 0: + return False + if not self._round_deltas or self._spent == 0: + return True + avg = sum(self._round_deltas) / len(self._round_deltas) + if avg <= 0: + return True + return remaining >= avg + + def remaining_fraction(self) -> float: + """Fraction of the token budget still unspent, in ``[0, 1]``. + + What a budget-aware :class:`~agentdescent.selection.SelectionPolicy` + anneals its exploration on. ``1.0`` when no budget was set -- a run + without a ceiling never runs short of one -- and clamped to ``[0, 1]`` + so an overshoot (the barrier's up-to-one-round lag) reads as "nothing + left" rather than a negative fraction that would invert the sign of an + exploration term. + """ + if self.max_tokens is None or self.max_tokens <= 0: + return 1.0 + return max(0.0, min(1.0, (self.max_tokens - self._spent) / self.max_tokens)) + + def observe(self, reward: Optional[float]) -> None: + """Record the reward of the round that just finished. + + Called at the top of each round, after :meth:`spend`, with the reward + the previous round produced (``None`` on the first round, and on a round + whose measurement failed). The efficiency it records is that round's + ``Δreward / Δtokens`` -- the run's own return per token, which + :meth:`diminishing_returns` compares against its own peak. + + A round that spent no tokens records nothing: ``Δreward / 0`` is not + infinity, it is a round whose cost was not measured, and letting it into + the sample would make an unmeasured run look infinitely efficient. + """ + if reward is None or not isinstance(reward, (int, float)): + return + reward = float(reward) + if self._prev_reward is None: + self._prev_reward = reward + self._prev_spent = self._spent + return + dtokens = self._spent - self._prev_spent + dreward = reward - self._prev_reward + self._prev_reward = reward + self._prev_spent = self._spent + if dtokens > 0: + self._efficiencies.append(dreward / dtokens) + + def diminishing_returns(self) -> bool: + """Whether the run's return per token has fallen off its own peak. + + The economic stop: keep buying compute while it is still paying, stop + when it is not -- *even with budget left*, which is the whole point. + ``max_tokens`` is a ceiling; this is the rule that decides not to spend + the ceiling. + + Returns ``False`` unless ``stop_on_diminishing`` is set, a budget is in + force (efficiency needs a token count) and at least + ``min_efficiency_samples`` rounds have been observed. It also returns + ``False`` when the run's peak efficiency is not positive: a run that + never improved has nothing to have declined *from*, and the patience + counter -- not this -- is the rule for "it never got better". + """ + if not self.stop_on_diminishing or self.max_tokens is None: + return False + samples = self._efficiencies + if len(samples) < self.min_efficiency_samples: + return False + peak = max(samples) + if peak <= 0: + return False + recent = sum(samples[-2:]) / 2.0 + return recent < self.efficiency_floor * peak + + # -- what the run reports -------------------------------------------------- + + def summary(self) -> dict: + """The governor's state, for the result and the status line. + + Always a dict (possibly all-``None``), so a caller does not have to + feature-detect — an inactive governor reports ``None`` fields, which + is the honest answer to "how much of the budget is left" when there + is no budget. + """ + if self.max_tokens is None: + return {"max_tokens": None, "spent": self._spent, + "remaining": None, "fusion_degraded": False, + "self_verify_degraded": False} + return { + "max_tokens": self.max_tokens, + "spent": self._spent, + "remaining": max(0, self.max_tokens - self._spent), + "fusion_degraded": not self.allow_fusion_tournament(), + "self_verify_degraded": not self.allow_self_verify(), + } + + +# --------------------------------------------------------------------------- +# Adaptive per-call thinking budget (o1-style test-time scaling) +# --------------------------------------------------------------------------- + + +@dataclass +class CallBudget: + """A mutable per-call max_tokens, set by the engine before each expansion. + + The engine's :class:`BudgetGovernor` decides how many tokens the *next* + model call should be allowed to spend, and writes that number here. A + :func:`budgeted_completion` wrapper reads it and delegates to the underlying + adapter with the right ``max_tokens``. + + The o1 insight, applied to search: a reasoning model's per-call thinking + budget should be **proportional to the expected value of that expansion** — + a promising, under-explored parent deserves more room to think than a + dead-end one the search has already tried five times. The engine knows the + parent's score and the remaining budget; the adapter knows how to call the + model; this object is the one place they meet. + + ``base`` is the default the adapter was constructed with. The engine never + raises ``next`` above ``base`` (a model call that costs more than the + caller's configured ceiling is a bill the caller did not agree to), and + never drops it below ``base // 4`` (a reasoning model starved below a + quarter of its budget returns empty content — measured, see + :func:`agentdescent.agents.claude`). + """ + + #: The adapter's configured max_tokens — the ceiling the engine never + #: exceeds. + base: int = 4096 + #: What the next call should use. Set by the engine before each expansion; + #: ``None`` means "use ``base``" (the adapter's own default), which is the + #: state a run without adaptive budgeting stays in. + next: Optional[int] = None + #: The floor: a reasoning model below this returns nothing. Measured on + #: deepseek-v4-flash at 1024: four of eight reflection prompts came back + #: empty. ``base // 4`` is the conservative reading of that. + floor: Optional[int] = None + + def __post_init__(self) -> None: + if self.floor is None: + self.floor = max(256, self.base // 4) + + def allocate(self, score: float, budget_remaining: float, + *, max_boost: float = 1.0) -> None: + """Set ``next`` for the upcoming expansion. + + The allocation is a product of two factors: + + * **Score factor** (``0.5 + 0.5 * score``): a high-scoring parent + gets the full ``base``; a zero-scoring one gets half. The parent's + score is the search's own estimate of how promising it is, so this + is "spend more thinking where the search says it is worth it." + + * **Budget factor** (``0.5 + 0.5 * budget_remaining``): early in the + run, when there is budget to exploit what exploration finds, the + engine allows full thinking; late, when there is not, it tightens. + This is the same annealing :class:`CostEfficient` uses on the + exploration bonus, applied to the thinking budget — both are "spend + more while you still have time to act on it." + + The product ranges from ``base * 0.25`` (dead-end parent, no budget + left) to ``base`` (promising parent, full budget), never above ``base`` + and never below ``floor``. ``max_boost`` > 1 is allowed but only when + the caller explicitly sets it — a budget that lets the engine raise + the per-call ceiling above the adapter's configured ``max_tokens`` is + a bill the caller has to opt into. + """ + s = 0.5 + 0.5 * max(0.0, min(1.0, float(score))) + b = 0.5 + 0.5 * max(0.0, min(1.0, float(budget_remaining))) + ceiling = int(self.base * max_boost) + self.next = max(self.floor, min(ceiling, int(self.base * s * b))) + + +def budgeted_completion( + factory: Callable[[int], "Completion"], + budget: CallBudget, +) -> "Completion": + """Wrap a ``max_tokens -> Completion`` factory into an adaptive Completion. + + The factory is called once per distinct ``max_tokens`` value (cached), so + a run that stays at the default pays one construction and a run that + adapts pays one per step it adapts to — not one per call. + + Usage:: + + from agentdescent.budget import CallBudget, budgeted_completion + from agentdescent.agents import claude, LLMAgent + + budget = CallBudget(base=8192) + agent = LLMAgent(budgeted_completion( + lambda mt: claude(model="claude-sonnet-4-5", max_tokens=mt, usage=u), + budget)) + + evolve(tasks, reward, agent=agent, max_tokens=500_000, + call_budget=budget, usage=u) + + Without ``call_budget`` the engine never touches ``budget.next`` and the + factory is called once with ``budget.base`` — the same cost as constructing + ``claude(max_tokens=8192)`` directly. + """ + + cache: dict[int, "Completion"] = {} + + def complete(prompt: str) -> str: + mt = budget.next if budget.next is not None else budget.base + if mt not in cache: + cache[mt] = factory(mt) + return cache[mt](prompt) + + return complete diff --git a/agentdescent/cli.py b/agentdescent/cli.py index d1b43b39..5b3b694e 100644 --- a/agentdescent/cli.py +++ b/agentdescent/cli.py @@ -64,9 +64,10 @@ def _status_line(st: runstore.RunStatus) -> str: reward = "" if st.best_reward is None else f" best={st.best_reward:.3f}" rounds = f"{st.round}/{st.rounds}" if st.rounds else f"{st.round}" usd = "" if st.usd is None else f" ${st.usd:.2f}" + tok = f" tok={st.tokens:,}" if st.tokens else "" tail = f" [{st.stop_reason}]" if st.stop_reason else "" err = f" error: {st.error}" if st.error else "" - return (f"{st.run_id} {st.state:<9} round {rounds:<6} calls={st.calls}{reward}{usd}" + return (f"{st.run_id} {st.state:<9} round {rounds:<6} calls={st.calls}{tok}{reward}{usd}" f" {st.kind or '?'}:{st.target or '?'}{tail}{err}") diff --git a/agentdescent/evolution.py b/agentdescent/evolution.py index 2e264bd2..a099baed 100644 --- a/agentdescent/evolution.py +++ b/agentdescent/evolution.py @@ -59,6 +59,7 @@ from .selection import SingleHead from .scheduler import AuditScheduler from .staleness import StalenessPolicy +from .budget import BudgetGovernor, CallBudget # --------------------------------------------------------------------------- @@ -926,6 +927,26 @@ def _wants_population(policy) -> bool: return policy is not None and type(policy) is not SingleHead +def _set_budget_remaining(aggregator: Any, governor: BudgetGovernor) -> None: + """Offer the governor's remaining budget fraction to the aggregator. + + A :class:`~agentdescent.population.PopulationAggregator` carries a + ``budget_remaining`` slot its selection context reads, so a + :class:`~agentdescent.selection.CostEfficient` policy can anneal exploration + as the token budget is spent. A plain ``Aggregator`` -- and any custom one -- + has no such slot and is left untouched, which is the correct default: a + selection policy that never asked for a budget must not start seeing one. + """ + if not governor.active: + return + # Only an aggregator that declares the slot gets it. A plain `Aggregator` + # must not gain an attribute nothing reads, and a `PopulationAggregator` + # always has it (set in `__init__`), so `hasattr` is the right test. + if not hasattr(aggregator, "budget_remaining"): + return + aggregator.budget_remaining = governor.remaining_fraction() + + def _cost_fields(meter: Meter) -> Dict[str, Any]: """The meter's counters, keyed as :class:`EvolutionResult` fields. @@ -1000,6 +1021,11 @@ class RoundInfo: #: comparison has to be able to ask "where was each configuration after N #: calls", which needs the number at every round rather than only at the end. calls: int = 0 + #: Tokens consumed by the end of this round, cumulative + #: (``prompt_tokens + completion_tokens``). A reasoning model can spend 40k + #: tokens on a single hidden-thinking call, so the total is the only budget + #: that maps to cost on a real deployment. + tokens: int = 0 # -- what the merge did, this round -------------------------------------- # @@ -1203,6 +1229,7 @@ class EvolutionResult: error: Optional[str] = None #: Why the run ended -- ``"target_reward"`` / ``"patience"`` / ``"rounds"`` / #: ``"max_seconds"`` / ``"max_iters"`` / ``"max_rollouts"`` / ``"max_calls"`` / + #: ``"max_tokens"`` / #: ``"stop_when"`` (the caller's own budget, see ``evolve(stop_when=)``) #: / ``"error"``. Without it a budget #: expiry is indistinguishable from convergence: ``error`` is ``None`` for @@ -1299,6 +1326,14 @@ class EvolutionResult: #: Every fusion tournament the run held, when the fusion policy recorded #: them (the shipped one does). Read it through :meth:`fusion_stats`. fusion_trials: List["FusionTrial"] = field(default_factory=list) + #: What the budget governor did, when a ``max_tokens`` cap was set. ``None`` + #: when no budget was given (the run ran without a governor at all). The + #: dict is the same shape :meth:`~agentdescent.budget.BudgetGovernor.summary` + #: returns: ``max_tokens``, ``spent``, ``remaining``, ``fusion_degraded``, + #: ``self_verify_degraded``. Read it to tell "the run finished because its + #: budget ran out" from "it converged" — and to know whether the tail of + #: the search ran degraded (fewer ranking passes) or at full quality. + budget: Optional[Dict[str, Any]] = None def fusion_stats(self) -> "FusionStats": """How often merging beat the best single diff -- and how badly it lost. @@ -1448,7 +1483,8 @@ def save(self, path: str) -> None: "n_items": h.n_items, "committed": h.committed, "rejected": h.rejected, "reasons": h.reasons, "elapsed_s": h.elapsed_s, "rollouts": h.rollouts, - "calls": h.calls, "considered": h.considered, + "calls": h.calls, "tokens": h.tokens, + "considered": h.considered, "discarded_stale": h.discarded_stale, "conflicts_dropped": h.conflicts_dropped, "fused": h.fused} for h in self.history @@ -1704,6 +1740,7 @@ def record_round(self, *, index: int, reward: float, n_items: int, info = RoundInfo( index, reward, n_items, committed, len(reports) - committed, reasons, elapsed_s=m.elapsed_s, rollouts=m.rollouts, calls=m.calls, + tokens=m.prompt_tokens + m.completion_tokens, considered=sum(x.considered for x in reports), discarded_stale=sum(x.discarded_stale for x in reports), conflicts_dropped=sum(x.conflicts_dropped for x in reports), @@ -1889,6 +1926,18 @@ def _build_engine(tasks, reward, *, agent, run, propose, strategy, initial_state # Wrap the actors before anything else can capture them: `run` is closed over # by the runtime, by every worker and by the verifier's `eval_fn`, so a later # wrap would miss whichever reference was taken first. + # + # The `usage` the meter reads defaults to the agent's own, when the agent + # has one and the caller did not pass `usage=`: the adapters that report + # tokens (`claude`, `openai_compatible`) take a `Usage` in their constructor, + # and the caller who built `LLMAgent(claude(usage=u))` has already told the + # model where to report. Forcing them to pass `u` again here — the only way + # the meter (and with it `max_tokens`, and the `tokens=` column on every + # round) would see the same numbers — was a foot-gun the whole token budget + # silently missed: the run reported `tokens=0` while the bill went to + # `u`, an object only the caller held. + if usage is None: + usage = getattr(agent, "usage", None) meter = Meter(usage=usage) if usage is not None else Meter() run, propose = measured(run, meter), measured(propose, meter) @@ -2180,6 +2229,44 @@ def evolve( max_seconds: Optional[float] = None, max_rollouts: Optional[int] = None, max_calls: Optional[int] = None, + #: Hard cap on total tokens consumed (``prompt + completion``). A reasoning + #: model can spend 40k tokens on hidden thinking in one call, so ``max_calls`` + #: and ``max_rollouts`` do not bound cost: a 20-round run with 4 workers is + #: 80 rollouts and ~160 model calls, but at 40k tokens each that is 6.4M + #: tokens -- the bill, not the count, is what a deployment needs to control. + #: Checked at the round barrier alongside the other budgets; the run stops + #: with ``stop_reason="max_tokens"`` and reports the spend it incurred. + max_tokens: Optional[int] = None, + #: Stop when the run's *measured* return per token has fallen off its own + #: peak by more than ``efficiency_floor`` (see + #: :meth:`~agentdescent.budget.BudgetGovernor.diminishing_returns`). The + #: economic stop: keep buying compute while it pays, stop when it does not, + #: *even with budget left*. Off by default -- it is a spending rule, not a + #: safety one, and a run whose reward only rises late would be cut short. + #: Needs ``max_tokens`` to be set (efficiency needs a cost to divide by). + stop_on_diminishing_returns: bool = False, + #: How far below the peak counts as diminishing, when the stop above is on. + #: Self-calibrating against the run's own best rate, so this is a *ratio* + #: and not an absolute: a reward is in ``[0, 1]`` and a token count in the + #: millions, and their quotient has no interpretable scale. + efficiency_floor: float = 0.25, + #: An adaptive per-call thinking budget (o1-style test-time scaling). + #: Pass a :class:`~agentdescent.budget.CallBudget` whose ``base`` matches + #: the adapter's configured ``max_tokens``, and the engine will adjust the + #: per-call ceiling based on each parent's score and the remaining budget — + #: promising parents get more thinking, dead-end ones get less, and the + #: allocation tightens as the token budget is spent. The adapter must be + #: wrapped with :func:`~agentdescent.budget.budgeted_completion` for this to + #: reach it. ``None`` (default) means the per-call ceiling is the adapter's + #: own and never changes. + call_budget: Optional["CallBudget"] = None, + #: An adaptive per-call thinking budget (o1-style test-time scaling). Pass + #: a :class:`~agentdescent.budget.CallBudget` whose ``base`` matches the + #: adapter's configured ``max_tokens``; the engine adjusts the per-call + #: ceiling based on each parent's score and the remaining budget — promising + #: parents get more thinking, dead-ends get less. The adapter must be + #: wrapped with :func:`~agentdescent.budget.budgeted_completion`. ``None`` + #: (default) means the per-call ceiling never changes. self_verify: bool = True, held_out_frac: float = 0.4, repo_path: Optional[str] = None, @@ -2384,6 +2471,38 @@ def evolve( The async path has no barrier and enforces both per rollout, so it overshoots by at most the rollouts already in flight. + max_tokens: + Hard cap on total tokens consumed (``prompt_tokens + completion_tokens``). + A reasoning model can spend 40k tokens on hidden thinking in one call, + so ``max_calls`` and ``max_rollouts`` do not bound cost: a 20-round run + with 4 workers is 80 rollouts and ~160 model calls, but at 40k tokens + each that is 6.4M tokens -- the bill, not the count, is what a + deployment needs to control. Checked at the round barrier alongside the + other budgets; the async path checks per-rollout for tighter control. + Stops with ``stop_reason="max_tokens"``. ``None`` (default) means + unbounded. + stop_on_diminishing_returns: + Stop when the run's own return per token has fallen off its peak by more + than ``efficiency_floor``. The economic rule: keep buying compute while + it pays, stop when it does not -- *even with budget left*. Off by + default, because a run whose reward only rises late would be cut short + by it, and because it needs ``max_tokens`` (return per token needs a cost + to divide by). Stops with ``stop_reason="diminishing_returns"``. + efficiency_floor: + How far below the peak counts as diminishing, when the stop above is on. + Self-calibrating against the run's own best rate, so it is a *ratio*, not + an absolute quantity: a reward is in ``[0, 1]`` and a token count is in + the millions, and their quotient has no interpretable scale. + call_budget: + An adaptive per-call thinking budget (o1-style test-time scaling). + Pass a :class:`~agentdescent.budget.CallBudget` whose ``base`` matches + the adapter's configured ``max_tokens``, and the engine adjusts the + per-call ceiling based on each parent's score and the remaining budget + — promising parents get more thinking, dead-end ones get less, and the + allocation tightens as the token budget is spent. The adapter must be + wrapped with :func:`~agentdescent.budget.budgeted_completion` for this + to reach it. ``None`` (default) means the per-call ceiling is the + adapter's own and never changes. self_verify: Re-run the trajectory with the diff applied to record a local before/after delta. Doubles the rollouts spent per proposal; ports that @@ -2564,7 +2683,7 @@ def evolve( max_seconds=20.0 if max_seconds is None else max_seconds, max_iters=(max_rollouts if max_rollouts is not None else rounds * max(1, n_workers)), - max_calls=max_calls, held_out_frac=held_out_frac, + max_calls=max_calls, max_tokens=max_tokens, held_out_frac=held_out_frac, repo_path=repo_path, agg_config=agg_config, staleness_policy=staleness_policy, aggregator_factory=aggregator_factory, oracle_budget=oracle_budget, cheap_eval_tasks=cheap_eval_tasks, fusion_tournament=fusion_tournament, @@ -2576,7 +2695,10 @@ def evolve( eval_concurrency=eval_concurrency, pipelined_gate=pipelined_gate, gate_workers=gate_workers, on_round=on_round, stop_when=stop_when, verbose=verbose, usage=usage, - policies=policies, checkpointing=checkpointing) + policies=policies, checkpointing=checkpointing, + stop_on_diminishing_returns=stop_on_diminishing_returns, + efficiency_floor=efficiency_floor, + call_budget=call_budget) if pipelined_gate: # The mirror of the block above, and the same reasoning: a knob accepted @@ -2674,6 +2796,13 @@ def evolve( dead_rounds = 0 # consecutive rounds where every worker failed deadline = time.time() + max_seconds if max_seconds else None stop_reason = "rounds" + # The cost-aware governor: degrades optional spend (fusion tournaments, + # self-verify) before the token wall, and projects whether the *next* + # round fits — so a run that cannot afford round N+1 ends at N with a + # clean merge instead of being cut mid-round. Inert without max_tokens. + governor = BudgetGovernor(max_tokens=max_tokens, + stop_on_diminishing=stop_on_diminishing_returns, + efficiency_floor=efficiency_floor) for r in range(rounds): if deadline is not None and time.time() >= deadline: stop_reason = "max_seconds" @@ -2686,16 +2815,74 @@ def evolve( # produced -- so the run overshoots by up to a round and reports the # spend it actually incurred instead of the one it was asked for. spent = eng.meter.snapshot() + tokens_spent = spent.prompt_tokens + spent.completion_tokens + governor.spend(tokens_spent) + # Record the *previous* round's reward so the governor can measure the + # run's own return per token. `history[-1]` is that round: this runs at + # the top of round r, after r-1 has closed. `None` on the first round. + governor.observe(history[-1].held_out_reward if history else None) over = ((max_rollouts is not None and spent.rollouts >= max_rollouts and "max_rollouts") or (max_calls is not None and spent.calls >= max_calls - and "max_calls")) + and "max_calls") or + (max_tokens is not None and tokens_spent >= max_tokens + and "max_tokens")) + # Tell a budget-aware selection policy how much of the budget is left, so + # it can anneal exploration before the round body selects a parent. A + # plain `Aggregator` has no such slot and is left alone; a + # `PopulationAggregator` reads it in `step()`'s SelectionContext. + _set_budget_remaining(aggregator, governor) if over: stop_reason = over if verbose: print(f"round {r:>3} stopping: {over} reached " - f"({spent.rollouts} rollouts / {spent.calls} calls)") + f"({spent.rollouts} rollouts / {spent.calls} calls " + f"/ {tokens_spent} tokens)") + break + # The economic stop: the run's own return per token has fallen off its + # peak. Opt-in (``stop_on_diminishing_returns``), and only when a token + # budget is in force -- efficiency needs a cost to divide by. This is + # the rule that decides *not* to spend the ceiling, which is the whole + # difference between a budget and an allocation. + if governor.diminishing_returns(): + stop_reason = "diminishing_returns" + if verbose: + print(f"round {r:>3} stopping: diminishing returns " + f"(return per token off its peak; " + f"{tokens_spent:,} tokens spent)") break + # The projection: a round that the remaining budget cannot afford is + # not dispatched. This lands *before* the snapshot, so the run ends on + # the last round it fully paid for rather than one round into the wall + # — the difference between a clean merge and a partial one. Reported + # as "max_tokens" (the budget that fired), not a separate reason: the + # caller's question is "why did my run stop early" and the answer is + # the token budget either way. + if not governor.affords_next_round(): + stop_reason = "max_tokens" + if verbose: + s = governor.summary() + print(f"round {r:>3} stopping: projected next round exceeds " + f"max_tokens ({s['spent']:,} spent, ~" + f"{s['remaining']:,} left)") + break + # The governor's soft degradation: at the soft floor (75% of budget), + # turn off the fusion tournament — it is the most expensive piece of + # optional ranking spend, and its absence costs selection refinement, + # never a commit gate. Mutating the fusion policy's ``tournament`` + # flag at the round barrier is safe: ``_tournament`` reads it in + # ``step()``, which runs *after* this point in the same round. + if governor.active: + fp = getattr(aggregator, "fusion_policy", None) + if fp is not None and hasattr(fp, "tournament"): + should = governor.allow_fusion_tournament() + if fp.tournament != should: + fp.tournament = should + if verbose: + tag = "on" if should else "off" + print(f"round {r:>3} fusion tournament {tag} " + f"(token budget: {governor.summary()['spent']:,}/" + f"{governor.max_tokens:,})") try: snap = ledger.snapshot(Ledger.DEV) except LedgerFailure as e: @@ -2817,6 +3004,14 @@ def _run_unit_inner(unit) -> None: health.record_success() if score >= solved_threshold: return + # o1-style test-time scaling: if an adaptive call budget is + # installed, set its per-call ceiling for this expansion based on + # the parent's score and the remaining token budget. A promising + # parent gets more thinking; a dead-end one gets less; the + # allocation tightens as the budget is spent. Inert without + # ``call_budget`` — a plain ``propose`` never sees it. + if call_budget is not None: + call_budget.allocate(score, governor.remaining_fraction()) proposal = _checked_proposal( propose(mine.render(), task, output, score), task) if not proposal: @@ -2844,8 +3039,12 @@ def _run_unit_inner(unit) -> None: f"{outside[0]!r}, outside its section {unit.section}") return # The self-verify rollout doubles the cost of every proposal, so it is - # opt-out here exactly as it is on the async path. - if self_verify: + # opt-out here exactly as it is on the async path. The governor can + # also turn it off mid-run when the token budget is nearly spent: + # the delta it produces feeds the acceptance test's tie-breaker + # (``observe_delta``), never the commit gates, so skipping it costs + # ranking precision on the advantage signal and nothing else. + if self_verify and governor.allow_self_verify(): after = _checked_reward( reward(task, run(mine.apply(diff).render(), task)), task) delta = after - score @@ -3002,7 +3201,8 @@ def _bounded(u=None): if verbose: print(f"round {r:>3} reward={info.held_out_reward:.3f} on " f"{len(held_out)} size={info.n_items} " - f"+{info.committed}/-{info.rejected}") + f"+{info.committed}/-{info.rejected}" + + (f" tokens={info.tokens:,}" if info.tokens else "")) if early_stop is not None: stop_reason = early_stop if verbose: @@ -3053,6 +3253,7 @@ def _bounded(u=None): ledger_log=_safe_log(ledger), error=run_error, stop_reason="error" if run_error else stop_reason, fusion_trials=_fusion_trials(aggregator), + budget=governor.summary() if governor.active else None, **_cost_fields(eng.meter)) eng.cleanup() return result diff --git a/agentdescent/population.py b/agentdescent/population.py index 05965737..a0890cb4 100644 --- a/agentdescent/population.py +++ b/agentdescent/population.py @@ -54,6 +54,21 @@ __all__ = ["PopulationAggregator", "population_factory"] +def _state_size(state: Dict[str, str]) -> int: + """A candidate's content size, as a deterministic cost proxy. + + The prompt an expansion sends is dominated by the artifact it renders, so + content size is a first-order proxy for the input tokens that expansion + costs -- and unlike measured tokens it is attributable to a single candidate + without racing a shared meter. A cost-aware policy divides its exploration + bonus by this, so a cheap candidate gets more speculative looks than an + expensive one at the same rank. It is a *proxy* and the policy treats a + missing cost as the pool mean, so an imperfect scale cannot turn an + unmeasured candidate into a free one. + """ + return sum(len(str(k)) + len(str(v)) for k, v in state.items()) + + class PopulationAggregator(Aggregator): """The shipped merge pipeline plus an archive and a selection policy.""" @@ -83,6 +98,12 @@ def __init__(self, ledger, verifier, audit, config, staleness_policy=None, #: the archive lock: `step()` is the merger's, one thread, while the #: lock guards the archive against the workers that `ingest`. self._selections = 0 + #: Fraction of the run's budget still unspent, in ``[0, 1]``, for the + #: ``SelectionContext`` a budget-aware policy anneals on. The engine sets + #: it each round (from its governor); ``1.0`` is the no-budget default, + #: which is what a run without a token ceiling should report -- it never + #: runs short of one. + self.budget_remaining: float = 1.0 # -- the archive --------------------------------------------------------- @@ -110,7 +131,17 @@ def _candidates(self) -> List[Candidate]: version=int(entry["version"]), state=dict(entry["state"]), score=float(entry["score"]), - selected=int(entry["selected"])) + selected=int(entry["selected"]), + # A deterministic cost proxy: the artifact's rendered + # content size, which dominates the prompt and so the + # input tokens of expanding it. It is not the measured + # token spend -- that is not attributable to a single + # parent from a shared meter under concurrency -- and + # the policy treats an unmeasured cost as the pool + # mean rather than zero, so a proxy that is wrong in + # scale is far less harmful than one that is wrong in + # *presence*. + cost=float(_state_size(entry["state"]))) for entry in self._archive ] @@ -211,7 +242,8 @@ def step(self) -> List[MergeReport]: # the same entry forever. `Beam(4)` was `Beam(1)`, and `ParetoFrontier` # sat on whichever front member was admitted first, usually the seed. ctx = SelectionContext(head=head_candidate, candidates=tuple(candidates), - round=self._selections, n_workers=1) + round=self._selections, n_workers=1, + budget_remaining=self.budget_remaining) self._selections += 1 chosen = list(self.selection.select(ctx, 1)) if not chosen: diff --git a/agentdescent/runstore.py b/agentdescent/runstore.py index 2e8dc56f..60651bfe 100644 --- a/agentdescent/runstore.py +++ b/agentdescent/runstore.py @@ -98,6 +98,11 @@ class RunStatus: committed: int = 0 calls: int = 0 rollouts: int = 0 + #: Total tokens consumed (prompt + completion), when the agent reports + #: them. 0 for an opaque `run`/`propose` with no `usage=` shared — which is + # why it is an int and not Optional: "not reported" and "reported zero" + # read the same here, and neither is a state a person acts on. + tokens: int = 0 usd: Optional[float] = None pid: Optional[int] = None started: Optional[float] = None @@ -460,6 +465,9 @@ def on_round(info: RoundInfo) -> None: # A number in `status` is read by a person deciding whether to cancel. changes["committed"] = st.committed + info.committed changes["rollouts"] = info.rollouts + # RoundInfo.tokens is the same running-total shape as rollouts; the + # meter sums prompt + completion, so this is the spend so far. + changes["tokens"] = info.tokens if usd_per_call is not None: changes["usd"] = round(usage.calls * usd_per_call, 4) rd.update_status(**changes) diff --git a/agentdescent/selection.py b/agentdescent/selection.py index 963a2762..f3edafce 100644 --- a/agentdescent/selection.py +++ b/agentdescent/selection.py @@ -129,6 +129,15 @@ class Candidate: #: says how good it is, and the two are different claims about a candidate #: that is slow today because it is a first draft of the right idea. prior: Optional[float] = None + #: What it costs to expand this candidate, in the caller's unit -- tokens on + #: a token-budgeted run, wall-clock seconds otherwise, or whatever a custom + #: aggregator measures. ``None`` means *unknown*, not zero: the same + #: distinction ``score`` and ``prior`` make, and for the same reason. A + #: policy that ranks by score-per-cost must not read a missing cost as + #: "free" -- that would make an unmeasured candidate the most attractive + #: one on the board. A policy that cannot measure cost leaves this ``None`` + #: and does not read it. + cost: Optional[float] = None @dataclass(frozen=True) @@ -147,6 +156,15 @@ class SelectionContext: candidates: Sequence[Candidate] = () round: int = 0 n_workers: int = 1 + #: Fraction of the run's budget still unspent, in ``[0, 1]``. ``1.0`` when the + #: caller set no budget (a run with no ceiling never runs short of one), and + #: what a budget-aware policy anneals its exploration by: a placement made + #: with 10% left has less time to pay off than one made at the start, so the + #: exploration it buys is worth proportionally less. This is the o1-style + #: test-time-scaling decision expressed as a search parameter -- spend to + #: explore while there is budget to exploit what you find, and stop paying + #: for exploration once there is not. + budget_remaining: float = 1.0 @runtime_checkable @@ -630,3 +648,105 @@ def select(self, ctx: SelectionContext, n: int) -> Sequence[Candidate]: # to the same arm would make a batch worth one rollout of information. chosen = ranked[:n] or [ctx.head] return [chosen[i % len(chosen)] for i in range(n)] + + +class CostEfficient(FlatPuct): + """FlatPuct that pays for exploration out of a budget. **Test-time scaling.** + + The o1 insight, applied to search: *a token is worth more where the search + still has time to act on what it buys.* o1 spends more thinking on the + problem it is unsure of; this spends the next expansion on the candidate + whose exploration buys the most quality per unit of remaining budget. The + formula keeps :class:`FlatPuct`'s rank exploitation and adds two things: + + * **Cost discounts exploration, not exploitation.** The exploration term is + divided by the candidate's cost, so a speculative expansion of an + expensive candidate has to beat a speculative expansion of a cheap one. + The exploitation term is left alone: a candidate the ranking already + places at the top is grounded, and cost must not disqualify it -- charging + cost against exploitation would bias the search against complex solutions + (whose artifacts are longer and costlier to expand) in favour of short + ones, which is not a quality judgment and must not be smuggled in as one. + + * **Exploration anneals with the remaining budget.** The bonus is scaled by + ``ctx.budget_remaining``: at the start it is ``1.0`` and the policy + explores as FlatPuct does, and as the budget is spent the bonus shrinks + toward zero, so the last expansions exploit what the run actually found. + This is what an anytime search must do -- an exploration made with 5% of + the budget left cannot pay off before the run ends -- and it is the + difference between spending a budget and *allocating* it. + + ``cost_exponent`` (``alpha``) sets how hard cost is charged: ``0`` ignores it + and is FlatPuct exactly, ``1`` is full value-per-cost. The default is ``1``. + ``anneal=False`` pins the exploration bonus at full strength, which is the + ablation that separates "cost-aware" from "budget-aware". + + **Unknown cost is the mean, not zero.** A candidate whose cost was never + measured is charged the pool's mean, for the same reason :attr:`Candidate.prior` + falls back to the mean: reading a missing number as ``0`` would make every + unmeasured candidate free to explore, which is exactly backwards. With no + costs measured at all the divisor is ``1`` for every candidate and this + policy is :class:`FlatPuct` to the floating-point bit — a property the + tests assert, so the mechanism cannot silently change a cost-blind run. + """ + + def __init__(self, c_puct: float = 1.0, cost_exponent: float = 1.0, + prior_exponent: float = 0.0, anneal: bool = True) -> None: + super().__init__(c_puct=c_puct, prior_exponent=prior_exponent) + self.cost_exponent = max(0.0, float(cost_exponent)) + self.anneal = anneal + + def cost_factors(self, rows: Sequence[Candidate]) -> List[float]: + """``cost**alpha`` per row, normalised so the mean is 1. + + Normalising by the mean is what keeps the divisor scale-free: the + exploration term is a rank-scale quantity and dividing it by a raw token + count (thousands) would collapse exploration to zero. Dividing by the + *mean cost* means an average-cost candidate keeps the FlatPuct divisor + of 1, a cheaper one gets more exploration, a costlier one less -- and a + pool with one uniform cost is unchanged. + """ + if self.cost_exponent == 0.0: + return [1.0] * len(rows) + raw = [c.cost if c.cost is not None and c.cost > 0 else None for c in rows] + known = [x for x in raw if x is not None] + mean = (sum(known) / len(known)) if known else 1.0 + effective = [x if x is not None else mean for x in raw] + factors = [max(1e-9, (x / mean)) ** self.cost_exponent for x in effective] + return factors + + def select(self, ctx: SelectionContext, n: int) -> Sequence[Candidate]: + rows = list(ctx.candidates) + if len(rows) <= 1: + return super().select(ctx, n) + ranks = self._rank_scores(rows) + visits = [c.selected for c in rows] + priors = self._priors(rows) + costs = self.cost_factors(rows) + # The annealing factor is the remaining budget when asked to anneal, and + # 1.0 otherwise. `budget_remaining` is clamped: a caller that passes a + # negative fraction (a bug) must not invert the exploration term. + scale = (max(0.0, min(1.0, ctx.budget_remaining)) if self.anneal else 1.0) + by_version = {c.version: i for i, c in enumerate(rows)} + picked: List[Candidate] = [] + for _ in range(n): + total = sum(visits) + best, best_puct = 0, -math.inf + for i, row in enumerate(rows): + explore = (self.c_puct * priors[i] + * math.sqrt(total) / (1 + visits[i]) + * scale / costs[i]) + puct = ranks[i] + explore + if puct > best_puct: + best, best_puct = i, puct + picked.append(rows[best]) + # Same ancestor reservation FlatPuct does, so a run of this policy is + # one rule with a cost term rather than a second selection rule. + seen: Set[int] = set() + node: Optional[int] = best + while node is not None and node not in seen: + seen.add(node) + visits[node] += 1 + parent = rows[node].parent + node = by_version.get(parent) if parent is not None else None + return picked diff --git a/docs/api.md b/docs/api.md index 0b938ef5..967952d4 100644 --- a/docs/api.md +++ b/docs/api.md @@ -13,7 +13,7 @@ means the parameter has none. Each section links to the page that explains *why* the module is shaped the way it is; this page is the *what*. -332 public names across 54 modules. +336 public names across 55 modules. --- @@ -61,7 +61,8 @@ EvolutionResult( sandboxes_created: int = 0, sandboxes_reused: int = 0, sandbox_failures: int = 0, - fusion_trials: List['FusionTrial'] = + fusion_trials: List['FusionTrial'] = , + budget: Optional[Dict[str, Any]] = None ) -> None ``` @@ -170,6 +171,7 @@ RoundInfo( elapsed_s: float = 0.0, rollouts: int = 0, calls: int = 0, + tokens: int = 0, considered: int = 0, discarded_stale: int = 0, conflicts_dropped: int = 0, @@ -220,6 +222,10 @@ evolve( max_seconds: Optional[float] = None, max_rollouts: Optional[int] = None, max_calls: Optional[int] = None, + max_tokens: Optional[int] = None, + stop_on_diminishing_returns: bool = False, + efficiency_floor: float = 0.25, + call_budget: Optional['CallBudget'] = None, self_verify: bool = True, held_out_frac: float = 0.4, repo_path: Optional[str] = None, @@ -271,6 +277,10 @@ evolve( | `max_seconds` | `Optional[float]` | `None` | Wall-clock budget. `None` (default) means unbounded; the async path uses `20.0` when unset. | | `max_rollouts` | `Optional[int]` | `None` | The budget in the two units a comparison has to hold fixed: rollouts completed, and actor invocations (`run` + `propose`). `rounds` is not one of them -- configurations differ in how much model a round buys, so a budget fixed in rounds hands the wider configuration more model and then reports the extra model as a win for parallelism. Either bound stops the run with `stop_reason` `"max_rollouts"` / `"max_calls"`. **Checked at the round barrier, so a run overshoots by up to one round.** A round is dispatched or it is not; stopping halfway would leave a half-merged round, and the states a comparison compares are the ones a merge produced. So a budget is a *bound on where to stop*, never the number to compare on: read the spend the run actually reported (`result.rollouts`, `result.usage.calls`), which is what `baselines` does -- it refuses to call two arms equal-budget when their measured spends differ. The async path has no barrier and enforces both per rollout, so it overshoots by at most the rollouts already in flight. | | `max_calls` | `Optional[int]` | `None` | As `max_rollouts`. | +| `max_tokens` | `Optional[int]` | `None` | Hard cap on total tokens consumed (`prompt_tokens + completion_tokens`). A reasoning model can spend 40k tokens on hidden thinking in one call, so `max_calls` and `max_rollouts` do not bound cost: a 20-round run with 4 workers is 80 rollouts and ~160 model calls, but at 40k tokens each that is 6.4M tokens -- the bill, not the count, is what a deployment needs to control. Checked at the round barrier alongside the other budgets; the async path checks per-rollout for tighter control. Stops with `stop_reason="max_tokens"`. `None` (default) means unbounded. | +| `stop_on_diminishing_returns` | `bool` | `False` | Stop when the run's own return per token has fallen off its peak by more than `efficiency_floor`. The economic rule: keep buying compute while it pays, stop when it does not -- *even with budget left*. Off by default, because a run whose reward only rises late would be cut short by it, and because it needs `max_tokens` (return per token needs a cost to divide by). Stops with `stop_reason="diminishing_returns"`. | +| `efficiency_floor` | `float` | `0.25` | How far below the peak counts as diminishing, when the stop above is on. Self-calibrating against the run's own best rate, so it is a *ratio*, not an absolute quantity: a reward is in `[0, 1]` and a token count is in the millions, and their quotient has no interpretable scale. | +| `call_budget` | `Optional['CallBudget']` | `None` | An adaptive per-call thinking budget (o1-style test-time scaling). Pass a `CallBudget` whose `base` matches the adapter's configured `max_tokens`, and the engine adjusts the per-call ceiling based on each parent's score and the remaining budget — promising parents get more thinking, dead-end ones get less, and the allocation tightens as the token budget is spent. The adapter must be wrapped with `budgeted_completion` for this to reach it. `None` (default) means the per-call ceiling is the adapter's own and never changes. | | `self_verify` | `bool` | `True` | Re-run the trajectory with the diff applied to record a local before/after delta. Doubles the rollouts spent per proposal; ports that score candidates only on held-out should pass `False`. | | `held_out_frac` | `float` | `0.4` | Fraction of `tasks` reserved for held-out scoring, in `(0, 1)`. | | `repo_path` | `Optional[str]` | `None` | Where the git-backed ledger lives. Omit for a throwaway repo that is removed when this call returns (not held until interpreter exit, so a sweep does not accumulate one git repo per run); **passing the same path again resumes** that ledger, and a caller-supplied path is never deleted. Git runs with an isolated config, so a personal `~/.gitconfig` (`commit.gpgsign`, `core.hooksPath`) cannot fail the ledger's own bookkeeping commits. | @@ -2681,6 +2691,23 @@ Archive( Keep the `k` best-scoring candidates and spread the workers over them. +### `CostEfficient(...)` + +FlatPuct that pays for exploration out of a budget. **Test-time scaling.** + +```python +CostEfficient( + c_puct: float = 1.0, + cost_exponent: float = 1.0, + prior_exponent: float = 0.0, + anneal: bool = True +) -> None +``` + +| method | what it does | +|---|---| +| `cost_factors(rows: Sequence[Candidate]) -> List[float]` | `cost**alpha` per row, normalised so the mean is 1. | + ### `MCTS(exploration: float = 1.4) -> None` UCT over the candidate tree: one evolve step is one rollout. @@ -2711,7 +2738,8 @@ SelectionContext( head: Candidate, candidates: Sequence[Candidate] = (), round: int = 0, - n_workers: int = 1 + n_workers: int = 1, + budget_remaining: float = 1.0 ) -> None ``` @@ -3075,6 +3103,10 @@ async_evolve( max_seconds: float = 20.0, max_iters: Optional[int] = None, max_calls: Optional[int] = None, + max_tokens: Optional[int] = None, + stop_on_diminishing_returns: bool = False, + efficiency_floor: float = 0.25, + call_budget: Optional[CallBudget] = None, target_reward: Optional[float] = None, patience: Optional[int] = None, max_worker_errors: int = 3, @@ -3124,6 +3156,10 @@ async_evolve( | `max_seconds` | `float` | `20.0` | Wall-clock budget for the **production phase only**. Two things still happen after it, so budget for them: a bounded shutdown (`shutdown_grace`, since an in-flight rollout cannot be cancelled) and **one held-out scoring pass** to compute `final_reward`. That pass is memoised per (artifact, task), so it is free when the final head was already scored by a sweep and costs a full held-out sweep of the backend when it was not -- which is exactly the case when the budget was too short for any sweep to finish. | | `max_iters` | `Optional[int]` | `None` | Stop after this many worker rollouts in total (a budget, not a barrier). | | `max_calls` | `Optional[int]` | `None` | Stop after this many actor invocations (`run` + `propose`) in total. The second half of an equal-budget comparison: two configurations matched on rollouts still differ in model spend whenever one of them asks for more proposals per rollout, and the cheaper unit is the one a reader assumes was held fixed. Both bounds are checked as each rollout lands, so a run overshoots only by what was already in flight. | +| `max_tokens` | `Optional[int]` | `None` | Stop after this many tokens in total (`prompt + completion`), as the meter measured them. The third budget unit, and the one cost is measured in: `max_calls` and `max_iters` count invocations, and a reasoning model can spend 40k tokens on hidden thinking in a single one, so neither bounds the bill. Checked as each rollout lands (tighter than the synchronous path's round barrier), and a `BudgetGovernor` degrades optional spend -- fusion tournaments at 75% of the budget, self-verify at 90% -- before the wall. `None` (default) means unbounded, and no governor is constructed. | +| `stop_on_diminishing_returns` | `bool` | `False` | Stop when the run's measured return per token falls off its own peak by more than `efficiency_floor`. Off by default; needs `max_tokens`. See `evolve`. | +| `efficiency_floor` | `float` | `0.25` | How far below the peak counts as diminishing, when the stop above is on. A self-calibrating ratio, not an absolute rate. See `evolve`. | +| `call_budget` | `Optional[CallBudget]` | `None` | An adaptive per-call thinking budget (o1-style test-time scaling). See `evolve`; the allocation happens before each worker's `propose` call, gated by the same `score` and `governor.remaining_fraction()` the synchronous path uses. | | `target_reward` | `Optional[float]` | `None` | Stop as soon as a sweep's held-out reward reaches this. Compared against the real reward, never against an acceptance probability. | | `patience` | `Optional[int]` | `None` | Stop after this many consecutive merge sweeps that fail to beat the best held-out reward seen so far. The async analogue of the synchronous knob: there are no round barriers here, so a *sweep* (one drain-and-merge by the merger) is the unit. `None` disables it. | | `max_worker_errors` | `int` | `3` | Consecutive failed rollouts before a worker that has *never* succeeded gives up. Workers that have succeeded at least once never retire; they back off and keep trying until the run's own budget ends it. | @@ -3411,7 +3447,7 @@ ArmResult( ) -> None ``` -### `Budget(rollouts: int, calls: Optional[int] = None) -> None` +### `Budget(rollouts: int, calls: Optional[int] = None, tokens: Optional[int] = None) -> None` What every arm is allowed to spend. @@ -3554,6 +3590,10 @@ Whether a candidate is committed. Accumulate a deduped list of rules/lessons (append-only, content-addressed). +### `BudgetGovernor` + +Tracks token spend against a cap and decides what to degrade. + ### `CALIBRATION_CEILING` Where a calibration share sits when the improvement pool is still finding new error modes, and where it goes once the pool has stopped. Both are policy choices, not measurements: what a label is worth in each pool depends on whether you are trying to *fix* the verifier or to *correct for* it. @@ -3566,6 +3606,10 @@ Where a calibration share sits when the improvement pool is still finding new er Somewhere to keep evaluations. In one process, across many, or on disk. +### `CallBudget` + +A mutable per-call max_tokens, set by the engine before each expansion. + ### `Completion` `Callable[[str], str]` — the one contract every model and agent satisfies. @@ -3786,6 +3830,10 @@ Agentic backends -- a base agent that *navigates documents with tools*, not just The control every efficiency number in this repository is missing. +### `budgeted_completion` + +Wrap a `max_tokens -> Completion` factory into an adaptive Completion. + ### `dataloader` Dependency-free dataset loading -- the *data layer* for examples/experiments. diff --git a/tests/test_backward_compat.py b/tests/test_backward_compat.py new file mode 100644 index 00000000..fef90bbb --- /dev/null +++ b/tests/test_backward_compat.py @@ -0,0 +1,139 @@ +"""Backward compatibility: the new features are extensions, not changes. + +The contract this file exists to prove: **a run that configures none of the new +options behaves exactly as it did before they existed.** Every mechanism added +on this branch -- `max_tokens`, `checkpointing`, `stop_on_diminishing_returns`, +`BudgetGovernor`, `CostEfficient` -- must be inert until asked for. + +These are behavioural assertions over a deterministic run, not a diff against a +stored golden: the run is seeded and its reward history is a function of the +seed alone, so "same history with and without the feature" is the strongest +statement a test can make here. +""" + +from __future__ import annotations + +import os +import warnings +from typing import List, Optional + +import pytest + +from agentdescent.agents import Usage +from agentdescent.evolution import evolve, Task, EvolutionResult + + +def _tasks(n: int = 8) -> List[Task]: + return [Task(id=f"t{i}", prompt=f"task {i}") for i in range(n)] + + +class _Agent: + """Deterministic agent: same proposals every run, no randomness.""" + + def __init__(self) -> None: + self.usage = Usage() + + def solve(self, rendered: str, task: Task) -> str: + self.usage.record(prompt_tokens=100, completion_tokens=200) + return task.id + + def propose(self, rendered: str, task: Task, output: str, + reward: float) -> Optional[str]: + self.usage.record(prompt_tokens=100, completion_tokens=200) + # A deterministic proposal per (task, reward) so the run is repeatable. + return f"rule-{task.id}-{round(reward, 3)}" if reward < 0.999 else None + + +def _run(**kwargs) -> EvolutionResult: + agent = _Agent() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return evolve( + _tasks(8), lambda t, o: 0.5, agent=agent, + rounds=4, n_workers=1, seed=7, usage=agent.usage, **kwargs) + + +def _history(r: EvolutionResult) -> list: + return [(h.round, round(h.held_out_reward, 9), h.committed, h.rejected) + for h in r.history] + + +# --- the default run is the pre-branch run --- + + +def test_no_new_options_is_the_old_path(): + """The default run must not acquire a budget, a checkpoint, or a governor.""" + r = _run() + assert r.budget is None, "a run with no max_tokens must not report a budget" + assert r.stop_reason in ("rounds", "patience", "target_reward") + # The new per-round token field exists but is only meaningful when a usage + # is shared; here it is shared, so it is populated -- the field being present + # is not a behaviour change. + assert hasattr(r.history[0], "tokens") + + +def test_disabled_flags_equal_no_flags(): + """Passing the new options at their defaults must equal passing none.""" + baseline = _run() + explicit = _run(max_tokens=None, checkpointing=False, + stop_on_diminishing_returns=False) + assert _history(baseline) == _history(explicit) + assert baseline.final_reward == explicit.final_reward + assert baseline.stop_reason == explicit.stop_reason + + +def test_a_never_reached_budget_does_not_change_the_search(): + """A token budget large enough never to bind must leave quality identical. + + This is the property that makes the budget *additive*: it can end a run + early, and it can degrade optional spend when it is close, but with room to + spare it touches nothing. + """ + baseline = _run() + budgeted = _run(max_tokens=10**12) + assert _history(baseline) == _history(budgeted), \ + "a budget that never binds must not change the run" + assert budgeted.budget is not None # it *reports* one ... + assert budgeted.budget["fusion_degraded"] is False + assert budgeted.budget["self_verify_degraded"] is False + + +def test_no_checkpoint_files_without_the_flag(tmp_path): + """No checkpoint is written unless `checkpointing=True`.""" + agent = _Agent() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + evolve(_tasks(8), lambda t, o: 0.5, agent=agent, + rounds=2, n_workers=1, repo_path=str(tmp_path / "repo"), + usage=agent.usage) + d = tmp_path / "repo" / "checkpoints" + assert not d.exists() or not any(d.iterdir()), \ + "a run without checkpointing=True must not write checkpoints" + + +def test_cost_efficient_is_opt_in(): + """The default selection is SingleHead; a run without a selection policy + must not construct or consult CostEfficient.""" + from agentdescent.policies import Policies + from agentdescent.selection import SingleHead, CostEfficient + + # Default policies use SingleHead; this is a structural check that the new + # policy is not the default anywhere. + pol = Policies() + assert pol.selection is None or isinstance(pol.selection, SingleHead) + assert not isinstance(pol.selection, CostEfficient) + + +def test_governor_inert_without_a_budget(): + """The governor the engine constructs with no max_tokens says nothing and + changes nothing.""" + from agentdescent.budget import BudgetGovernor + + g = BudgetGovernor(max_tokens=None, stop_on_diminishing=True) + assert not g.active + assert g.remaining_fraction() == 1.0 + assert not g.diminishing_returns() + assert g.affords_next_round() + assert g.allow_self_verify() + assert g.allow_fusion_tournament() + assert not g.over_budget() diff --git a/tests/test_baselines.py b/tests/test_baselines.py index 8c6c52b3..2ac53242 100644 --- a/tests/test_baselines.py +++ b/tests/test_baselines.py @@ -144,9 +144,12 @@ def test_fixing_calls_instead_moves_the_divergence_rather_than_removing_it(): assert any(unit == "rollouts" for _, unit, _, _ in comparison.confounded) -def test_the_fixed_unit_must_be_one_of_the_two(): +def test_the_fixed_unit_must_be_one_of_the_three(): + # rollouts, calls, tokens are the three units the engine can hold fixed. + for unit in ("rollouts", "calls", "tokens"): + compare([], fixed=unit) # does not raise with pytest.raises(ValueError): - compare([], fixed="tokens") + compare([], fixed="dollars") def test_an_arm_given_twice_the_budget_is_refused_not_averaged(): @@ -389,3 +392,81 @@ def test_fork_seeds_never_collide_with_the_other_arms(): assert not (fork_seeds & {0, 1, 2}), \ f"fork reused a plain arm seed: {sorted(fork_seeds & {0, 1, 2})}" assert len(fork_seeds) == 6, "and the two seeds' branches stay distinct" + + +# -- tokens: the third unit -------------------------------------------------- + + +def _token_arm(name, tokens, values): + """Arms that differ only in token spend, for the token-unit checks.""" + return [ArmResult(arm=name, seed=i, width=1, rollouts=10, calls=10, + prompt_tokens=t // 2, completion_tokens=t - t // 2, + wallclock=1.0, wallclock_parallel=1.0, + dev_reward=v, test_reward=v) + for i, (t, v) in enumerate(zip(tokens, values))] + + +def test_budget_split_splits_tokens_too(): + """A budget split across independent runs must divide every unit it holds.""" + b = Budget(rollouts=30, calls=60, tokens=9_000) + share = b.split(3) + assert share.rollouts == 10 + assert share.calls == 20 + assert share.tokens == 3_000 + + +def test_budget_split_leaves_none_tokens_none(): + """An unset token budget stays unset when split.""" + assert Budget(rollouts=30).split(2).tokens is None + + +def test_arm_result_tokens_is_the_sum(): + arm = _token_arm("a", [600], [0.5])[0] + assert arm.tokens == arm.prompt_tokens + arm.completion_tokens == 600 + + +def test_compare_can_fix_tokens(): + """`fixed="tokens"` is a valid unit, and an arm that spent far more is + reported as unequal rather than silently compared.""" + even = _token_arm("a", [1000, 1000, 1000], [0.6, 0.6, 0.6]) + even += _token_arm("b", [1000, 1000, 1000], [0.7, 0.7, 0.7]) + comparison = compare(even, fixed="tokens") + assert not comparison.unequal, "equal token spend must not be flagged" + + uneven = _token_arm("a", [1000, 1000, 1000], [0.6, 0.6, 0.6]) + uneven += _token_arm("b", [3000, 3000, 3000], [0.7, 0.7, 0.7]) + flagged = compare(uneven, fixed="tokens") + assert any(arm == "b" and unit == "tokens" + for arm, unit, _, _ in flagged.unequal), \ + "an arm that spent 3x the tokens while fixing tokens must be flagged" + + +def test_tokens_are_confounded_when_fixing_rollouts(): + """Fixing rollouts cannot fix tokens: an arm that thinks more per call + spends more, and the table has to say so.""" + even_rollouts = _token_arm("a", [1000, 1000, 1000], [0.6, 0.6, 0.6]) + even_rollouts += _token_arm("b", [5000, 5000, 5000], [0.7, 0.7, 0.7]) + comparison = compare(even_rollouts, fixed="rollouts") + assert any(arm == "b" and unit == "tokens" + for arm, unit, _, _ in comparison.confounded), \ + "a 5x token difference must appear as a confound" + + +def test_markdown_shows_tokens_or_a_dash(): + """The token column reads a number when reported and — when not, never 0.""" + reported = _token_arm("a", [600, 600, 600], [0.5, 0.5, 0.5]) + table = to_markdown(compare(reported)) + assert "tokens" in table + assert "600" in table + + unreported = _arm("a", [0.5, 0.5, 0.5]) # prompt_tokens=0 + table2 = to_markdown(compare(unreported)) + # The row's token cell is —, not 0. + assert "| — |" in table2 or "—" in table2 + + +def test_compare_empty_returns_empty_not_crash(): + """An empty sweep is a result to report, not an exception.""" + comparison = compare([]) + assert comparison.arms == {} + assert comparison.unequal == [] diff --git a/tests/test_budget_governor.py b/tests/test_budget_governor.py new file mode 100644 index 00000000..5451dd38 --- /dev/null +++ b/tests/test_budget_governor.py @@ -0,0 +1,219 @@ +"""Tests for the budget governor: graceful degradation before the token wall. + +Covers: +- Hard stop: the wall fires at 100% +- Fusion tournament degrades at the soft floor (75%) +- Self-verify degrades at the hard floor (90%) +- Projection: affords_next_round predicts correctly +- Inert without a budget +- Governor.summary() reports the right fields +""" + +from __future__ import annotations + +import pytest + +from agentdescent.budget import BudgetGovernor, SOFT_FLOOR, HARD_FLOOR + + +# --- basic thresholds --- + + +def test_governor_inert_without_budget(): + g = BudgetGovernor(max_tokens=None) + assert not g.active + assert not g.over_budget() + assert g.allow_fusion_tournament() + assert g.allow_self_verify() + assert g.affords_next_round() + s = g.summary() + assert s["max_tokens"] is None + assert s["remaining"] is None + assert s["fusion_degraded"] is False + assert s["self_verify_degraded"] is False + + +def test_hard_stop_at_budget(): + g = BudgetGovernor(max_tokens=1000) + g.spend(999) + assert not g.over_budget() + g.spend(1000) + assert g.over_budget() + g.spend(1200) + assert g.over_budget() + + +def test_fusion_degrades_at_soft_floor(): + g = BudgetGovernor(max_tokens=1000) + g.spend(740) + assert g.allow_fusion_tournament() + g.spend(750) + assert not g.allow_fusion_tournament() # at 75% + g.spend(800) + assert not g.allow_fusion_tournament() + + +def test_self_verify_degrades_at_hard_floor(): + g = BudgetGovernor(max_tokens=1000) + g.spend(890) + assert g.allow_self_verify() + g.spend(900) + assert not g.allow_self_verify() # at 90% + g.spend(950) + assert not g.allow_self_verify() + + +def test_fusion_degrades_before_self_verify(): + g = BudgetGovernor(max_tokens=1000) + g.spend(760) # past soft floor, before hard floor + assert not g.allow_fusion_tournament() + assert g.allow_self_verify() + + +# --- projection --- + + +def test_affords_next_round_with_history(): + g = BudgetGovernor(max_tokens=3000) + # Three rounds at 500 each = 1500 spent. Remaining 1500 / avg 500 = 3 rounds. + g.spend(500) + g.spend(1000) + g.spend(1500) + assert g.affords_next_round() + # Now simulate a huge round that leaves only 100 tokens. + g.spend(2900) + assert not g.affords_next_round() + + +def test_affords_next_round_no_history(): + g = BudgetGovernor(max_tokens=1000) + # No spend call yet — always afford (the run has not started). + assert g.affords_next_round() + + +def test_affords_next_round_zero_spend(): + g = BudgetGovernor(max_tokens=1000) + g.spend(0) # a run with no token reporting + assert g.affords_next_round() # can't project, so afford + + +def test_affords_next_round_uneven_spend(): + g = BudgetGovernor(max_tokens=2000) + g.spend(100) + g.spend(500) # delta 400 + g.spend(1500) # delta 1000 + # avg delta = (100+400+1000)/3 = 500 + # remaining = 2000 - 1500 = 500 + # 500 >= 500 → exactly affords + assert g.affords_next_round() + # one more round at avg 500 → 2000 spent, 0 left + g.spend(2000) + assert not g.affords_next_round() + + +# --- summary --- + + +def test_summary_fields(): + g = BudgetGovernor(max_tokens=1000, soft_floor=0.5, hard_floor=0.8) + g.spend(600) + s = g.summary() + assert s["max_tokens"] == 1000 + assert s["spent"] == 600 + assert s["remaining"] == 400 + assert s["fusion_degraded"] is True # 600 > 500 + assert s["self_verify_degraded"] is False # 600 < 800 + + +def test_summary_all_none_without_budget(): + g = BudgetGovernor() + s = g.summary() + assert s["max_tokens"] is None + assert s["spent"] == 0 + assert s["remaining"] is None + + +# --- custom floors --- + + +def test_custom_floors(): + g = BudgetGovernor(max_tokens=1000, soft_floor=0.5, hard_floor=0.8) + g.spend(400) + assert g.allow_fusion_tournament() # 40% < 50% + assert g.allow_self_verify() + g.spend(510) + assert not g.allow_fusion_tournament() # 51% > 50% + assert g.allow_self_verify() # 51% < 80% + g.spend(810) + assert not g.allow_self_verify() # 81% > 80% + + +# --- diminishing returns (the economic stop) --- + + +def test_diminishing_off_by_default(): + """No opt-in, no economic stop -- even with falling efficiency.""" + g = BudgetGovernor(max_tokens=1000) + for rew, tok in [(0.1, 100), (0.2, 200), (0.3, 300), (0.35, 400), (0.36, 500), (0.36, 600)]: + g.spend(tok) + g.observe(rew) + assert not g.diminishing_returns() + + +def test_diminishing_needs_a_budget(): + """Efficiency is return per *token*; without a token budget there is nothing + to divide by, so the rule cannot fire.""" + g = BudgetGovernor(max_tokens=None, stop_on_diminishing=True) + for rew, tok in [(0.1, 100), (0.2, 200), (0.3, 300), (0.35, 400), (0.36, 500), (0.36, 600)]: + g.spend(tok) + g.observe(rew) + assert not g.diminishing_returns() + + +def test_diminishing_needs_enough_samples(): + g = BudgetGovernor(max_tokens=1000, stop_on_diminishing=True, + min_efficiency_samples=4) + g.spend(100); g.observe(0.1) + g.spend(200); g.observe(0.2) + g.spend(300); g.observe(0.3) # only 2 samples yet + assert not g.diminishing_returns() + + +def test_diminishing_fires_when_return_per_token_collapses(): + g = BudgetGovernor(max_tokens=10_000, stop_on_diminishing=True, + min_efficiency_samples=4, efficiency_floor=0.25) + for rew, tok in [(0.10, 100), (0.20, 200), (0.30, 300), (0.35, 400), + (0.36, 500), (0.36, 600)]: + g.spend(tok) + g.observe(rew) + assert g.diminishing_returns() + + +def test_diminishing_does_not_fire_while_still_productive(): + """A run whose efficiency is steady must not be stopped.""" + g = BudgetGovernor(max_tokens=10_000, stop_on_diminishing=True, + min_efficiency_samples=4, efficiency_floor=0.25) + # Constant efficiency: every 100 tokens buys 0.1 reward. + for i in range(6): + g.spend(100 * (i + 1)) + g.observe(0.1 * (i + 1)) + assert not g.diminishing_returns() + + +def test_diminishing_never_fires_when_never_improved(): + """A run that never improved has no peak to decline from -- patience, not + this, is the rule for that.""" + g = BudgetGovernor(max_tokens=1000, stop_on_diminishing=True, + min_efficiency_samples=4) + for i in range(6): + g.spend(100 * (i + 1)) + g.observe(0.5) # flat + assert not g.diminishing_returns() + + +def test_observe_ignores_none_and_missing_cost_rounds(): + """A round with no reward, or one that spent no tokens, records nothing.""" + g = BudgetGovernor(max_tokens=1000, stop_on_diminishing=True) + g.spend(0); g.observe(None) + g.spend(0); g.observe(0.5) # no tokens -> no sample + assert g._efficiencies == [] diff --git a/tests/test_call_budget.py b/tests/test_call_budget.py new file mode 100644 index 00000000..8c6df87c --- /dev/null +++ b/tests/test_call_budget.py @@ -0,0 +1,205 @@ +"""Tests for adaptive per-call thinking budget (o1-style test-time scaling). + +Covers: +- CallBudget.allocate: score × budget_remaining → per-call max_tokens +- Bounds: never above base, never below floor +- budgeted_completion: caches per max_tokens, delegates correctly +- Inert without call_budget: the engine never touches it +- End-to-end: evolve with call_budget runs without crashing +""" + +from __future__ import annotations + +import warnings +from typing import Any, Optional + +import pytest + +from agentdescent.agents import Usage +from agentdescent.budget import CallBudget, budgeted_completion +from agentdescent.evolution import evolve, Task + + +# --- CallBudget.allocate --- + + +def test_allocate_promising_parent_full_budget(): + b = CallBudget(base=8192) + b.allocate(0.9, 1.0) + assert b.next is not None + assert b.next <= b.base + # 0.5 + 0.5*0.9 = 0.95; 0.5 + 0.5*1.0 = 1.0; 8192*0.95*1.0 = 7782 + assert b.next == 7782 + + +def test_allocate_dead_end_parent_low_budget(): + b = CallBudget(base=8192) + b.allocate(0.0, 0.0) + # 0.5 * 0.5 = 0.25; 8192*0.25 = 2048 = floor + assert b.next == b.floor + + +def test_allocate_never_above_base(): + b = CallBudget(base=4096) + for score in [0.0, 0.5, 1.0]: + for remaining in [0.0, 0.5, 1.0]: + b.allocate(score, remaining) + assert b.next <= b.base + + +def test_allocate_never_below_floor(): + b = CallBudget(base=4096) + for score in [0.0, -1.0]: + for remaining in [0.0, -1.0]: + b.allocate(score, remaining) + assert b.next >= b.floor + + +def test_allocate_clamps_score(): + b = CallBudget(base=4096) + b.allocate(2.0, 2.0) # above 1 + assert b.next <= b.base + b.allocate(-5.0, -5.0) + assert b.next >= b.floor + + +def test_floor_is_base_div_4(): + assert CallBudget(base=8192).floor == 2048 + assert CallBudget(base=1024).floor == 256 # minimum + + +def test_custom_floor(): + """A custom floor only kicks in when the calculated value falls below it.""" + b = CallBudget(base=8192, floor=1000) + # score=0, remaining=0 → 8192*0.25 = 2048 > floor → 2048, not 1000. + b.allocate(0.0, 0.0) + assert b.next == 2048 + # With a very small base and a high floor, the floor binds. + small = CallBudget(base=100, floor=50) + small.allocate(0.0, 0.0) + # 100*0.25 = 25 < 50 → floor binds + assert small.next == 50 + + +def test_next_none_by_default(): + b = CallBudget(base=4096) + assert b.next is None + + +# --- budgeted_completion --- + + +def test_budgeted_completion_uses_base_when_next_is_none(): + calls = [] + + def factory(mt): + calls.append(mt) + return lambda prompt: f"mt={mt}" + + b = CallBudget(base=4096) + bc = budgeted_completion(factory, b) + assert bc("hello") == "mt=4096" + assert calls == [4096] + + +def test_budgeted_completion_adapts_when_next_is_set(): + calls = [] + + def factory(mt): + calls.append(mt) + return lambda prompt: f"mt={mt}" + + b = CallBudget(base=4096) + bc = budgeted_completion(factory, b) + bc("first") # base + b.next = 2048 + bc("second") # adapted + assert 4096 in calls + assert 2048 in calls + + +def test_budgeted_completion_caches_per_max_tokens(): + calls = [] + + def factory(mt): + calls.append(mt) + return lambda prompt: f"mt={mt}" + + b = CallBudget(base=4096) + bc = budgeted_completion(factory, b) + bc("a") + bc("b") # same mt, cached + assert len(calls) == 1 + b.next = 2048 + bc("c") # new mt + assert len(calls) == 2 + bc("d") # cached + assert len(calls) == 2 + + +# --- inert without call_budget --- + + +def test_call_budget_unused_does_not_crash(): + """A run without call_budget must not touch it.""" + tasks = [Task(id=f"t{i}", prompt=f"task {i}") for i in range(8)] + + def run(rendered, task): + return task.id + + def propose(rendered, task, output, reward): + return f"rule-{task.id}" if reward < 0.999 else None + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve(tasks, lambda t, o: 0.5, run=run, propose=propose, + rounds=2, n_workers=1) + assert result.stop_reason == "rounds" + + +# --- end to end --- + + +def test_evolve_with_call_budget_runs(): + """evolve(call_budget=...) plumbs through and the budget is allocated.""" + tasks = [Task(id=f"t{i}", prompt=f"task {i}") for i in range(8)] + usage = Usage() + allocations = [] + + budget = CallBudget(base=4096) + + def factory(mt): + def complete(prompt): + usage.record(prompt_tokens=100, completion_tokens=mt // 10) + return f"rule: {prompt[:10]}" + + return complete + + from agentdescent.evolution import LLMAgent + + agent = LLMAgent(budgeted_completion(factory, budget)) + + # Track allocations + orig_allocate = budget.allocate + + def tracking_allocate(score, remaining, **kw): + orig_allocate(score, remaining, **kw) + allocations.append((round(score, 3), round(remaining, 3), budget.next)) + + budget.allocate = tracking_allocate + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve(tasks, lambda t, o: 0.5, agent=agent, + rounds=3, n_workers=1, max_tokens=10_000_000, + call_budget=budget, usage=usage) + + assert result.stop_reason in ("rounds", "patience", "max_tokens") + # The engine called allocate at least once per round that proposed. + if allocations: + # All allocations are within bounds. + assert all(a[2] <= budget.base for a in allocations) + assert all(a[2] >= budget.floor for a in allocations) + # Allocations vary (the score is constant 0.5, so they should be + # similar but the budget_remaining changes). + assert len(allocations) >= 1 diff --git a/tests/test_cost_aware_selection.py b/tests/test_cost_aware_selection.py new file mode 100644 index 00000000..1ee2fc5f --- /dev/null +++ b/tests/test_cost_aware_selection.py @@ -0,0 +1,318 @@ +"""Tests for cost-aware selection: the o1-style test-time-scaling decision. + +Covers: +- Cost discounts exploration, never exploitation +- Unknown cost is the pool mean (never "free") +- Annealing: exploration shrinks as budget_remaining falls +- cost_exponent=0 is FlatPuct exactly (the ablation and the safety property) +- Uniform costs are unchanged whatever the exponent +- The population layer fills Candidate.cost and passes budget_remaining +- End-to-end: a BudgetGovernor drives budget_remaining into the aggregator +""" + +from __future__ import annotations + +import pytest + +from agentdescent.selection import ( + Candidate, + FlatPuct, + CostEfficient, + SelectionContext, +) + + +def _cand(version, score, cost=None, selected=0, parent=None): + return Candidate(artifact_id="a", version=version, score=score, + cost=cost, selected=selected, parent=parent) + + +def _ctx(cands, *, budget_remaining=1.0, round=0): + return SelectionContext(head=cands[0], candidates=tuple(cands), + round=round, n_workers=1, + budget_remaining=budget_remaining) + + +# --- the cost divisor --- + + +def test_uniform_cost_leaves_selection_unchanged(): + """Every candidate the same cost must give the same order as FlatPuct.""" + cands = [_cand(0, 0.5, cost=100), _cand(1, 0.9, cost=100), _cand(2, 0.7, cost=100)] + cheap = CostEfficient(cost_exponent=1.0) + flat = FlatPuct() + assert ([c.version for c in cheap.select(_ctx(cands), 1)] + == [c.version for c in flat.select(_ctx(cands), 1)]) + + +def test_cost_exponent_zero_is_flatpuct_exactly(): + """alpha=0 is the ablation: it must be FlatPuct to the bit, so the cost term + cannot change a run that did not ask for it.""" + cands = [_cand(0, 0.5, cost=1), _cand(1, 0.9, cost=10_000), _cand(2, 0.7, cost=5)] + blind = CostEfficient(cost_exponent=0.0) + flat = FlatPuct() + assert ([c.version for c in blind.select(_ctx(cands), 3)] + == [c.version for c in flat.select(_ctx(cands), 3)]) + + +def test_unknown_cost_is_the_mean_not_free(): + """A candidate with no measured cost must not become the cheapest one.""" + factors = CostEfficient(cost_exponent=1.0).cost_factors( + [_cand(0, 0.5, cost=100), _cand(1, 0.5, cost=None), _cand(2, 0.5, cost=100)]) + # The unknown one gets the mean (100), so all three factors are 1.0. + assert factors == [1.0, 1.0, 1.0] + + +def test_all_costs_unknown_is_flatpuct(): + """No measured cost anywhere -> the divisor is 1 for all -> FlatPuct.""" + cands = [_cand(0, 0.5, cost=None), _cand(1, 0.9, cost=None), _cand(2, 0.7, cost=None)] + cheap = CostEfficient(cost_exponent=1.0) + flat = FlatPuct() + assert ([c.version for c in cheap.select(_ctx(cands), 1)] + == [c.version for c in flat.select(_ctx(cands), 1)]) + + +# --- cost discounts exploration, not exploitation --- + + +def test_cost_factors_discount_expensive_candidates(): + """The divisor is cost^alpha normalised to mean 1: cheap < 1 < expensive.""" + factors = CostEfficient(cost_exponent=1.0).cost_factors( + [_cand(0, 0.5, cost=50), _cand(1, 0.5, cost=100), _cand(2, 0.5, cost=200)]) + # mean = 116.67; cheap is below 1, expensive above. + assert factors[0] < 1.0 < factors[2] + assert factors[0] < factors[1] < factors[2] + + +def test_expensive_candidate_loses_the_exploration_bonus(): + """With exploration dominating (large c_puct, both visited so the sqrt term + is live), the cheap low-rank candidate beats the pricey high-rank one.""" + # Both visited, so `sqrt(total) > 0` and the exploration term is live. + cheap_low_rank = _cand(0, 0.4, cost=10, selected=1) + pricey_high_rank = _cand(1, 0.9, cost=1_000, selected=1) + policy = CostEfficient(c_puct=100.0, cost_exponent=1.0) + picked = policy.select(_ctx([cheap_low_rank, pricey_high_rank]), 1) + assert picked[0].version == cheap_low_rank.version, \ + "with exploration dominant, cost must decide and pick the cheap one" + + # The ablation: cost_exponent=0 (cost-blind FlatPuct) picks the high rank. + blind = CostEfficient(c_puct=100.0, cost_exponent=0.0) + picked_blind = blind.select(_ctx([cheap_low_rank, pricey_high_rank]), 1) + assert picked_blind[0].version == pricey_high_rank.version, \ + "without cost, the rank must decide" + + +def test_exploitation_is_not_discounted_by_cost(): + """A clearly better candidate is still picked even if it is the pricier one + -- cost must never disqualify the grounded pick.""" + better_pricey = _cand(0, 0.95, cost=10_000, selected=0) + worse_cheap = _cand(1, 0.10, cost=1, selected=0) + policy = CostEfficient(c_puct=1.0, cost_exponent=1.0) + picked = policy.select(_ctx([better_pricey, worse_cheap]), 1) + assert picked[0].version == better_pricey.version + + +# --- annealing --- + + +def test_exploration_anneals_with_remaining_budget(): + """With budget left, a fresh (unvisited) candidate gets explored; with none + left, the policy exploits the ranked best instead.""" + best = _cand(0, 0.9, cost=100, selected=3) # explored, high score + fresh = _cand(1, 0.5, cost=100, selected=0) # never tried, lower score + + full_budget = CostEfficient(c_puct=2.0, cost_exponent=0.0, anneal=True) + # At the start the exploration term is large enough to try the fresh one. + first = full_budget.select( + _ctx([best, fresh], budget_remaining=1.0), 1)[0] + assert first.version == fresh.version, \ + "with a full budget the unexplored candidate must get a look" + + no_budget = CostEfficient(c_puct=2.0, cost_exponent=0.0, anneal=True) + last = no_budget.select( + _ctx([best, fresh], budget_remaining=0.0), 1)[0] + assert last.version == best.version, \ + "with no budget left the policy must exploit, not explore" + + +def test_anneal_false_pins_exploration(): + """anneal=False is the ablation: same pick at any budget_remaining.""" + best = _cand(0, 0.9, cost=100, selected=3) + fresh = _cand(1, 0.5, cost=100, selected=0) + policy = CostEfficient(c_puct=2.0, cost_exponent=0.0, anneal=False) + a = policy.select(_ctx([best, fresh], budget_remaining=1.0), 1)[0] + b = policy.select(_ctx([best, fresh], budget_remaining=0.0), 1)[0] + assert a.version == b.version + + +def test_negative_budget_remaining_does_not_invert(): + """A caller bug (negative fraction) must not flip the exploration sign.""" + best = _cand(0, 0.9, cost=100, selected=3) + fresh = _cand(1, 0.5, cost=100, selected=0) + policy = CostEfficient(c_puct=2.0, cost_exponent=0.0, anneal=True) + # Clamped to 0 -> behaves like no budget (exploit), not like -1. + picked = policy.select(_ctx([best, fresh], budget_remaining=-5.0), 1)[0] + assert picked.version == best.version + + +# --- selection shape --- + + +def test_returns_n_picks_and_reserves_visits(): + cands = [_cand(0, 0.5, cost=10), _cand(1, 0.6, cost=10), _cand(2, 0.7, cost=10)] + policy = CostEfficient() + picks = policy.select(_ctx(cands), 3) + assert len(picks) == 3 + # Each pick is a real candidate. + assert all(p in cands for p in picks) + + +def test_single_candidate_returns_head(): + one = _cand(0, 0.5, cost=10) + policy = CostEfficient() + assert policy.select(_ctx([one]), 2) == [one, one] + + +# --- population wiring --- + + +def test_population_aggregator_fills_candidate_cost(): + """The archive's candidates carry a content-size cost proxy.""" + import threading + from agentdescent.population import PopulationAggregator + from agentdescent.selection import Archive + + agg = PopulationAggregator.__new__(PopulationAggregator) + agg.selection = Archive(sampling="novelty") + agg.population_artifact = "a" + agg._archive = [ + {"state": {"rule": "x" * 100}, "score": 0.5, "version": 1, "selected": 0}, + {"state": {"rule": "y" * 1000}, "score": 0.6, "version": 2, "selected": 0}, + ] + agg._keys = set() + agg._archive_lock = threading.Lock() + agg._selections = 0 + agg.budget_remaining = 1.0 + + cands = agg._candidates() + assert all(c.cost is not None for c in cands), "cost proxy not filled" + # The bigger artifact costs more. + assert cands[1].cost > cands[0].cost + + +def test_population_passes_budget_remaining_to_context(): + """PopulationAggregator.step() hands its budget_remaining to the selection + context, so a budget-aware policy sees it.""" + import threading + from agentdescent.population import PopulationAggregator + + captured = {} + + class Capturing: + def select(self, ctx, n): + captured["budget_remaining"] = ctx.budget_remaining + return [ctx.head] * n + + agg = PopulationAggregator.__new__(PopulationAggregator) + agg.selection = Capturing() + agg.population_artifact = "a" + agg._archive = [ + {"state": {"r": "1"}, "score": 0.5, "version": 1, "selected": 0}, + {"state": {"r": "2"}, "score": 0.6, "version": 2, "selected": 0}, + ] + agg._keys = set() + agg._archive_lock = threading.Lock() + agg._selections = 0 + agg.budget_remaining = 0.42 + # `step()` calls a lot of super() machinery; call the selection part the + # way step() does, so the assertion is about the context construction. + cands = agg._candidates() + from agentdescent.selection import SelectionContext + ctx = SelectionContext(head=cands[0], candidates=tuple(cands), + round=agg._selections, n_workers=1, + budget_remaining=agg.budget_remaining) + agg.selection.select(ctx, 1) + assert captured["budget_remaining"] == 0.42 + + +# --- governor -> budget_remaining --- + + +def test_governor_remaining_fraction(): + from agentdescent.budget import BudgetGovernor + g = BudgetGovernor(max_tokens=1000) + assert g.remaining_fraction() == 1.0 + g.spend(250) + assert g.remaining_fraction() == pytest.approx(0.75) + g.spend(1000) + assert g.remaining_fraction() == 0.0 + g.spend(1500) # overshoot (barrier lag) + assert g.remaining_fraction() == 0.0 # clamped, not negative + + +def test_governor_remaining_fraction_no_budget(): + from agentdescent.budget import BudgetGovernor + assert BudgetGovernor(max_tokens=None).remaining_fraction() == 1.0 + + +def test_set_budget_remaining_only_touches_declaring_aggregators(): + from agentdescent.evolution import _set_budget_remaining + from agentdescent.budget import BudgetGovernor + + g = BudgetGovernor(max_tokens=1000) + g.spend(500) + + class Declaring: + budget_remaining = 1.0 + + class Plain: + pass + + d, p = Declaring(), Plain() + _set_budget_remaining(d, g) + _set_budget_remaining(p, g) + assert d.budget_remaining == pytest.approx(0.5) + assert not hasattr(p, "budget_remaining"), \ + "a plain aggregator must not gain an attribute nothing reads" + + +# --- end to end --- + + +def test_evolve_with_cost_efficient_and_token_budget(): + """The whole chain: evolve(selection=CostEfficient(), max_tokens=N) runs, + the governor drives budget_remaining, and the policy anneals on it.""" + import warnings + from agentdescent.agents import Usage + from agentdescent.evolution import evolve, Task + from agentdescent.policies import Policies + from agentdescent.selection import CostEfficient + + class Agent: + def __init__(self): + self.usage = Usage() + + def solve(self, rendered, task): + self.usage.record(prompt_tokens=100, completion_tokens=200) + return task.id + + def propose(self, rendered, task, output, reward): + self.usage.record(prompt_tokens=100, completion_tokens=200) + return f"rule-{task.id}" if reward < 0.999 else None + + tasks = [Task(id=f"t{i}", prompt=f"task {i}") for i in range(8)] + agent = Agent() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve( + tasks, lambda t, o: 0.5, agent=agent, + rounds=6, n_workers=1, + max_tokens=6000, + policies=Policies(selection=CostEfficient(c_puct=1.0)), + usage=agent.usage, + ) + # The run completed (or stopped on the budget) without crashing, and the + # budget was tracked. + assert result.stop_reason in ("rounds", "max_tokens", "patience") + assert result.budget is not None + assert result.budget["spent"] > 0 diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py new file mode 100644 index 00000000..40511a57 --- /dev/null +++ b/tests/test_token_budget.py @@ -0,0 +1,363 @@ +"""Tests for ``max_tokens`` budget enforcement. + +Covers: +- ``stop_reason == "max_tokens"`` when the cap is reached (sync path) +- ``stop_reason == "max_tokens"`` on the async path +- Token consumption reported in ``RoundInfo.tokens`` +- ``result.usage.total_tokens`` reflects actual spend +- A generous cap does not interfere with normal termination +- Interaction with ``max_calls`` and ``max_seconds`` +- Zero-token runs (no model) do not trigger the cap +""" + +from __future__ import annotations + +import warnings +from typing import Any, Dict, List, Optional + +import pytest + +from agentdescent.agents import Usage +from agentdescent.evolution import evolve, Task, RoundInfo, EvolutionResult + + +# --- helpers --------------------------------------------------------------- + + +def _tasks(n: int = 8) -> List[Task]: + return [Task(id=f"t{i}", prompt=f"task {i}") for i in range(n)] + + +class _CountingAgent: + """An agent that reports a fixed token cost per call. + + ``prompt_tokens`` and ``completion_tokens`` are set per call, so the total + is deterministic and controllable — essential for testing a token budget. + """ + + def __init__(self, prompt: int = 100, completion: int = 200) -> None: + self.prompt = prompt + self.completion = completion + self.calls = 0 + self._usage = Usage() + + def solve(self, rendered: str, task: Task) -> str: + self.calls += 1 + self._usage.record( + prompt_tokens=self.prompt, completion_tokens=self.completion) + return task.id + + def propose(self, rendered: str, task: Task, output: str, + reward: float) -> Optional[str]: + self.calls += 1 + self._usage.record( + prompt_tokens=self.prompt, completion_tokens=self.completion) + return f"rule-{task.id}" if reward < 0.999 else None + + @property + def usage(self) -> Usage: + return self._usage + + +def _run_evolve(*, max_tokens=None, rounds=3, asynchronous=False, **kw): + """Run a minimal evolve with a counting agent and token tracking.""" + agent = _CountingAgent(prompt=100, completion=200) + # 300 tokens per call (100 prompt + 200 completion). + defaults = dict( + rounds=rounds, + n_workers=1, + max_tokens=max_tokens, + asynchronous=asynchronous, + usage=agent.usage, # share so meter sees the tokens + ) + if asynchronous: + defaults["max_seconds"] = 30 + defaults.update(kw) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve(_tasks(8), lambda t, o: 0.5, agent=agent, **defaults) + return result, agent + + +# --- sync path ------------------------------------------------------------- + + +def test_sync_max_tokens_stops_run(): + """A token budget that is reached stops the run with stop_reason="max_tokens".""" + # 300 tokens/call, 2 calls/round (solve + propose) = 600 tokens/round. + # 3 rounds = 1800 tokens. A cap of 1200 should stop after round 2. + result, agent = _run_evolve(max_tokens=1200, rounds=5) + assert result.stop_reason == "max_tokens" + assert result.usage.total_tokens >= 1200 + # The run did not complete all 5 rounds. + assert len(result.history) < 5 + + +def test_sync_max_tokens_generous_cap_does_not_interfere(): + """A generous cap lets the run complete normally.""" + result, agent = _run_evolve(max_tokens=10_000_000, rounds=2) + assert result.stop_reason == "rounds" + assert len(result.history) == 2 + + +def test_sync_max_tokens_zero_when_no_usage(): + """A run with no token reporting (no Usage shared) does not trigger the cap.""" + def run(rendered, task): + return task.id + + def propose(rendered, task, output, reward): + return f"rule-{task.id}" if reward < 0.999 else None + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve( + _tasks(8), lambda t, o: 0.5, + run=run, propose=propose, + rounds=2, n_workers=1, + max_tokens=1, # would trigger if tokens were counted + ) + # No Usage object shared → meter reports 0 tokens → cap never fires. + assert result.stop_reason == "rounds" + + +def test_round_info_has_tokens_field(): + """RoundInfo.tokens is populated and cumulative.""" + result, agent = _run_evolve(max_tokens=10_000_000, rounds=3) + assert len(result.history) == 3 + tokens = [r.tokens for r in result.history] + # Cumulative: should be non-decreasing. + assert all(tokens[i] <= tokens[i + 1] for i in range(len(tokens) - 1)) + # The last round should reflect the total spend. + assert tokens[-1] == result.usage.total_tokens + + +def test_sync_max_tokens_interaction_with_max_calls(): + """When both max_tokens and max_calls are set, whichever fires first wins.""" + # 300 tokens/call, max_calls=3, max_tokens=10000 → max_calls fires first. + result, agent = _run_evolve(max_tokens=10_000, rounds=5) + # Now with a tight token cap. + result2, _ = _run_evolve(max_tokens=600, rounds=5) + # 600 tokens = 2 calls = 1 round. max_tokens should fire. + assert result2.stop_reason == "max_tokens" + + +def test_result_usage_exposes_total_tokens(): + """EvolutionResult.usage.total_tokens is the sum of prompt + completion.""" + result, agent = _run_evolve(max_tokens=10_000_000, rounds=2) + assert result.usage.total_tokens == result.usage.prompt_tokens + result.usage.completion_tokens + assert result.usage.total_tokens > 0 + + +# --- async path ------------------------------------------------------------ + + +def test_async_max_tokens_stops_run(): + """The async path checks per-rollout and stops with stop_reason="max_tokens".""" + result, agent = _run_evolve( + max_tokens=900, rounds=10, asynchronous=True, + max_seconds=30) + # 300 tokens/call. 900 tokens = 3 calls. Should stop early. + assert result.stop_reason == "max_tokens" + assert result.usage.total_tokens >= 900 + + +def test_async_max_tokens_generous_cap_completes(): + """A generous cap on the async path does not interfere.""" + result, agent = _run_evolve( + max_tokens=10_000_000, rounds=2, asynchronous=True, + max_seconds=30) + assert result.stop_reason in ("rounds", "max_iters", "max_seconds") + + +# --- serialization --------------------------------------------------------- + + +def test_round_info_tokens_serializes(): + """RoundInfo preserves the tokens field through EvolutionResult.to_dict.""" + info = RoundInfo( + round=3, held_out_reward=0.75, n_items=5, + committed=2, rejected=1, reasons={}, + elapsed_s=10.0, rollouts=6, calls=12, + tokens=3600) + + # Verify the field survives __dict__ access (what to_dict uses). + assert info.__dict__["tokens"] == 3600 + + # Verify a real run populates it. + result, agent = _run_evolve(max_tokens=10_000_000, rounds=1) + assert result.history[0].tokens > 0 + assert result.history[0].tokens == result.usage.total_tokens + + +def test_agent_usage_adopted_automatically(): + """When the agent carries a Usage and the caller did not pass usage=, + the meter adopts the agent's — token budgets and the tokens= column work + without the caller threading the same object through two parameters.""" + from agentdescent.agents import Usage + + agent = _CountingAgent(prompt=100, completion=200) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve( + _tasks(8), lambda t, o: 0.5, + agent=agent, + rounds=2, n_workers=1, + max_tokens=10_000_000, + # NOTE: no usage= here — the agent's must be adopted. + ) + assert result.usage is agent.usage + assert result.usage.total_tokens > 0 + assert result.history[0].tokens > 0 + + +def test_agent_usage_adopted_budget_fires(): + """A max_tokens cap fires even without evolve(usage=...) when the agent + reports tokens through its own Usage.""" + agent = _CountingAgent(prompt=100, completion=200) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve( + _tasks(8), lambda t, o: 0.5, + agent=agent, + rounds=5, n_workers=1, + max_tokens=900, # 300/call → stops early + # no usage= — adoption must make this fire + ) + assert result.stop_reason == "max_tokens" + assert result.usage is agent.usage + + +def test_explicit_usage_wins_over_agent_usage(): + """An explicit usage= takes precedence over the agent's — the caller who + passes both has said which counter is the truth.""" + from agentdescent.agents import Usage + + agent = _CountingAgent(prompt=100, completion=200) + mine = Usage() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve( + _tasks(8), lambda t, o: 0.5, + agent=agent, + rounds=1, n_workers=1, + usage=mine, + ) + assert result.usage is mine + assert result.usage is not agent.usage + + +def test_no_usage_anywhere_reports_zero(): + """A bare run/propose pair with no Usage anywhere reports tokens=0 and a + token budget that cannot fire — the honest reading, not an error.""" + def run(rendered, task): + return task.id + + def propose(rendered, task, output, reward): + return f"rule-{task.id}" if reward < 0.999 else None + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve( + _tasks(8), lambda t, o: 0.5, + run=run, propose=propose, + rounds=1, n_workers=1, + max_tokens=1, + ) + assert result.usage.total_tokens == 0 + assert result.stop_reason == "rounds" + assert result.history[0].tokens == 0 + + +def test_async_governor_degrades_self_verify(): + """The async path's governor turns off self-verify at the hard floor. + + We verify this indirectly: a run with a tight token budget should produce + evidence cards with delta=0 (self-verify skipped) once the budget crosses + the hard floor, whereas a run without a budget produces non-zero deltas. + + Since we cannot easily inspect individual cards, we verify the observable + consequence: the run stops with stop_reason='max_tokens' and the governor + was consulted (the verbose output would mention it, but we check the + budget fired at all — the governor's check is the only path that produces + stop_reason='max_tokens' on the async path).""" + agent = _CountingAgent(prompt=100, completion=200) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + # 300 tokens/call. With 2 workers and max_tokens=600, + # the first 2 calls (600 tokens) hit the budget. + result = evolve( + _tasks(8), lambda t, o: 0.5, + agent=agent, + rounds=10, n_workers=1, + max_tokens=600, + asynchronous=True, + max_seconds=30, + usage=agent.usage, + ) + assert result.stop_reason == "max_tokens" + + +def test_async_governor_fusion_degrades(): + """The async path's governor degrades the fusion tournament at the soft floor. + + With fusion_tournament=True and a tight token budget, the governor should + turn off the tournament before the budget is exhausted. We verify by + checking that the run still completes with the right stop reason and + the tournament was available at least initially (fusion_trials is populated + by DefaultFusion and accessible on the result).""" + from agentdescent.agents import Usage + agent = _CountingAgent(prompt=100, completion=200) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = evolve( + _tasks(8), lambda t, o: 0.5, + agent=agent, + rounds=10, n_workers=1, + max_tokens=2000, + fusion_tournament=True, + asynchronous=True, + max_seconds=30, + usage=agent.usage, + ) + # The run should have stopped due to the token budget. + assert result.stop_reason == "max_tokens" + # fusion_trials may or may not be populated (depends on whether any merge + # had survivors before degradation), but the run must not crash. + assert hasattr(result, "fusion_trials") + + +def test_result_budget_field_populated(): + """EvolutionResult.budget carries the governor's summary when max_tokens is set.""" + result, _ = _run_evolve(max_tokens=10_000_000, rounds=2) + assert result.budget is not None + assert result.budget["max_tokens"] == 10_000_000 + assert result.budget["spent"] > 0 + assert result.budget["remaining"] > 0 + assert result.budget["fusion_degraded"] is False # generous cap, no degradation + assert result.budget["self_verify_degraded"] is False + + +def test_result_budget_none_without_max_tokens(): + """EvolutionResult.budget is None when no max_tokens was given.""" + result, _ = _run_evolve(rounds=1) + assert result.budget is None + + +def test_result_budget_shows_degradation(): + """A tight budget produces fusion_degraded=True in the result.""" + # 300 tokens/call, 2 calls/round = 600 tokens/round. + # With max_tokens=700, soft_floor=75% → 525 tokens → fusion degrades round 1. + result, _ = _run_evolve(max_tokens=700, rounds=5) + assert result.budget is not None + assert result.budget["fusion_degraded"] is True + + +def test_diminishing_returns_flag_plumbs_through(): + """`stop_on_diminishing_returns=True` reaches the governor; a run with a + never-binding budget and constant reward does not fire it (no peak to + decline from), so the run ends the ordinary way.""" + result, _ = _run_evolve(max_tokens=10**12, rounds=4, + stop_on_diminishing_returns=True) + assert result.stop_reason in ("rounds", "max_tokens", + "diminishing_returns", "patience") + assert result.budget is not None