feat(budgetallocator): cost-aware search - spend a token budget by value, not by exhaustion

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

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",
 ]

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()