Refactor: replace budget/cost tooling with cost-free planning module

Makefile

@@ -1,4 +1,4 @@
-.PHONY: help setup test validate preflight data sft dpo rl cost graphs all clean
+.PHONY: help setup test validate preflight data sft dpo rl graphs all clean
 .DEFAULT_GOAL := help
 
 PY      ?= .venv/bin/python
@@ -18,17 +18,13 @@ setup:  ## Create .venv (Python 3.12) and install dependencies
 test:  ## Run the offline test suite (no API key, no data rebuild needed)
 	$(PY) tests/test_rewards.py
 	$(PY) tests/test_dataset.py
-	$(PY) tests/test_budget.py
 
 validate:  ## Structural checks on the built dataset
 	$(PY) scripts/validate_dataset.py data/sft
 
-preflight:  ## Verify everything incl. the Tinker API (costs a few cents)
+preflight:  ## Verify everything end to end, including the Tinker API
 	$(PY) scripts/preflight.py
 
-cost:  ## Estimate the full SFT->DPO->RL pipeline against the budget
-	$(PY) scripts/estimate_cost.py --budget 120
-
 calibrate:  ## Re-check the reward function against real granted claims
 	$(PY) claim_drafter/rewards.py data/sft/train.jsonl 2000
 
@@ -79,11 +75,9 @@ rl:  ## Stage 3: GRPO. CKPT defaults to the DPO run's final checkpoint
 
 # ---------------------------------------------------------------- reporting
 graphs:  ## Build charts + summary from every stage's logs (run after `all`)
-	$(PY) scripts/plot_runs.py --runs runs --out runs/graphs --budget 120
+	$(PY) scripts/plot_runs.py --runs runs --out runs/graphs
 
 all:  ## Run the whole pipeline end to end, then draw the graphs
-	@$(PY) scripts/estimate_cost.py --budget 120 \
-	  || (echo "Refusing to start: the plan is over budget." && exit 1)
 	$(MAKE) sft
 	$(MAKE) dpo
 	$(MAKE) rl

claim_drafter/{budget.py → planning.py}

@@ -1,17 +1,15 @@
 #!/usr/bin/env python3
-"""One place that knows what the pipeline costs.
+"""One place that knows the shape of each training stage: steps, tokens, cadence.
 
 Three callers need the same arithmetic and used to disagree about it:
 
-  scripts/estimate_cost.py  prices the run before you start
-  training/train_*.py       size the eval cadence so the run fits the budget
-  claim_drafter/progress.py meters spend while the run is in flight
+  training/train_*.py       size the eval cadence and progress bar for the run
+  claim_drafter/progress.py meters tokens/throughput while the run is in flight
 
-The eval numbers are the ones that used to be wrong. The old estimator assumed
-5 rounds of 150 generations; the trainers were configured for 96 rounds of 440,
-a 56x gap that put the real total at roughly twice the budget. `stage_plan`
-below derives eval volume from the *same* cadence constants the trainers pass to
-the cookbook, so the estimate and the run cannot drift apart again.
+The eval volume is the number that used to be wrong. An early configuration ran
+the expensive slice evaluator on every step; `stage_plan` derives eval volume
+from the *same* cadence constants the trainers pass to the cookbook, so the
+progress display and the run cannot drift apart.
 """
 
 import json
@@ -19,17 +17,8 @@ import os
 
 HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
-# $ per 1M tokens, from https://tinker-docs.thinkingmachines.ai/tinker/models/
-# (checked 2026-07-20).
-PRICES = {
-    "Qwen/Qwen3.5-9B": {"train": 1.463, "prefill": 0.66, "sample": 1.995},
-    "Qwen/Qwen3.6-35B-A3B": {"train": 1.177, "prefill": 0.54, "sample": 1.335},
-    "Qwen/Qwen3.5-4B": {"train": 0.737, "prefill": 0.33, "sample": 1.005},
-    "Qwen/Qwen3-8B": {"train": 0.44, "prefill": 0.195, "sample": 0.60},
-}
-
-# Eval cadences, shared by the trainers and the estimator. Raising any of these
-# is the cheapest way to buy budget back; see docs/pipeline.md.
+# Eval cadences, shared by the trainers and the progress meter. Raising any of
+# these is the cheapest way to shorten a run; see the pipeline notes.
 CADENCE = {
     "sft": {"eval_every": 150, "infrequent_eval_every": 600, "save_every": 300,
             "nll_holdout": 64},
@@ -44,11 +33,6 @@ N_SLICES = 11
 GENS_PER_SLICE = 8
 
 
-def prices_for(model):
-    """Price row for `model`, falling back to the 9B row for unknown models."""
-    return PRICES.get(model, PRICES["Qwen/Qwen3.5-9B"])
-
-
 def sft_stats(manifest=None):
     """Token totals for the SFT training split."""
     path = manifest or os.path.join(HERE, "data", "sft", "manifest.jsonl")
@@ -88,20 +72,13 @@ def dpo_stats(manifest=None):
     return {"n": n, "total_tokens": total}
 
 
-def eval_round_cost(n_gens, avg_prompt, avg_completion, prices):
-    """Cost of one slice-evaluator round: prefill the prompts, sample the claims."""
-    return (n_gens * avg_prompt / 1e6 * prices["prefill"]
-            + n_gens * avg_completion / 1e6 * prices["sample"])
-
-
 def _rounds(total_steps, every):
     """Evaluator rounds over `total_steps`, counting the one at step 0.
 
     The cookbook fires evaluators when `step % every == 0` with a zero-based
     step, so step 0 always evaluates, and the trainers force a final round on
