diff --git a/src/openultrasast/cli.py b/src/openultrasast/cli.py index fa8b922..7479bfd 100644 --- a/src/openultrasast/cli.py +++ b/src/openultrasast/cli.py @@ -223,6 +223,12 @@ def main(argv: list[str] | None = None) -> int: command.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) command.add_argument("--slice", choices=SLICE_NAMES, default="all") command.add_argument("--json", action="store_true") + if name == "baseline": + command.add_argument( + "--resume", + action="store_true", + help="keep the pairs a previous run of this baseline already measured and paid for", + ) if name == "candidates": command.add_argument("--generator", choices=("ir", "ruleset"), default="ir", help="which pass enumerates the candidate sites") if name in {"baseline", "round", "score"}: @@ -1352,6 +1358,7 @@ def _learning_run(args: argparse.Namespace, cases: Sequence[PairCase], taxonomy: slices=(args.slice,) if args.slice != "all" else DEFAULT_SLICES, spent_usd=spent, thinking=config.learning.thinking, + resume=args.resume, ) print(f"learning baseline cost: ${report.cost_usd:.4f}") print(f"learning baseline {model}: {len(report.floors)} families, k={report.k_runs} -> {out / 'baseline'}") diff --git a/src/openultrasast/learning/rounds.py b/src/openultrasast/learning/rounds.py index 68bbfb7..cc08b1e 100644 --- a/src/openultrasast/learning/rounds.py +++ b/src/openultrasast/learning/rounds.py @@ -103,6 +103,7 @@ def run_baseline( spent_usd: Callable[[], float] | None = None, cutoff: str = "", thinking: bool = False, + resume: bool = False, ) -> BaselineReport: """Clone one detector into every family, measure each family **with its own configuration**, write the artifacts. @@ -120,28 +121,40 @@ def run_baseline( write_default_configs(configs_dir, taxonomy, prompt=prompt or _SYSTEM_PROMPT, version="0") configs = load_family_configs(configs_dir, taxonomy) selected = [case for case in cases if case.slice in set(slices)] + # Every pair is written down as it is measured. Round zero over the vfc slice is 1760 detector runs and hours of + # paid calls; accumulating in memory and persisting once at the end means a timeout, a kill or a machine going + # away costs all of it. An exception already cost only one pair; nothing else did. + checkpoint = out_dir / "baseline" / _artifact_key(model, slices, thinking) / "scores.jsonl" + done = _checkpointed(checkpoint) if resume else {} + if not resume and checkpoint.exists(): + checkpoint.unlink() scores = [] for case in selected: name = _family_of(case, taxonomy) config = configs.get(name) if config is None: continue # a family with no configuration is not measured rather than measured with someone else's + if case.name in done: + scores.append(done[case.name]) + continue try: - scores.append(score_case(case, scan_factory(config), taxonomy=taxonomy, runs=k_runs, family=name)) + score = score_case(case, scan_factory(config), taxonomy=taxonomy, runs=k_runs, family=name) + scores.append(score) + _append_score(checkpoint, score) except Exception: # noqa: BLE001 — a provider hiccup costs one pair, not the whole run # Round zero is the most expensive thing this harness does. Raising here threw away everything already # measured and paid for, so the only honest response to a flaky endpoint was to run it all again. - scores.append( - PairFamilyScore( - pair=case.name, - family=name, - slice=case.slice, - fix_date=case.fix_date, - runs=(), - outcome="unscorable", - unscorable_reason="detector_unreachable", - ) + unreachable = PairFamilyScore( + pair=case.name, + family=name, + slice=case.slice, + fix_date=case.fix_date, + runs=(), + outcome="unscorable", + unscorable_reason="detector_unreachable", ) + scores.append(unreachable) + _append_score(checkpoint, unreachable) metrics = aggregate(scores, taxonomy=taxonomy, cutoff=cutoff) floors = {name: _floor(name, block, [item for item in scores if item.family == name], k_runs) for name, block in metrics.items()} report = BaselineReport( @@ -160,6 +173,59 @@ def run_baseline( return report +def _append_score(path: Path, score: PairFamilyScore) -> None: + """One line per pair, flushed, so an interrupted run keeps what it measured.""" + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "pair": score.pair, + "family": score.family, + "slice": score.slice, + "fix_date": score.fix_date, + "outcome": score.outcome, + "runs": list(score.runs), + "flips": score.flips, + "other_findings_vuln": score.other_findings_vuln, + "other_findings_fixed": score.other_findings_fixed, + "fabricated_vuln": score.fabricated_vuln, + "fabricated_fixed": score.fabricated_fixed, + "rung": score.rung, + "unscorable_reason": score.unscorable_reason, + } + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload, sort_keys=True) + "\n") + handle.flush() + + +def _checkpointed(path: Path) -> dict[str, PairFamilyScore]: + """What a previous run of this exact baseline already measured and paid for.""" + if not path.is_file(): + return {} + out: dict[str, PairFamilyScore] = {} + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + continue # a line half-written when the run died is not a measurement + out[str(row["pair"])] = PairFamilyScore( + pair=str(row["pair"]), + family=str(row["family"]), + slice=str(row.get("slice", "")), + fix_date=str(row.get("fix_date", "")), + runs=tuple(row.get("runs", ())), + outcome=row["outcome"], + flips=int(row.get("flips", 0)), + other_findings_vuln=int(row.get("other_findings_vuln", 0)), + other_findings_fixed=int(row.get("other_findings_fixed", 0)), + fabricated_vuln=int(row.get("fabricated_vuln", 0)), + fabricated_fixed=int(row.get("fabricated_fixed", 0)), + rung=row.get("rung", "suspicion"), + unscorable_reason=row.get("unscorable_reason"), + ) + return out + + def score_case(case: PairCase, scan: Scan, *, taxonomy: FamilyTaxonomy, runs: int, family: str | None = None) -> PairFamilyScore: """Run one detector over both sides of one pair ``runs`` times and score it for its family. diff --git a/tests/test_learning_rounds.py b/tests/test_learning_rounds.py index 0b4d827..04e554a 100644 --- a/tests/test_learning_rounds.py +++ b/tests/test_learning_rounds.py @@ -5,6 +5,8 @@ from __future__ import annotations import json from pathlib import Path +import pytest + from openultrasast.benchmark import ExpectedFinding from openultrasast.findings import StaticFinding from openultrasast.learning.families import load_families @@ -295,3 +297,81 @@ def test_two_slices_on_one_model_are_two_baselines_not_one_overwriting_the_other table = compare_baselines(out) assert set(table) == {"deepseek-v4-flash", "deepseek-v4-flash (vfc)"} assert "memory" in table["deepseek-v4-flash (vfc)"] and "injection" in table["deepseek-v4-flash"] + + +def test_round_zero_keeps_every_pair_it_has_already_paid_for(tmp_path: Path) -> None: + """A four-hour run that is interrupted must not lose four hours of paid calls. + + Measured the hard way: the 176-pair vfc baseline hit its timeout and wrote nothing, because scores were + accumulated in memory and persisted once at the end. An exception was already handled; a timeout, a kill and a + machine going away were not, and they cost the same money.""" + from openultrasast.learning.rounds import run_baseline + + taxonomy = load_families() + out = tmp_path / "out" + scored: list[str] = [] + + def build(_config): # type: ignore[no-untyped-def] + def scan(root: Path) -> list[StaticFinding]: + scored.append(root.name) + # K=3 means three vulnerable-side scans per pair, so this gives up part way through the third pair + if len([item for item in scored if item == "vuln"]) > 6: + raise KeyboardInterrupt("the operator gave up") + return _steady(root) + + return scan + + cases = [_case(tmp_path, name) for name in ("a", "b", "c", "d")] + with pytest.raises(KeyboardInterrupt): + run_baseline( + cases, + taxonomy=taxonomy, + configs_dir=tmp_path / "configs", + scan_factory=build, + model="m", + out_dir=out, + k_runs=3, + ) + checkpoint = out / "baseline" / "m" / "scores.jsonl" + assert checkpoint.is_file(), "the pairs already measured were not written down" + rows = [json.loads(line) for line in checkpoint.read_text().splitlines() if line.strip()] + assert [row["pair"] for row in rows] == ["a", "b"] + assert all(row["outcome"] == "pair_correct" for row in rows) + + +def test_a_resumed_baseline_does_not_pay_twice(tmp_path: Path) -> None: + from openultrasast.learning.rounds import run_baseline + + taxonomy = load_families() + out = tmp_path / "out" + cases = [_case(tmp_path, name) for name in ("a", "b")] + run_baseline( + cases, + taxonomy=taxonomy, + configs_dir=tmp_path / "configs", + scan_factory=_always(_steady), + model="m", + out_dir=out, + k_runs=3, + ) + calls: list[str] = [] + + def counting(_config): # type: ignore[no-untyped-def] + def scan(root: Path) -> list[StaticFinding]: + calls.append(root.name) + return _steady(root) + + return scan + + report = run_baseline( + cases, + taxonomy=taxonomy, + configs_dir=tmp_path / "configs", + scan_factory=counting, + model="m", + out_dir=out, + k_runs=3, + resume=True, + ) + assert calls == [], "every pair was already measured and paid for" + assert report.metrics["injection"]["scorable"] == 2