| """scripts/p2/contrasts.py: declarative K0/K1/K2/K4/K6 driver. Pure CPU, fake runners, nothing heavy, no GPU, no Lean. |
| |
| Run: CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest tests/test_contrasts.py -q |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import shlex |
| from pathlib import Path |
|
|
| import pytest |
|
|
| from test_p2_support import gen_record, problem_record, write_jsonl, load_script |
|
|
| K = load_script("p2_contrasts", "scripts/p2/contrasts.py") |
| EI = K.EI |
| MEAN = 5000.0 |
| BUDGET = 8 * 5000 * 20 * 1 |
|
|
|
|
| @pytest.fixture |
| def world(tmp_workdir): |
| d = Path(tmp_workdir) |
| for name in ("goedel", "kimina"): |
| (d / name).mkdir() |
| (d / name / "config.json").write_text("{}") |
| write_jsonl(d / "pool.jsonl", [problem_record(f"p{i:02d}") for i in range(10)]) |
| write_jsonl(d / "mutated.jsonl", [problem_record(f"m{i:02d}", sig=f"(x : ℕ) (h : x = {i + 50}) : x = {i + 50}") for i in range(5)]) |
| write_jsonl(d / "clf.jsonl", [problem_record(f"c{i:02d}", sig=f"(x : ℕ) (h : x = {i + 70}) : x = {i + 70}") for i in range(5)]) |
| write_jsonl(d / "rnd.jsonl", [problem_record(f"r{i:02d}", sig=f"(x : ℕ) (h : x = {i + 90}) : x = {i + 90}") for i in range(5)]) |
| write_jsonl(d / "valid.jsonl", [problem_record(f"mf{i}", "minif2f", "valid", f"(y : ℕ) (h : y = {i + 100}) : y = {i + 100}") for i in range(4)]) |
| write_jsonl(d / "test.jsonl", [problem_record("mt0", "minif2f", "test", "(y : ℕ) (h : y = 777) : y = 777")]) |
| write_jsonl(d / "mobench.jsonl", [problem_record("mo0", "mobench", "test", "(z : ℕ) (h : z = 7) : z = 7")]) |
| rl = d / "rl.toml" |
| rl.write_text(Path(K.REPO / "configs" / "p4_cispo_lean.toml").read_text()) |
| return d |
|
|
|
|
| def base_argv(d, *contrasts, **kw): |
| a = ["--contrast", *contrasts, "--out-dir", str(d / "out"), "--goedel-model", str(d / "goedel"), "--kimina-model", str(d / "kimina"), |
| "--pool", str(d / "pool.jsonl"), "--eval-problems", str(d / "valid.jsonl"), "--protected-problems", str(d / "test.jsonl"), str(d / "mobench.jsonl"), |
| "--mean-completion-tokens", str(MEAN), "--token-budget", str(BUDGET), "--rl-config", str(d / "rl.toml"), "--n", "8", "--rounds", "1"] |
| for k, v in kw.items(): |
| a += [f"--{k.replace('_', '-')}"] + ([] if v is True else (list(map(str, v)) if isinstance(v, (list, tuple)) else [str(v)])) |
| return a |
|
|
|
|
| def plan_of(d, *contrasts, **kw): |
| args = K.parse_args(base_argv(d, *contrasts, **kw)) |
| return K.build_plan(args), args |
|
|
|
|
| class FakeRunner: |
| def __init__(self, fail=None): |
| self.calls, self.fail = [], fail |
|
|
| def run(self, argv, env, log): |
| self.calls.append(list(argv)) |
| if self.fail and self.fail in " ".join(argv): |
| return 7 |
| return 0 |
|
|
|
|
| |
|
|
|
|
| def test_a_test_split_scored_file_is_refused_and_there_is_no_escape_hatch(world, capsys): |
| argv = base_argv(world, "K1") |
| argv[argv.index("--eval-problems") + 1] = str(world / "test.jsonl") |
| assert K.main(argv + ["--dry-run"]) == K.EXIT_REFUSED |
| assert "TEST split" in capsys.readouterr().err |
| with pytest.raises(SystemExit): |
| K.parse_args(argv + ["--allow-test-selection"]) |
|
|
|
|
| def test_missing_protected_problems_and_unreadable_scored_file_are_refused(world, capsys): |
| argv = base_argv(world, "K1") |
| i = argv.index("--protected-problems") |
| assert K.main(argv[:i] + argv[i + 3:] + ["--dry-run"]) == K.EXIT_REFUSED |
| bad = base_argv(world, "K1") |
| bad[bad.index("--eval-problems") + 1] = str(world / "nope.jsonl") |
| assert K.main(bad + ["--dry-run"]) == K.EXIT_REFUSED |
| assert "cannot read the scored benchmark" in capsys.readouterr().err |
|
|
|
|
| def test_guard_passes_a_dev_split_and_every_ei_command_carries_protected_problems(world): |
| plan, args = plan_of(world, "K1", "K2", "K6", "K0", mutated_pool=world / "mutated.jsonl", classifier_pool=world / "clf.jsonl", random_pool=world / "rnd.jsonl") |
| assert plan.guard["splits"] == ["valid"] and plan.guard["allow_test_selection"] is False |
| ei_steps = [s for s in plan.steps if s.kind == "ei"] |
| assert ei_steps |
| for s in ei_steps: |
| i = s.argv.index("--protected-problems") |
| assert str(world / "test.jsonl") in s.argv[i:] and str(world / "mobench.jsonl") in s.argv[i:] |
|
|
|
|
| @pytest.mark.parametrize("kw,msg", [ |
| (dict(seeds=[0, 1, 2]), "exactly 2 seeds"), |
| ]) |
| def test_k0_needs_exactly_two_seeds(world, kw, msg): |
| with pytest.raises(K.ContrastError, match=msg): |
| plan_of(world, "K0", **kw) |
|
|
|
|
| def test_missing_mean_tokens_budget_and_pools_are_refused(world): |
| a = base_argv(world, "K1") |
| i = a.index("--mean-completion-tokens") |
| with pytest.raises(K.ContrastError, match="mean-completion-tokens"): |
| K.build_plan(K.parse_args(a[:i] + a[i + 2:])) |
| with pytest.raises(K.ContrastError, match="--mutated-pool"): |
| plan_of(world, "K2") |
| with pytest.raises(K.ContrastError, match="K4 needs EI arms|requires"): |
| plan_of(world, "K4") |
|
|
|
|
| |
|
|
|
|
| def test_budget_arithmetic(): |
| assert K.planned_tokens_ei(2, 10, 8, 100) == 16000 |
| assert K.derive_ei_problems(16000, 2, 8, 100) == 10 |
| assert K.derive_ei_problems(16999, 2, 8, 100) == 10 |
| with pytest.raises(K.ContrastError): |
| K.derive_ei_problems(10, 1, 8, 100) |
| assert K.derive_rl_steps(1_000_000, 16, 16, 1000) == 4 |
| assert K.derive_rl_steps(1, 16, 16, 1000) == 1 |
| assert K.token_match(100, 104, 0.05) and not K.token_match(100, 110, 0.05) |
|
|
|
|
| def test_k1_arms_are_matched_on_planned_tokens_and_share_the_prompt_set(world): |
| plan, args = plan_of(world, "K1") |
| ei = [a for a in plan.arms if a["arm"] == "ei_sft"] |
| rl = [a for a in plan.arms if a["arm"] == "rl"] |
| assert len(ei) == len(rl) == 3 and {a["seed"] for a in ei} == {0, 1, 2} |
| for e, r in zip(ei, rl): |
| assert e["tokens_planned"] == BUDGET and e["max_problems"] == 20 |
| |
| step_tokens = r["prompts_per_step"] * r["group"] * MEAN |
| assert abs(r["tokens_planned"] - BUDGET) <= step_tokens / 2 |
| assert r["band_file"].endswith("round_01/select/problems.jsonl") and f"/s{e['seed']}/" in r["band_file"] |
| rows = K.budget_table(plan, 0.5) |
| assert rows[0]["contrast"] == "K1" and rows[0]["matched"].startswith("planned") |
| cfgs = [s for s in plan.steps if s.kind == "write_config" and s.arm == "rl"] |
| text = cfgs[0].writes[next(iter(cfgs[0].writes))] |
| import tomllib |
| t = tomllib.loads(text) |
| assert t["train"]["optimizer_steps"] == rl[0]["optimizer_steps"] and t["train"]["seed"] == 0 |
| assert t["data"]["band_file"] == rl[0]["band_file"] |
|
|
|
|
| def test_set_toml_keys_is_section_aware_keeps_comments_and_refuses_missing_keys(): |
| src = 'a = 1 # top\n[x]\na = 2 # inner\nb = "s"\n[y]\na = 3\n' |
| out = K.set_toml_keys(src, {("x", "a"): 9, ("", "a"): 7, ("y", "a"): "z"}) |
| assert out == 'a = 7 # top\n[x]\na = 9 # inner\nb = "s"\n[y]\na = "z"\n' |
| with pytest.raises(K.ContrastError): |
| K.set_toml_keys(src, {("x", "nope"): 1}) |
|
|
|
|
| def test_sft_config_carries_the_seed_and_the_global_token_mean_setting(world): |
| plan, _ = plan_of(world, "K0") |
| cfgs = [s for s in plan.steps if s.kind == "write_config"] |
| import tomllib |
| seeds = sorted(tomllib.loads(next(iter(s.writes.values())))["train"]["seed"] for s in cfgs) |
| assert seeds == [0, 1] |
| assert tomllib.loads(next(iter(cfgs[0].writes.values())))["train"]["loss_normalization"] == "global_token_mean" |
|
|
|
|
| def test_measured_tokens_and_measure_verdict(world): |
| plan, args = plan_of(world, "K6", classifier_pool=world / "clf.jsonl", random_pool=world / "rnd.jsonl", seeds=[0]) |
| arms = {a["arm"]: a for a in plan.arms} |
| for name, n_tok in (("random", 100), ("classifier", 104)): |
| gf = Path(arms[name]["gen_files"][0]) |
| write_jsonl(gf, [dict(gen_record(f"p{i}", 0), completion_ids=list(range(n_tok))) for i in range(3)]) |
| assert K.measured_generated_tokens([arms["random"]["gen_files"][0]]) == 300 |
| assert K.measure(plan, args, say=lambda *_: None) == K.EXIT_OK |
| write_jsonl(Path(arms["classifier"]["gen_files"][0]), [dict(gen_record(f"p{i}", 0), completion_ids=list(range(150))) for i in range(3)]) |
| assert K.measure(plan, args, say=lambda *_: None) == K.EXIT_UNMATCHED |
|
|
|
|
| |
|
|
|
|
| def test_ei_command_round_trips_through_expert_iterations_own_parser(world): |
| plan, args = plan_of(world, "K1") |
| s = next(x for x in plan.steps if x.kind == "ei") |
| cfg = EI.parse_args(s.argv[2:]) |
| assert cfg.max_problems == 20 and cfg.rounds == 1 and cfg.n == 8 and cfg.seed == 0 |
| assert cfg.protected_problems and cfg.allow_test_selection is False |
| assert cfg.train_config.endswith("ei_sft_s0.sft.toml") |
|
|
|
|
| def test_k2_is_ready_now_that_the_temperature_flags_exist(world, capsys): |
| """K2 was BLOCKED until lean_gen --temperature / expert_iteration --gen-temperature landed (merged 2026-10-19). |
| The driver detects the flags from the real argparse definitions, so K2 must now plan as ready.""" |
| plan, args = plan_of(world, "K2", mutated_pool=world / "mutated.jsonl") |
| assert {a["arm"] for a in plan.arms} == {"temp_matched", "mutated"} |
| assert all(a["temperature"] == 1.2 for a in plan.arms) |
| assert any("--gen-temperature" in p and ": present" in p for p in plan.prerequisites) |
| assert not any("MISSING (pending)" in p for p in plan.prerequisites) |
| assert not K.blocked_steps(plan) |
| K.dry_run(plan, args) |
| out = capsys.readouterr().out |
| assert "BLOCKED" not in out and "--gen-temperature 1.2" in out |
|
|
|
|
| def test_k4_alpha_exports_scale_lora_alpha_and_reuse_the_arms_own_eval_for_k1_0(world): |
| plan, args = plan_of(world, "K1", "K4", seeds=[0], k4_from=["K1"]) |
| ex = [s for s in plan.steps if s.kind == "alpha_export"] |
| assert len(ex) == 2 |
| alphas = sorted(float(s.argv[s.argv.index("--alpha") + 1]) for s in ex) |
| assert alphas == [64.0, 89.6] |
| assert all("P2_EXPORT_BASE_SNAPSHOT" in s.env and s.argv[s.argv.index("--ckpt") + 1].startswith(EI.LATEST) for s in ex) |
| rep = next(s for s in plan.steps if s.sid.startswith("K4.") and s.kind == "report") |
| assert rep.argv.count("--run") == 3 and "@a1.0" in " ".join(rep.argv) |
| evs = [s for s in plan.steps if s.kind == "eval" and s.contrast == "K4"] |
| assert len(evs) == 2 and all("--model-rev" in " ".join(s.argv) and "-K4-K1-ei_sft-s0-a0." in " ".join(s.argv) for s in evs) |
| assert any("RL arms are not LoRA" in p for p in plan.prerequisites) |
|
|
|
|
| def test_k0_report_pairs_the_two_seeds_and_k1_reports_per_seed(world): |
| plan, _ = plan_of(world, "K0") |
| reps = [s for s in plan.steps if s.kind == "report"] |
| assert len(reps) == 1 and "ei_s0" in " ".join(reps[0].argv) and "ei_s1" in " ".join(reps[0].argv) |
| plan, _ = plan_of(world, "K1", pkpo_config=world / "rl.toml") |
| reps = [s for s in plan.steps if s.kind == "report"] |
| assert len(reps) == 6 |
| assert {a["arm"] for a in plan.arms} == {"ei_sft", "rl", "rl_pkpo"} |
|
|
|
|
| |
|
|
|
|
| def test_dry_run_prints_commands_budget_and_touches_nothing(world, capsys): |
| out = world / "out" |
| assert K.main(base_argv(world, "K1", "K6", classifier_pool=world / "clf.jsonl", random_pool=world / "rnd.jsonl") + ["--dry-run"]) == K.EXIT_OK |
| text = capsys.readouterr().out |
| assert "expert_iteration.py" in text and "grpo.py" in text and "eval_report.py" in text and "# budget K1" in text |
| assert not out.exists() |
|
|
|
|
| def test_print_manifest_is_valid_json_with_the_bookkeeping(world, capsys): |
| assert K.main(base_argv(world, "K1", seeds=[0]) + ["--print-manifest"]) == K.EXIT_OK |
| m = json.loads(capsys.readouterr().out) |
| assert m["guard"]["allow_test_selection"] is False and m["budget"][0]["contrast"] == "K1" |
| assert {s["kind"] for s in m["steps"]} >= {"write_config", "ei", "rl", "eval", "report"} |
| assert m["inputs"]["mean_completion_tokens"] == MEAN and len(m["plan_sha"]) == 64 |
|
|
|
|
| def test_run_materialises_configs_resumes_per_step_and_refuses_a_changed_step(world): |
| plan, args = plan_of(world, "K0") |
|
|
| class R(FakeRunner): |
| def run(self, argv, env, log): |
| super().run(argv, env, log) |
| if "expert_iteration" in " ".join(argv): |
| m = Path(argv[argv.index("--run-dir") + 1]) / "round_01" / "round_manifest.json" |
| m.parent.mkdir(parents=True, exist_ok=True) |
| m.write_text("{}") |
| if "eval_report" in " ".join(argv): |
| rd = Path(argv[argv.index("--out-dir") + 1]) |
| rd.mkdir(parents=True, exist_ok=True) |
| (rd / "report.json").write_text("{}") |
| return 0 |
|
|
| fr = R() |
| assert K.run_plan(plan, args, fr, say=lambda *_: None) == K.EXIT_OK |
| kinds = ["ei" if "expert_iteration" in " ".join(c) else "report" for c in fr.calls] |
| assert kinds == ["ei", "ei", "report"] |
| assert (Path(args.out_dir) / "contrasts_manifest.json").exists() |
| assert all(Path(p).exists() for s in plan.steps for p in s.writes) |
| fr2 = R() |
| assert K.run_plan(plan, args, fr2, say=lambda *_: None) == K.EXIT_OK and fr2.calls == [] |
| plan2, args2 = plan_of(world, "K0", n=4) |
| assert K.run_plan(plan2, args2, R(), say=lambda *_: None) == K.EXIT_FAIL |
|
|
|
|
| def test_a_failing_step_stops_the_chain_and_leaves_no_marker(world): |
| plan, args = plan_of(world, "K0") |
| fr = FakeRunner(fail="expert_iteration") |
| assert K.run_plan(plan, args, fr, say=lambda *_: None) == K.EXIT_FAIL and len(fr.calls) == 1 |
| assert not list((Path(args.out_dir) / "steps").glob("*.ei.done.json")) |
|
|