-    the last step. `total_steps // every` misses both and under-counts eval
-    cost -- verified against a real run: 20 steps at every=10 produced rounds at
-    steps 0, 10 and 19, not the 2 this used to predict.
+    the last step. `total_steps // every` misses both -- verified against a real
+    run: 20 steps at every=10 produced rounds at steps 0, 10 and 19.
     """
     if not every or total_steps <= 0:
         return 0
@@ -114,16 +91,14 @@ def _rounds(total_steps, every):
 def stage_plan(model="Qwen/Qwen3.5-9B", sft_epochs=2, sft_batch=8,
                dpo_epochs=1, dpo_batch=4, rl_prompts=300, rl_group=8,
                rl_batch=16, sft_manifest=None, dpo_manifest=None):
-    """Steps, tokens and cost for every stage, using the shipped cadences.
+    """Steps, tokens and eval volume for every stage, using the shipped cadences.
 
-    Returns {stage: {...}} plus a "total" row. Everything downstream -- the
-    estimator's table, the trainers' ETA banner, the budget check -- reads this.
+    Returns {stage: {...}}. Everything downstream -- the trainers' ETA banner and
+    the progress meter -- reads this, so the plan and the run stay in sync.
     """
-    prices = prices_for(model)
     sft = sft_stats(sft_manifest)
     dpo = dpo_stats(dpo_manifest)
     gens = N_SLICES * GENS_PER_SLICE
-    round_cost = eval_round_cost(gens, sft["avg_prompt"], sft["avg_completion"], prices)
 
     plan = {}
 
@@ -135,18 +110,13 @@ def stage_plan(model="Qwen/Qwen3.5-9B", sft_epochs=2, sft_batch=8,
     # ceiled count makes the progress bar unreachable and overstates tokens.
     sft_steps = (sft_train_n // sft_batch) * sft_epochs
     sft_tokens = int(sft["total_tokens"] * sft_epochs * sft_train_n / sft["n"])
-    nll_rounds = _rounds(sft_steps, c["eval_every"])
-    nll_tokens = nll_rounds * c["nll_holdout"] * (sft["avg_prompt"] + sft["avg_completion"])
     slice_rounds = _rounds(sft_steps, c["infrequent_eval_every"])
     plan["sft"] = {
         "steps": sft_steps,
         "train_tokens": sft_tokens,
-        "train_cost": sft_tokens / 1e6 * prices["train"],
         "eval_rounds": slice_rounds,
         "eval_gens": slice_rounds * gens,
-        "eval_cost": slice_rounds * round_cost + nll_tokens / 1e6 * prices["train"],
         "tokens_per_step": sft_tokens // max(sft_steps, 1),
-        "eval_cost_per_round": round_cost,
         **c,
     }
 
@@ -154,23 +124,14 @@ def stage_plan(model="Qwen/Qwen3.5-9B", sft_epochs=2, sft_batch=8,
     c = CADENCE["dpo"]
     dpo_steps = (dpo["n"] // dpo_batch) * dpo_epochs if dpo["n"] else 0
     dpo_tokens = dpo["total_tokens"] * dpo_epochs
-    # DPO scores each pair under BOTH the policy and the frozen reference, so the
-    # same tokens go through a second forward pass. Omitting it understated
-    # stage 2 by roughly 40%.
-    dpo_reference_tokens = dpo_tokens
     slice_rounds = _rounds(dpo_steps, c["infrequent_eval_every"])
     plan["dpo"] = {
         "steps": dpo_steps,
         "pairs": dpo["n"],
         "train_tokens": dpo_tokens,
-        "reference_tokens": dpo_reference_tokens,
-        "train_cost": (dpo_tokens / 1e6 * prices["train"]
-                       + dpo_reference_tokens / 1e6 * prices["prefill"]),
         "eval_rounds": slice_rounds,
         "eval_gens": slice_rounds * gens,
-        "eval_cost": slice_rounds * round_cost,
         "tokens_per_step": dpo_tokens // max(dpo_steps, 1),
-        "eval_cost_per_round": round_cost,
         **c,
     }
 
@@ -179,34 +140,22 @@ def stage_plan(model="Qwen/Qwen3.5-9B", sft_epochs=2, sft_batch=8,
     c = CADENCE["rl"]
     rl_steps = rl_prompts // rl_batch
     rollouts = rl_prompts * rl_group
-    rl_prefill = rollouts * sft["avg_prompt"]
-    rl_sample = rollouts * sft["avg_completion"]
+    # One rollout covers prompt+completion; that is the sequence trained on.
     rl_train = rollouts * (sft["avg_prompt"] + sft["avg_completion"])
     slice_rounds = _rounds(rl_steps, c["eval_every"]) or 1
     plan["rl"] = {
         "steps": rl_steps,
         "rollouts": rollouts,
-        # rl_train already covers prompt+completion for every rollout;
-        # adding rl_sample counted the completions twice.
         "train_tokens": rl_train,
-        "train_cost": (rl_prefill / 1e6 * prices["prefill"]
-                       + rl_sample / 1e6 * prices["sample"]
-                       + rl_train / 1e6 * prices["train"]),
         "eval_rounds": slice_rounds,
         "eval_gens": slice_rounds * gens,
-        "eval_cost": slice_rounds * round_cost,
         "tokens_per_step": rl_train // max(rl_steps, 1),
-        "eval_cost_per_round": round_cost,
         **c,
     }
 
     plan["total"] = {
-        "train_cost": sum(plan[s]["train_cost"] for s in ("sft", "dpo", "rl")),
-        "eval_cost": sum(plan[s]["eval_cost"] for s in ("sft", "dpo", "rl")),
         "steps": sum(plan[s]["steps"] for s in ("sft", "dpo", "rl")),
         "eval_gens": sum(plan[s]["eval_gens"] for s in ("sft", "dpo", "rl")),
     }
-    plan["total"]["cost"] = plan["total"]["train_cost"] + plan["total"]["eval_cost"]
-    plan["prices"] = prices
     plan["model"] = model
     return plan

claim_drafter/progress.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Live console progress, ETA and running cost for the three training stages.
+"""Live console progress, ETA and throughput for the three training stages.
 
 The cookbook's own console output is a Rich table printed on *every* step --
 2,416 tables for stage 1 alone, which scrolls the useful information away.
@@ -11,8 +11,8 @@ to build its MultiplexLogger. We drop PrettyPrintLogger from the returned
 multiplex and append a StageProgress in its place.
 
 Side effect worth knowing: this writes `runs/<stage>/progress.jsonl`, one record
-per step, carrying wall-clock, throughput and cumulative cost. metrics.jsonl has
-the learning curves but no timing or cost, so scripts/plot_runs.py reads both.
+per step, carrying wall-clock and throughput. metrics.jsonl has the learning
+curves but no timing, so scripts/plot_runs.py reads both.
 
 No new dependencies -- the bar is ~40 lines of ANSI rather than tqdm/rich.
 """
@@ -57,24 +57,15 @@ class StageProgress:
     tests can exercise the formatting.
     """
 
-    def __init__(self, stage, total_steps, log_dir, price_train,
+    def __init__(self, stage, total_steps, log_dir,
                  tokens_per_step=0, eval_every=0, save_every=0,
-                 eval_cost_per_round=0.0, eval_seconds_per_round=None,
-                 stream=None, cost_per_step=None):
+                 eval_seconds_per_round=None, stream=None):
         self.stage = stage
         self.total_steps = max(int(total_steps or 0), 0)
         self.log_dir = log_dir
-        self.price_train = price_train
         self.tokens_per_step = tokens_per_step
-        # DPO bills each pair twice -- once through the policy at the train
-        # rate, once through the frozen reference at the prefill rate -- so
-        # tokens x train-price under-reports it. The plan knows the real
-        # number; use it when given.
-        self.cost_per_step = (cost_per_step if cost_per_step is not None
-                              else tokens_per_step / 1e6 * price_train)
         self.eval_every = eval_every
         self.save_every = save_every
-        self.eval_cost_per_round = eval_cost_per_round
         self.eval_seconds_per_round = eval_seconds_per_round
         self.stream = stream or sys.stderr
         self.tty = hasattr(self.stream, "isatty") and self.stream.isatty()
@@ -84,8 +75,6 @@ class StageProgress:
         self.sec_per_step = None      # EMA, seconds
         self.steps_seen = 0
         self.evals_seen = 0
-        self.train_cost = 0.0
-        self.eval_cost = 0.0
         self._fh = None
         if log_dir:
             os.makedirs(log_dir, exist_ok=True)
@@ -139,10 +128,8 @@ class StageProgress:
         self.last_step_at = now
         self.steps_seen += 1
 
-        self.train_cost += self.cost_per_step
         if is_eval:
             self.evals_seen += 1
-            self.eval_cost += self.eval_cost_per_round
 
         elapsed = now - self.start
         remaining = self._eta(step)
@@ -163,8 +150,6 @@ class StageProgress:
             "steps           %d" % self.steps_seen,
             "wall clock      %s" % _fmt_duration(total),
             "eval rounds     %d" % self.evals_seen,
-            "estimated cost  $%.2f  (train $%.2f + eval $%.2f)"
-            % (self.train_cost + self.eval_cost, self.train_cost, self.eval_cost),
             "logs            %s" % self.log_dir,
         ])
         if self._fh:
@@ -207,7 +192,6 @@ class StageProgress:
         if self.sec_per_step:
             parts.append("%.1fs/it" % self.sec_per_step)
         parts.append("%s<%s" % (_fmt_duration(elapsed), _fmt_duration(remaining)))
-        parts.append("$%.2f" % (self.train_cost + self.eval_cost))
 
         label, away = self._next_event(step)
         if label and self.sec_per_step:
@@ -250,9 +234,6 @@ class StageProgress:
             "eta_s": round(remaining, 1) if remaining is not None else None,
             "sec_per_step": round(self.sec_per_step, 3) if self.sec_per_step else None,
             "tokens_cum": self.steps_seen * self.tokens_per_step,
-            "cost_train_usd": round(self.train_cost, 4),
-            "cost_eval_usd": round(self.eval_cost, 4),
-            "cost_total_usd": round(self.train_cost + self.eval_cost, 4),
             "is_eval": is_eval,
         }
         if loss is not None:
@@ -275,25 +256,15 @@ class StageProgress:
         self.stream.flush()
 
 
-def announce(stage, total_steps, tokens_total, price_train, eval_rounds,
-             eval_gens_per_round, eval_cost, stream=None, train_cost=None):
-    """Print the pre-flight estimate for a stage before the first step runs.
-
-    `train_cost` comes from the plan when the caller has it. Recomputing it from
-    tokens is only right when every token is billed at the train rate, which is
-    not true for DPO: each pair also goes through the frozen reference model at
-    the prefill rate, so recomputing understated stage 2 by about 45%.
-    """
+def announce(stage, total_steps, tokens_total, eval_rounds,
+             eval_gens_per_round, stream=None):
+    """Print the pre-flight summary for a stage before the first step runs."""
     stream = stream or sys.stderr
-    if train_cost is None:
-        train_cost = tokens_total / 1e6 * price_train
     lines = [
         "STAGE %s -- starting" % stage.upper(),
         "steps            %s" % "{:,}".format(total_steps),
         "train tokens     %s" % "{:,}".format(int(tokens_total)),
         "eval             %d rounds x %d generations" % (eval_rounds, eval_gens_per_round),
-        "estimated cost   $%.2f  (train $%.2f + eval $%.2f)"
-        % (train_cost + eval_cost, train_cost, eval_cost),
         "duration         unknown until the first steps land; the bar below",
         "                 shows a live ETA once it has measured s/it",
     ]
@@ -305,10 +276,9 @@ def announce(stage, total_steps, tokens_total, price_train, eval_rounds,
     stream.flush()
 
 
-def install(stage, total_steps, log_dir, price_train, tokens_per_step=0,
-            eval_every=0, save_every=0, eval_cost_per_round=0.0,
-            eval_seconds_per_round=None, replace_pretty=True,
-            cost_per_step=None):
+def install(stage, total_steps, log_dir, tokens_per_step=0,
+            eval_every=0, save_every=0,
+            eval_seconds_per_round=None, replace_pretty=True):
     """Attach a StageProgress to whatever logger the cookbook builds next.
 
     Must be called before the trainer's `main()`. Returns the StageProgress so
@@ -319,11 +289,9 @@ def install(stage, total_steps, log_dir, price_train, tokens_per_step=0,
 
     progress = StageProgress(
         stage=stage, total_steps=total_steps, log_dir=log_dir,
-        price_train=price_train, tokens_per_step=tokens_per_step,
+        tokens_per_step=tokens_per_step,
         eval_every=eval_every, save_every=save_every,
-        eval_cost_per_round=eval_cost_per_round,
         eval_seconds_per_round=eval_seconds_per_round,
-        cost_per_step=cost_per_step,
     )
 
     original = ml_log.setup_logging

scripts/estimate_cost.py

@@ -1,79 +0,0 @@
-#!/usr/bin/env python3
-"""Cost the full SFT -> DPO -> RL pipeline against a fixed budget.
-
-    python3 scripts/estimate_cost.py [--budget 120] [--model Qwen/Qwen3.5-9B]
-
-All arithmetic lives in claim_drafter/budget.py, which the trainers also import
-for their eval cadences -- so what this prints is what the run will actually do.
-Before that refactor this script hard-coded an eval assumption (5 rounds of 150
-generations) that the trainers did not share (96 rounds of 440), and understated
-the total by about half the budget.
-
-Caveat carried from the original: Tinker documents the "train" meter as "forward
-and backward pass for gradient computation" but does not state whether it bills
-every token or only weight=1 tokens. Everything below assumes the full sequence,
-which is the conservative reading.
-"""
-
-import argparse
-import os
-import sys
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from claim_drafter import budget as B
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--budget", type=float, default=120.0)
-    ap.add_argument("--model", default="Qwen/Qwen3.5-9B")
-    ap.add_argument("--sft-epochs", type=int, default=2)
-    ap.add_argument("--dpo-epochs", type=int, default=1)
-    ap.add_argument("--rl-prompts", type=int, default=300)
-    ap.add_argument("--rl-group", type=int, default=8, help="GRPO samples per prompt")
-    args = ap.parse_args()
-
-    plan = B.stage_plan(model=args.model, sft_epochs=args.sft_epochs,
-                        dpo_epochs=args.dpo_epochs, rl_prompts=args.rl_prompts,
-                        rl_group=args.rl_group)
-    p = plan["prices"]
-
-    print("Model: %s   (train $%.3f / prefill $%.3f / sample $%.3f per 1M)"
-          % (args.model, p["train"], p["prefill"], p["sample"]))
-    print("\n%-34s %8s %13s %8s %8s" % ("stage", "steps", "train tokens", "train $", "eval $"))
-    print("-" * 76)
-
-    labels = {
-        "sft": "1. SFT  (%d epochs)" % args.sft_epochs,
-        "dpo": "2. DPO  (%d epoch, %s pairs)" % (args.dpo_epochs,
-                                                 "{:,}".format(plan["dpo"]["pairs"])),
-        "rl":  "3. RL   (%d prompts x %d samples)" % (args.rl_prompts, args.rl_group),
-    }
-    for stage in ("sft", "dpo", "rl"):
-        r = plan[stage]
-        print("%-34s %8s %13s %8.2f %8.2f"
-              % (labels[stage], "{:,}".format(r["steps"]),
-                 "{:,}".format(r["train_tokens"]), r["train_cost"], r["eval_cost"]))
-
-    t = plan["total"]
-    print("-" * 76)
-    print("%-34s %8s %13s %8.2f %8.2f"
-          % ("TOTAL", "{:,}".format(t["steps"]), "", t["train_cost"], t["eval_cost"]))
-    print("%-34s %39s %8.2f" % ("", "combined", t["cost"]))
-
-    print("\nEval volume: %s generations across %d rounds"
-          % ("{:,}".format(t["eval_gens"]),
-             sum(plan[s]["eval_rounds"] for s in ("sft", "dpo", "rl"))))
-    print("Budget $%.2f -> %s"
-          % (args.budget,
-             ("fits, $%.2f left" % (args.budget - t["cost"])) if t["cost"] <= args.budget
-             else ("OVER by $%.2f" % (t["cost"] - args.budget))))
-    print("Stage share: SFT %.0f%%  DPO %.0f%%  RL %.0f%%"
-          % tuple(100.0 * (plan[s]["train_cost"] + plan[s]["eval_cost"]) / t["cost"]
-                  for s in ("sft", "dpo", "rl")))
-
-    return 0 if t["cost"] <= args.budget else 1
-
-
-if __name__ == "__main__":
-    sys.exit(main())

scripts/plot_runs.py

@@ -1,12 +1,12 @@
 #!/usr/bin/env python3
 """Turn the logs from every stage into charts, once all three have finished.
 
-    python3 scripts/plot_runs.py [--runs runs] [--out runs/graphs] [--budget 120]
+    python3 scripts/plot_runs.py [--runs runs] [--out runs/graphs]
 
 Reads two files per stage:
 
     runs/<stage>/metrics.jsonl   the cookbook's learning curves
-    runs/<stage>/progress.jsonl  wall-clock, throughput and spend (progress.py)
+    runs/<stage>/progress.jsonl  wall-clock and throughput (progress.py)
 
 Stages that have not run yet are skipped with a note rather than an error, so
 this is safe to run mid-pipeline -- you just get fewer panels.
@@ -195,14 +195,10 @@ def plot_slice_breakdown(plt, data, out):
     return save(fig, out, "03_slice_breakdown.png")
 
 
-def plot_timing_and_cost(plt, data, out, budget_usd):
-    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
-    ax_rate, ax_cost, ax_wall = axes
+def plot_timing(plt, data, out):
+    fig, axes = plt.subplots(1, 2, figsize=(11, 4))
+    ax_rate, ax_wall = axes
 
-    # Spend is cumulative across the whole pipeline, so each stage continues
-    # where the previous one stopped -- on both axes. Restarting x at 0 per
-    # stage made a single long run look like three short ones.
-    cost_base = x_base = 0.0
     for stage in STAGES:
         if stage not in data:
             continue
@@ -210,31 +206,20 @@ def plot_timing_and_cost(plt, data, out, budget_usd):
         xs, ys = series(p, "sec_per_step")
         if ys:
             ax_rate.plot(xs, ys, lw=1, label=stage)
-        cxs, cys = series(p, "cost_total_usd")
-        if cys:
-            ax_cost.plot([x_base + i for i in range(len(cys))],
-                         [cost_base + y for y in cys], lw=1.5, label=stage)
-            cost_base += cys[-1]
-            x_base += len(cys)
         exs, eys = series(p, "elapsed_s")
         if eys:
             ax_wall.plot(exs, [y / 60.0 for y in eys], lw=1.2, label=stage)
-    if budget_usd:
-        ax_cost.axhline(budget_usd, ls="--", c="crimson", lw=1,
-                        label="budget $%.0f" % budget_usd)
 
     ax_rate.set_title("Seconds per step")
     ax_rate.set_xlabel("step")
-    ax_cost.set_title("Cumulative estimated spend ($)")
-    ax_cost.set_xlabel("logged step (all stages)")
     ax_wall.set_title("Wall clock (minutes)")
     ax_wall.set_xlabel("step")
     for ax in axes:
         ax.grid(alpha=0.3)
         if ax.get_legend_handles_labels()[0]:
             ax.legend(fontsize=8)
-    fig.suptitle("Throughput, spend and wall clock")
-    return save(fig, out, "04_timing_and_cost.png")
+    fig.suptitle("Throughput and wall clock")
+    return save(fig, out, "04_timing.png")
 
 
 def plot_slice_trajectories(plt, data, out, metric="parse_rate"):
@@ -274,10 +259,11 @@ def plot_slice_trajectories(plt, data, out, metric="parse_rate"):
 
 
 def plot_convergence(plt, data, out):
-    """Where the held-out loss stopped improving, and what the rest cost.
+    """Where the held-out loss stopped improving, and how much run remained after.
 
     This is the panel that answers 'should I have trained this long', which the
-    aggregate learning curve does not.
+    aggregate learning curve does not. The second panel measures the run after
+    convergence in wall-clock time.
     """
     if "sft" not in data:
         return None
@@ -294,13 +280,13 @@ def plot_convergence(plt, data, out):
     conv_step = xs[conv_i]
 
     prog = data["sft"]["progress"]
-    cost_x = [r.get("step", 0) for r in prog if r.get("cost_total_usd") is not None]
-    cost_y = [r["cost_total_usd"] for r in prog if r.get("cost_total_usd") is not None]
-    total_cost = cost_y[-1] if cost_y else 0.0
-    cost_at_conv = 0.0
-    for sx, sy in zip(cost_x, cost_y):
+    t_x = [r.get("step", 0) for r in prog if r.get("elapsed_s") is not None]
+    t_y = [r["elapsed_s"] / 60.0 for r in prog if r.get("elapsed_s") is not None]
+    total_min = t_y[-1] if t_y else 0.0
+    min_at_conv = 0.0
+    for sx, sy in zip(t_x, t_y):
         if sx <= conv_step:
-            cost_at_conv = sy
+            min_at_conv = sy
 
     fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.2))
 
@@ -319,24 +305,24 @@ def plot_convergence(plt, data, out):
     ax.grid(alpha=0.25)
 
     ax = axes[1]
-    if cost_y:
-        ax.plot(cost_x, cost_y, color=C_SERIES_1, linewidth=2)
+    if t_y:
+        ax.plot(t_x, t_y, color=C_SERIES_1, linewidth=2)
         ax.axvline(conv_step, color=C_MID, linewidth=2, linestyle="--")
-        wasted = max(0.0, total_cost - cost_at_conv)
-        ax.fill_between([x for x in cost_x if x >= conv_step],
-                        [cost_at_conv] * sum(1 for x in cost_x if x >= conv_step),
-                        [y for x, y in zip(cost_x, cost_y) if x >= conv_step],
+        wasted = max(0.0, total_min - min_at_conv)
+        ax.fill_between([x for x in t_x if x >= conv_step],
+                        [min_at_conv] * sum(1 for x in t_x if x >= conv_step),
+                        [y for x, y in zip(t_x, t_y) if x >= conv_step],
                         color=C_WARN, alpha=0.18)
-        ax.annotate("$%.2f spent after convergence\n(of $%.2f total)" % (wasted, total_cost),
+        ax.annotate("%.0f min after convergence\n(of %.0f min total)" % (wasted, total_min),
                     xy=(0.42, 0.18), xycoords="axes fraction",
                     fontsize=10, color=C_WARN)
-    ax.set_title("Cumulative spend", color=C_INK)
+    ax.set_title("Cumulative wall clock", color=C_INK)
     ax.set_xlabel("training step", color=C_INK_2)
-    ax.set_ylabel("USD", color=C_INK_2)
+    ax.set_ylabel("minutes", color=C_INK_2)
     ax.grid(alpha=0.25)
 
     fig.suptitle("Did the run earn its length?", color=C_INK)
-    return save(fig, out, "06_convergence_and_spend.png")
+    return save(fig, out, "06_convergence.png")
 
 
 def plot_slice_deltas(plt, data, out, metric="reward"):
@@ -574,7 +560,7 @@ def verdict(data):
             notes.append(
                 "- **Held-out loss converged at step %d of %d** (%.0f%% of the run). "
                 "The remaining steps did not improve it, so a shorter run would "
-                "have reached the same place for less." % (xs[conv_i], xs[-1], frac * 100))
+                "have reached the same place." % (xs[conv_i], xs[-1], frac * 100))
 
     # 2. Did sampled-generation quality move?
     for key, label in (("parse_rate", "parse rate"), ("reward", "reward")):
@@ -668,33 +654,28 @@ def verdict(data):
     return lines + (notes or ["Nothing anomalous in the logged metrics."])
 
 
-def summarise(data, budget_usd):
+def summarise(data):
     lines = ["# Run summary", ""]
-    total_cost = total_time = 0.0
-    lines.append("| stage | steps | wall clock | est. cost | eval rounds |")
-    lines.append("|---|---|---|---|---|")
+    total_time = 0.0
+    lines.append("| stage | steps | wall clock | eval rounds |")
+    lines.append("|---|---|---|---|")
     for stage in STAGES:
         if stage not in data:
-            lines.append("| %s | _not run_ | | | |" % stage.upper())
+            lines.append("| %s | _not run_ | | |" % stage.upper())
             continue
         p = data[stage]["progress"]
         if not p:
-            lines.append("| %s | ? | ? | ? | ? |" % stage.upper())
+            lines.append("| %s | ? | ? | ? |" % stage.upper())
             continue
         last = p[-1]
         secs = last.get("elapsed_s") or 0.0
-        cost = last.get("cost_total_usd") or 0.0
         total_time += secs
-        total_cost += cost
-        lines.append("| %s | %d | %dh%02dm | $%.2f | %d |"
+        lines.append("| %s | %d | %dh%02dm | %d |"
                      % (stage.upper(), last.get("step", 0) + 1,
-                        int(secs // 3600), int(secs % 3600 // 60), cost,
+                        int(secs // 3600), int(secs % 3600 // 60),
                         sum(1 for r in p if r.get("is_eval"))))
-    lines += ["", "**Total** %dh%02dm, $%.2f estimated"
-              % (int(total_time // 3600), int(total_time % 3600 // 60), total_cost)]
-    if budget_usd:
-        lines[-1] += " against a $%.0f budget (%s)" % (
-            budget_usd, "within" if total_cost <= budget_usd else "OVER")
+    lines += ["", "**Total wall clock** %dh%02dm"
+              % (int(total_time // 3600), int(total_time % 3600 // 60))]
     lines += verdict(data)
     return "\n".join(lines) + "\n"
 
@@ -703,7 +684,6 @@ def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--runs", default="runs")
     ap.add_argument("--out", default="runs/graphs")
-    ap.add_argument("--budget", type=float, default=120.0)
     args = ap.parse_args()
 
     try:
@@ -726,7 +706,7 @@ def main():
         plot_learning_curves(plt, data, args.out),
         plot_eval_quality(plt, data, args.out),
         plot_slice_breakdown(plt, data, args.out),
-        plot_timing_and_cost(plt, data, args.out, args.budget),
+        plot_timing(plt, data, args.out),
         plot_slice_trajectories(plt, data, args.out),
         plot_convergence(plt, data, args.out),
         plot_slice_deltas(plt, data, args.out),
@@ -737,7 +717,7 @@ def main():
 
     summary_path = os.path.join(args.out, "summary.md")
     with open(summary_path, "w") as f:
-        f.write(summarise(data, args.budget))
+        f.write(summarise(data))
 
     print("\nWrote %d charts to %s/" % (sum(1 for w in written if w), args.out))
     for w in written:
@@ -745,7 +725,7 @@ def main():
             print("  " + w)
     print("  " + summary_path)
     print()
-    print(summarise(data, args.budget))
+    print(summarise(data))
 
 
 if __name__ == "__main__":

tests/test_budget.py

@@ -1,134 +0,0 @@
-"""Tests for the cost plan and the progress meter.
-
-The load-bearing test is `test_plan_fits_the_budget`. The pipeline previously
-shipped with an eval cadence that cost more than the training it was measuring
--- 96 rounds of 440 generations, far more than the training it measured. Nothing caught it because the estimator and the trainers held
-separate copies of the eval assumption. Now both read claim_drafter/budget.py
-and this test asserts the total, so the two cannot drift apart again.
-
-    pytest tests/ -q          # or: python3 tests/test_budget.py
-"""
-
-import json
-import os
-import sys
-import tempfile
-
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from claim_drafter import budget
-from claim_drafter.progress import StageProgress, _fmt_duration
-
-BUDGET_USD = 120.0
-
-
-class _Sink:
-    """Stand-in for a non-TTY stream, so the bar renders without ANSI."""
-
-    def __init__(self):
-        self.text = ""
-
-    def write(self, s):
-        self.text += s
-
-    def flush(self):
-        pass
-
-
-def test_plan_fits_the_budget():
-    total = budget.stage_plan()["total"]["cost"]
-    assert total <= BUDGET_USD, "plan costs $%.2f, over the $%.0f budget" % (total, BUDGET_USD)
-
-
-def test_eval_is_a_minority_of_spend():
-    # Eval exists to measure training, so it should not rival it. This is the
-    # specific failure the old cadence had: eval was ~55% of the total.
-    t = budget.stage_plan()["total"]
-    assert t["eval_cost"] < 0.15 * t["train_cost"], \
-        "eval is $%.2f against $%.2f of training" % (t["eval_cost"], t["train_cost"])
-
-
-def test_eval_volume_stays_bounded():
-    t = budget.stage_plan()["total"]
-    assert t["eval_gens"] < 2000, "%d eval generations is a wall-clock problem" % t["eval_gens"]
-
-
-def test_every_stage_evaluates_at_least_once():
-    plan = budget.stage_plan()
-    for stage in ("sft", "dpo", "rl"):
-        assert plan[stage]["eval_rounds"] >= 1, "%s never runs the evaluator" % stage
-
-
-def test_cheaper_model_costs_less():
-    big = budget.stage_plan(model="Qwen/Qwen3.5-9B")["total"]["cost"]
-    small = budget.stage_plan(model="Qwen/Qwen3.5-4B")["total"]["cost"]
-    assert small < big
-
-
-def test_unknown_model_falls_back_rather_than_crashing():
-    assert budget.prices_for("Nonexistent/Model-1B") == budget.PRICES["Qwen/Qwen3.5-9B"]
-
-
-def test_tokens_per_step_is_positive_everywhere():
-    plan = budget.stage_plan()
-    for stage in ("sft", "dpo", "rl"):
-        assert plan[stage]["tokens_per_step"] > 0, stage
-
-
-def test_progress_writes_one_record_per_step():
-    d = tempfile.mkdtemp()
-    p = StageProgress("sft", 10, d, 1.463, tokens_per_step=1000, stream=_Sink())
-    for step in range(5):
-        p.log_metrics({"train_mean_nll": 2.0 - step * 0.1}, step=step)
-    p.close()
-    with open(os.path.join(d, "progress.jsonl")) as f:
-        records = [json.loads(l) for l in f if l.strip()]
-    assert len(records) == 5
-    assert records[-1]["step"] == 4
-    assert records[-1]["cost_total_usd"] > records[0]["cost_total_usd"]
-
-
-def test_progress_flags_eval_rounds():
-    d = tempfile.mkdtemp()
-    p = StageProgress("sft", 4, d, 1.463, eval_cost_per_round=0.97, stream=_Sink())
-    p.log_metrics({"train_mean_nll": 1.0}, step=0)
-    p.log_metrics({"train_mean_nll": 1.0, "overall/reward": 0.6}, step=1)
-    p.close()
-    with open(os.path.join(d, "progress.jsonl")) as f:
-        records = [json.loads(l) for l in f if l.strip()]
-    assert records[0]["is_eval"] is False
-    assert records[1]["is_eval"] is True
-    assert records[1]["cost_eval_usd"] > 0
-
-
-def test_eval_steps_do_not_poison_the_eta():
-    # An eval step is 10-100x slower than a training step. Folding it into the
-    # per-step average would inflate the ETA for every remaining step.
-    d = tempfile.mkdtemp()
-    p = StageProgress("sft", 100, d, 1.463, stream=_Sink())
-    p.log_metrics({"train_mean_nll": 1.0}, step=0)
-    p.log_metrics({"train_mean_nll": 1.0}, step=1)
-    baseline = p.sec_per_step
-    p.log_metrics({"train_mean_nll": 1.0, "overall/reward": 0.5}, step=2)
-    assert p.sec_per_step == baseline, "eval round changed the per-step average"
-
-
-def test_duration_formatting():
-    assert _fmt_duration(45) == "45s"
-    assert _fmt_duration(750) == "12m30s"
-    assert _fmt_duration(11220) == "3h07m"
-    assert _fmt_duration(None) == "--"
-
-
-if __name__ == "__main__":
-    failures = 0
-    for name, fn in sorted(globals().items()):
-        if name.startswith("test_") and callable(fn):
-            try:
-                fn()
-                print("  PASS  %s" % name)
-            except AssertionError as e:
-                failures += 1
-                print("  FAIL  %s: %s" % (name, e))
-    print("\n%d failure(s)" % failures)
-    sys.exit(1 if failures else 0)

training/train_dpo.py

@@ -46,11 +46,11 @@ def main():
     from tinker_cookbook.preference.preference_datasets import ComparisonBuilderFromJsonl
     from tinker_cookbook.supervised.types import ChatDatasetBuilderCommonConfig
     from tinker_cookbook.renderers import TrainOnWhat
-    from claim_drafter import budget, progress
+    from claim_drafter import planning, progress
     from claim_drafter.evaluator import ClaimEvaluatorBuilder
 
-    plan = budget.stage_plan(model=args.model, dpo_epochs=args.epochs,
-                             dpo_batch=args.batch_size)["dpo"]
+    plan = planning.stage_plan(model=args.model, dpo_epochs=args.epochs,
+                               dpo_batch=args.batch_size)["dpo"]
 
     common = ChatDatasetBuilderCommonConfig(
         model_name_for_tokenizer=args.model,
@@ -65,15 +65,12 @@ def main():
 
     total_steps = args.max_steps or plan["steps"]
     progress.announce(
-        "dpo", total_steps, plan["train_tokens"], budget.prices_for(args.model)["train"],
-        plan["eval_rounds"], budget.N_SLICES * budget.GENS_PER_SLICE, plan["eval_cost"],
-        train_cost=plan["train_cost"])
+        "dpo", total_steps, plan["train_tokens"],
+        plan["eval_rounds"], planning.N_SLICES * planning.GENS_PER_SLICE)
     progress.install(
-        "dpo", total_steps, args.log_path, budget.prices_for(args.model)["train"],
+        "dpo", total_steps, args.log_path,
         tokens_per_step=plan["tokens_per_step"],
-        eval_every=plan["infrequent_eval_every"], save_every=plan["save_every"],
-        eval_cost_per_round=plan["eval_cost_per_round"],
-        cost_per_step=plan["train_cost"] / max(plan["steps"], 1))
+        eval_every=plan["infrequent_eval_every"], save_every=plan["save_every"])
 
     config = train_dpo.Config(
         log_path=args.log_path,

training/train_rl.py

@@ -49,24 +49,21 @@ def main():
 
     from tinker_cookbook.rl import train as rl_train
     from tinker_cookbook.rl.train import KLReferenceConfig
-    from claim_drafter import budget, progress
+    from claim_drafter import planning, progress
     from claim_drafter.evaluator import ClaimEvaluatorBuilder
     from claim_drafter.rl_env import ClaimDatasetBuilder
 
-    plan = budget.stage_plan(model=args.model, rl_prompts=args.prompts,
+    plan = planning.stage_plan(model=args.model, rl_prompts=args.prompts,
                              rl_group=args.group_size,
                              rl_batch=args.batch_size)["rl"]
 
     progress.announce(
-        "rl", plan["steps"], plan["train_tokens"], budget.prices_for(args.model)["train"],
-        plan["eval_rounds"], budget.N_SLICES * budget.GENS_PER_SLICE, plan["eval_cost"],
-        train_cost=plan["train_cost"])
+        "rl", plan["steps"], plan["train_tokens"],
+        plan["eval_rounds"], planning.N_SLICES * planning.GENS_PER_SLICE)
     progress.install(
-        "rl", plan["steps"], args.log_path, budget.prices_for(args.model)["train"],
+        "rl", plan["steps"], args.log_path,
         tokens_per_step=plan["tokens_per_step"],
-        eval_every=plan["eval_every"], save_every=plan["save_every"],
-        eval_cost_per_round=plan["eval_cost_per_round"],
-        cost_per_step=plan["train_cost"] / max(plan["steps"], 1))
+        eval_every=plan["eval_every"], save_every=plan["save_every"])
 
     config = rl_train.Config(
         log_path=args.log_path,

training/train_sft.py

@@ -3,11 +3,10 @@
 
     python3 training/train_sft.py --epochs 2
 
-Reads TINKER_API_KEY from .env. See docs/experiment.md for why each value is set.
+Reads TINKER_API_KEY from .env.
 
-NOTE: this has not been executed against the live API (no key at authoring time).
-Run `python3 scripts/preflight.py` first — it exercises the same code path for a
-few cents and will surface any signature drift before you spend the budget.
+Run `python3 scripts/preflight.py` first — it exercises the same code path on a
+couple of examples and will surface any signature drift before a full run.
 """
 
 import argparse
@@ -40,11 +39,11 @@ def main():
     from tinker_cookbook.supervised.data import FromConversationFileBuilder
     from tinker_cookbook.supervised.types import ChatDatasetBuilderCommonConfig
     from tinker_cookbook.renderers import TrainOnWhat
-    from claim_drafter import budget, progress
+    from claim_drafter import planning, progress
     from claim_drafter.evaluator import ClaimEvaluatorBuilder
 
-    plan = budget.stage_plan(model=args.model, sft_epochs=args.epochs,
-                             sft_batch=args.batch_size)["sft"]
+    plan = planning.stage_plan(model=args.model, sft_epochs=args.epochs,
+                               sft_batch=args.batch_size)["sft"]
 
     common = ChatDatasetBuilderCommonConfig(
         model_name_for_tokenizer=args.model,
@@ -60,27 +59,24 @@ def main():
         train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE,
     )
 
-    # Two evaluators on two cadences, because they cost three orders of
-    # magnitude apart. The NLL holdout is a single forward pass -- cheap enough
-    # to run often, and the only thing that shows overfitting, which
+    # Two evaluators on two cadences, because they differ by three orders of
+    # magnitude in runtime. The NLL holdout is a single forward pass -- fast
+    # enough to run often, and the only thing that shows overfitting, which
     # train_mean_nll cannot. The slice evaluator samples 88 full claim sets and
-    # is the expensive one, so it runs rarely. Running the slice evaluator on
-    # the NLL cadence is what put the pipeline at ~2x budget.
+    # is the slow one, so it runs rarely. Running the slice evaluator on the NLL
+    # cadence is what made an early run take roughly twice as long.
     slice_evals = [] if args.no_eval else [
         ClaimEvaluatorBuilder(model_name_for_tokenizer=args.model)
     ]
 
     total_steps = args.max_steps or plan["steps"]
     progress.announce(
-        "sft", total_steps, plan["train_tokens"], budget.prices_for(args.model)["train"],
-        plan["eval_rounds"], budget.N_SLICES * budget.GENS_PER_SLICE, plan["eval_cost"],
-        train_cost=plan["train_cost"])
+        "sft", total_steps, plan["train_tokens"],
+        plan["eval_rounds"], planning.N_SLICES * planning.GENS_PER_SLICE)
     progress.install(
-        "sft", total_steps, args.log_path, budget.prices_for(args.model)["train"],
+        "sft", total_steps, args.log_path,
         tokens_per_step=plan["tokens_per_step"],
-        eval_every=plan["infrequent_eval_every"], save_every=plan["save_every"],
-        eval_cost_per_round=plan["eval_cost_per_round"],
-        cost_per_step=plan["train_cost"] / max(plan["steps"], 1))
+        eval_every=plan["infrequent_eval_every"], save_every=plan["save_every"])
 
     config = train.Config(
         log_path=args.log_path,