From fbc3220bbe1c8a4f8afa0a724fd41a51b390dd22 Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:04:33 +0200 Subject: [PATCH 1/2] feat(learn): A/B experiment runner -- registered manifests, paired arms, bootstrap and McNemar (decision-engine task 7) learn/experiments.py: YAML manifests under plane/experiments/.yaml, digest-checked at register, run and analyse time; `experiment`, `arm_outcome` and `experiment_result` rows; paired arms in a seeded repository order with a per-unit coin; arm A replay-only at $0; canary replicates under a fresh cache salt; repository-cluster paired bootstrap, exact McNemar, two O'Brien-Fleming looks, adopt/reject/equivalent/ inconclusive. `ousast learn experiment register|units|run|analyse`. Program: `ProgramSpec.sampling = retrieval_ensemble` (k samples at temperature 0 over k disjoint retrieved example sets; default unchanged); `Caller.salt`; `retrieve(exclude=)`. verify: OUSAST_VERIFY_CALLERS=off. Manifests exp-001-known-callers (plane arms, units pending) and exp-002-retrieval-ensemble (slice-2 incumbents, ceiling $4) -- not run. 18 scripted tests; module audit regenerated (151 modules). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo --- .kiro/specs/learned-decision-engine/tasks.md | 28 +- .kiro/steering/overview.md | 4 +- .../measurements/2026-09-08-module-audit.json | 29 +- docs/decision-engine.md | 1 + docs/memory.md | 5 +- plane/experiments/exp-001-known-callers.yaml | 34 + .../exp-002-retrieval-ensemble.yaml | 41 + src/openultrasast/learn/cli.py | 113 +++ src/openultrasast/learn/experiments.py | 898 ++++++++++++++++++ src/openultrasast/learn/program.py | 65 +- src/openultrasast/learn/retrieve.py | 8 +- src/openultrasast/plane/tasks/verify.py | 3 +- tests/test_learn_experiments.py | 485 ++++++++++ tests/test_plane_increment.py | 2 +- 14 files changed, 1681 insertions(+), 35 deletions(-) create mode 100644 plane/experiments/exp-001-known-callers.yaml create mode 100644 plane/experiments/exp-002-retrieval-ensemble.yaml create mode 100644 src/openultrasast/learn/experiments.py create mode 100644 tests/test_learn_experiments.py diff --git a/.kiro/specs/learned-decision-engine/tasks.md b/.kiro/specs/learned-decision-engine/tasks.md index 827f37b..941324c 100644 --- a/.kiro/specs/learned-decision-engine/tasks.md +++ b/.kiro/specs/learned-decision-engine/tasks.md @@ -116,12 +116,28 @@ a model or embedding call. - _Requirements: 5.1, 5.2, 8.3_ - [ ] 7. Experiments (exp-001 and program variants need budget go-ahead) - - `learn/experiments.py`, `register`, `experiment-outcome`, Run annotation checks, `OUSAST_VERIFY_CALLERS`; - exp-001 (known-callers line on/off, ~$7.7 expected, ceiling $15) registered, run and recorded -- an - inconclusive result is pre-registered as likely. - - Program-variant experiments (paired on candidates, repository cluster, arm A replayed from cache at $0): the - machinery and its test; the first variant (`k = 5` vs `k = 1`, or the confidence source) **(budget go-ahead: - ~$0.9-1.35 per arm-B pass of 300 candidates, ceiling $2 each)**. + - [x] 7.1 Machinery (no model in tests): `learn/experiments.py` -- YAML manifests under `plane/experiments/.yaml` + (hypothesis, arms as program-spec overrides or plane `task_env`, unit, pairing, repository cluster, metrics with + `[label=]` selectors and targets, MDE, two looks with O'Brien-Fleming boundaries, budget ceiling, families, folds, + frozen units file), `register` (`experiment` row with both digests; refuses an uncommitted, modified or re-declared + manifest), `units` (freeze the families' evaluation candidates, paired by candidate across pins; no model), `run` + (paired arms in a seeded repository order, a seeded coin per unit for which arm goes first, `arm_outcome` rows, + resumable, arm A replay-only at $0, canary replicates under a fresh cache salt, the first look during the run), + `analyse` (paired repository bootstrap, exact McNemar, the looks, adopt / reject / equivalent / inconclusive, + `experiment_result` row + `benchmarks/experiments//result.json`), `ousast learn experiment + register|units|run|analyse`; `OUSAST_VERIFY_CALLERS=off` in `plane/tasks/verify.py`; the retrieval-ensemble + sampling variant (`ProgramSpec.sampling = retrieval_ensemble`: `k` samples at temperature 0 over `k` disjoint + retrieved example sets; default unchanged). Tests: `test_learn_experiments.py` (18, scripted clients, FileStore). + - [x] 7.2 Registered, not run: `exp-001-known-callers` (plane arms, `flag_rate[label=1]`, MDE 0.15, ceiling $15; + units `pending` -- model-free candidate generation on the 41 development repositories is not built, and plane Run + annotation checks / the `experiment-outcome` task are not written) and `exp-002-retrieval-ensemble` (arm A today's + k = 5 temperature sampling, arm B the retrieval ensemble; primary canary agreement (target 0.90) and within-pair + AUC, secondary usd per candidate, AUC, ADVISORY recall/precision; families injection and access_control on the + slice-2 candidates; ceiling $4: A $1, B $3). + - [ ] 7.3 exp-001 run and recorded (~$7.7 expected, ceiling $15) -- an inconclusive result is pre-registered as + likely; needs the frozen units file first. + - [ ] 7.4 exp-002 run and recorded **(budget go-ahead: ~$2 expected, ceiling $4)**; further program variants + (`k = 5` vs `k = 1`, the confidence source) as new manifests. - _Requirements: 3.6, 4.1-4.4_ - [ ] 8. Integration diff --git a/.kiro/steering/overview.md b/.kiro/steering/overview.md index 02d0fc2..a430965 100644 --- a/.kiro/steering/overview.md +++ b/.kiro/steering/overview.md @@ -215,5 +215,5 @@ Counted 2026-10-02 from each spec's `tasks.md` checkboxes and `spec.json`. > and evolutionary proposals come from the plane memory. The HarnessX specs in the table above are history. As of 2026-09-09: 97 source modules, 90 load-bearing and 7 standalone capabilities; 814 tests. -150 source modules, 145 load-bearing and 5 standalone capabilities, 0 orphaned -(`benchmarks/measurements/2026-09-08-module-audit.json`, regenerated 2026-10-02); 1,964 tests collected. +151 source modules, 146 load-bearing and 5 standalone capabilities, 0 orphaned +(`benchmarks/measurements/2026-09-08-module-audit.json`, regenerated 2026-10-02); 2,009 tests collected. diff --git a/benchmarks/measurements/2026-09-08-module-audit.json b/benchmarks/measurements/2026-09-08-module-audit.json index 182a817..56cfd90 100644 --- a/benchmarks/measurements/2026-09-08-module-audit.json +++ b/benchmarks/measurements/2026-09-08-module-audit.json @@ -1,7 +1,7 @@ { - "modules": 150, + "modules": 151, "counts": { - "load_bearing": 145, + "load_bearing": 146, "standalone_capability": 5 }, "orphaned": [], @@ -355,6 +355,7 @@ "learn.evaluate", "learn.examples", "learn.excerpt", + "learn.experiments", "learn.features", "learn.folds", "learn.labels", @@ -390,7 +391,8 @@ "reason": "reachable from gate, map_gate, pair_gate, cli", "importers": [ "learn.cli", - "learn.evaluate" + "learn.evaluate", + "learn.experiments" ] }, { @@ -408,7 +410,8 @@ "reason": "reachable from gate, map_gate, pair_gate, cli", "importers": [ "learn.cli", - "learn.compile" + "learn.compile", + "learn.experiments" ] }, { @@ -419,6 +422,7 @@ "learn.cli", "learn.compile", "learn.evaluate", + "learn.experiments", "learn.folds", "learn.program", "learn.retrieve" @@ -435,6 +439,14 @@ "plane.harvest" ] }, + { + "module": "learn.experiments", + "classification": "load_bearing", + "reason": "reachable from gate, map_gate, pair_gate, cli", + "importers": [ + "learn.cli" + ] + }, { "module": "learn.features", "classification": "load_bearing", @@ -452,6 +464,7 @@ "learn.cli", "learn.compile", "learn.evaluate", + "learn.experiments", "learn.program", "learn.retrieve" ] @@ -487,7 +500,8 @@ "importers": [ "learn.cli", "learn.compile", - "learn.evaluate" + "learn.evaluate", + "learn.experiments" ] }, { @@ -831,6 +845,7 @@ "learn.compile", "learn.embeddings", "learn.examples", + "learn.experiments", "learn.program", "plane.egress", "plane.engine_alerts", @@ -855,6 +870,7 @@ "learn.cli", "learn.compile", "learn.embeddings", + "learn.experiments", "learn.program", "plane.generate", "plane.tasks.roles", @@ -920,7 +936,7 @@ { "module": "plane.memory", "classification": "load_bearing", - "reason": "harnessx-removal Req 6.1/6.4: the memory store (file, MinIO), ingest and fact reuse; `ousast plane remember`", + "reason": "harnessx-removal Req 6.1/6.4: the memory store (file, S3), ingest and fact reuse; `ousast plane remember`", "importers": [ "cli", "improve.memory", @@ -928,6 +944,7 @@ "learn.compile", "learn.embeddings", "learn.examples", + "learn.experiments", "learn.program", "plane.generate", "plane.tasks.loop", diff --git a/docs/decision-engine.md b/docs/decision-engine.md index 401986a..43135d1 100644 --- a/docs/decision-engine.md +++ b/docs/decision-engine.md @@ -79,6 +79,7 @@ v3 check are still ahead. | `ousast learn evaluate --program ID` | outer-fold evaluation, cross-fitted calibration, nested operating points | | `ousast learn curve --program ID` | the metric over memory size on a fixed candidate subset | | `ousast learn audit-leaks` | how well each signal alone separates the sides of a pair (no model) | +| `ousast learn experiment register\|units\|run\|analyse MANIFEST` | A/B experiments from a pre-registered manifest (`plane/experiments/.yaml`): record its digest, freeze the units (no model), run the arms paired (arm A may replay from the cache at $0), analyse (repository bootstrap, exact McNemar, two looks, the adoption verdict) | The memory store is the plane's (`OUSAST_MEMORY`; see [ops/ax/README.md](ops/ax/README.md)). The harvest that produced the plane signals runs as plane Runs generated by diff --git a/docs/memory.md b/docs/memory.md index c1a6b77..692c59e 100644 --- a/docs/memory.md +++ b/docs/memory.md @@ -30,7 +30,10 @@ replaces itself. The kinds (`KINDS`) and who writes them: | `features` | `remember` task, from the `features` task | one feature record per candidate and family | | `proposal_outcome` | `ousast improve --memory` (not with `--dry-run`), and the loop's `improve` step | what happened to each memory proposal in a round | | `example` | `ousast learn memory build` | a labelled case: label + feature record + excerpt sha (no rationale) | -| `label`, `decision`, `experiment`, `arm_outcome`, `experiment_result` | reserved for later decision-engine tasks | declared in `KINDS`; not written by the code on main yet | +| `experiment` | `ousast learn experiment register` (task `register`) and `run` (task `run`) | a registered manifest's digests and provenance; a run's counts-only summary. Repo `experiments/`, pin the registering commit | +| `arm_outcome` | `ousast learn experiment run` | one decision per (unit, arm, replicate): score, verdict, usd, order | +| `experiment_result` | `ousast learn experiment analyse --record` | the paired estimates, the looks and the adoption verdict (counts and intervals only) | +| `label`, `decision` | reserved for later decision-engine tasks | declared in `KINDS`; not written by the code on main yet | **Blobs** (`put_blob`) are content-addressed: a 64-hex name under a fixed prefix. Excerpts are written by `learn memory build` (`learn/examples.py`), embeddings by `learn memory embed` diff --git a/plane/experiments/exp-001-known-callers.yaml b/plane/experiments/exp-001-known-callers.yaml new file mode 100644 index 0000000..3f65a3e --- /dev/null +++ b/plane/experiments/exp-001-known-callers.yaml @@ -0,0 +1,34 @@ +# Experiment exp-001 (learned-decision-engine design section 6, Requirement 4.4): the known-callers prompt line. +# Registered before it runs; `ousast learn experiment register` records this file's sha256, and a changed file is a new +# experiment. Pre-registered expectation: with about 55 positive units the paired MDE is about 0.15, the observed gap +# 0.10 needs about 115 -- an inconclusive result is the likely and honest outcome; a follow-up needs new units. +id: exp-001-known-callers +hypothesis: > + Removing the known-callers line from the verify prompt changes the rate at which verify flags a true site in any + pass (plane-increment-2 recorded 17 of 20 declared sites flagged with the line vs 19 without, confounded). +kind: plane # arms are plane Run configurations; the plane's experiment-outcome task writes the arm_outcome rows +arms: + A: + label: verify as today (known-callers line on; CALLERS_CAP 8) + task_env: {} + B: + label: verify with the known-callers line off + task_env: {OUSAST_VERIFY_CALLERS: "off"} # everything else identical: model, PER_HUNT, MAX_STEPS, facts, runner image +unit: candidate +pairing: paired # every unit runs in both arms +cluster: repository # the resampling unit of the bootstrap +order: {randomise: per_case, seed: 1} +metric: + primary: "flag_rate[label=1]" # flagged in any pass, over the positive units (the 17-vs-19 metric) + secondary: ["flag_rate[label=0]", "usd_per_candidate"] # false flags on fixed-pin negatives; cost +test: {primary: cluster_bootstrap_paired_diff, secondary: exact_mcnemar, alpha: 0.05, sides: 2, resamples: 10000} +mde: 0.15 +stopping: {looks: [0.5, 1.0], boundary: obrien_fleming} # |z| >= 2.797 at half the repositories, 1.977 at the end +budget_usd: {total: 15, per_arm: 7.5} # ~$7.7 expected: 552 hunts x $0.014 +families: [injection, access_control, path, untrusted_destination] # the verify families +# The development corpus, not v2 (v2 generated the hypothesis): population v1 (11 cases), pmpro, MediaWiki, WP Statistics, +# and the 27 full-repository pointer pairs of the verify families -- 41 repositories. Candidates are generated model-free +# (fix-changed functions plus the top engine- and quick-alert functions, up to 8 per pin, on both pins) and frozen here; +# only units whose arm-A callers line is non-empty are treated. The file does not exist yet: registration records the +# units as `pending`, and the experiment cannot run until it is frozen, committed and registered again. +units_file: plane/experiments/exp-001-known-callers.units.jsonl diff --git a/plane/experiments/exp-002-retrieval-ensemble.yaml b/plane/experiments/exp-002-retrieval-ensemble.yaml new file mode 100644 index 0000000..7f0664e --- /dev/null +++ b/plane/experiments/exp-002-retrieval-ensemble.yaml @@ -0,0 +1,41 @@ +# Experiment exp-002 (learned-decision-engine design section 6, engine-level experiments): variance reduction. +# Registered before it runs; `ousast learn experiment register` records this file's sha256, and a changed file is a new +# experiment. The incumbents are the slice-2 compiled programs (benchmarks/measurements/2026-10-02-decision-engine-slice-2), +# whose canary agreement on a re-run was 0.83 (injection, 30 candidates) and 0.83 (access_control). +id: exp-002-retrieval-ensemble +hypothesis: > + Drawing the k = 5 samples from five disjoint retrieved example sets at temperature 0 (a retrieval ensemble) makes the + decision more repeatable than five temperature samples over one example set, without hurting the ranking within a + candidate pair, at no higher cost per candidate. +kind: program +programs: # the incumbents: one compiled program per family (slice 2, record.json `program.id`) + injection: b207a0710e0abf053b3b9397b25f9ee1cab4679144865e3a82c1b9be2c7b4982 + access_control: 745cb1d76bbb09dc1574d333d2118489be62e39886517894d3699fa565b0e8c2 +arms: + A: + label: today's sampling -- k = 5, sample 1 at temperature 0 and samples 2..5 at 0.7, one retrieved example set + spec: {} # the compiled program as it is + replay_only: false # replicate 0 replays the evaluation's cached responses at $0; replicate 1 (the canary) is paid + repeats: 2 + budget_usd: 1.0 + B: + label: retrieval ensemble -- k = 5 at temperature 0, each sample over a disjoint retrieved example set + spec: {sampling: retrieval_ensemble} # the one declared difference (ProgramSpec.sampling) + repeats: 2 + budget_usd: 3.0 +unit: candidate +pairing: paired # every candidate in both arms, same memory, same outer folds +cluster: repository +order: {randomise: per_case, seed: 2} +metric: + primary: + - {name: canary_agreement, target: 0.90} # replicate 0 and 1 give the same verdict; B must reach 0.90 + - within_pair_auc # the vulnerable side of a candidate pair scores above the fixed side + secondary: [usd_per_candidate, auc, advisory_recall, advisory_precision] +test: {primary: cluster_bootstrap_paired_diff, secondary: exact_mcnemar, alpha: 0.05, sides: 2, resamples: 10000} +mde: 0.10 # agreement 0.83 -> 0.93 on ~147 paired candidates +stopping: {looks: [0.5, 1.0], boundary: obrien_fleming} +budget_usd: {total: 4, per_arm: {A: 1.0, B: 3.0}} # ~$0.5-0.7 per arm-pass of 147 candidates at k = 5; 3 paid passes +families: [injection, access_control] +folds: {kind: outer, k: 5, seed: 0, compile_fraction: 0.25} # the incumbent evaluation's folds (learn/compile.toml seed) +units_file: plane/experiments/exp-002-retrieval-ensemble.units.jsonl # frozen by `ousast learn experiment units` diff --git a/src/openultrasast/learn/cli.py b/src/openultrasast/learn/cli.py index f8eadfd..a2f1eb2 100644 --- a/src/openultrasast/learn/cli.py +++ b/src/openultrasast/learn/cli.py @@ -54,6 +54,30 @@ def add_commands(learn_sub: Any) -> None: if name == "curve": sub.add_argument("--subset", type=int, default=200) + experiment = learn_sub.add_parser("experiment", help="A/B experiments: register a manifest, run its arms paired, analyse") + experiment_sub = experiment.add_subparsers(dest="experiment_command", required=True) + for name, text in ( + ("register", "record the manifest's and the units file's sha256 as the experiment row (both committed at HEAD)"), + ("units", "freeze the units of a program experiment from the memory (no model): the families' evaluation candidates"), + ("run", "run every arm on every unit, paired; arm A replays from the response cache when the manifest says so"), + ("analyse", "paired cluster bootstrap, exact McNemar, the two looks and the adoption verdict (counts and intervals only)"), + ): + sub = experiment_sub.add_parser(name, help=text) + sub.add_argument("manifest", type=Path, help="plane/experiments/.yaml") + sub.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)") + if name in ("units", "run"): + sub.add_argument("--profile", choices=("static", "plane"), default="static") + if name == "run": + sub.add_argument( + "--manifest-model", type=Path, default=Path("plane/models/deepseek-flash.yaml"), help="the chat Model (prices)" + ) + sub.add_argument("--no-early-stop", action="store_true", help="take no first look during the run") + if name in ("run", "analyse"): + sub.add_argument("--resamples", type=int, help="bootstrap resamples (default: the manifest's)") + sub.add_argument("--out", type=Path, help="write the counts-only report here (analyse --record: benchmarks/experiments//)") + if name == "analyse": + sub.add_argument("--record", action="store_true", help="also write the experiment_result row") + leaks = learn_sub.add_parser("audit-leaks", help="how well each signal alone separates the sides of a pair (no model)") leaks.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)") leaks.add_argument("--units", type=Path, help="the harvest units.json (unit -> case and side) for feature rows without a side") @@ -73,10 +97,99 @@ def run(args: argparse.Namespace) -> int: return _memory_embed(args) if args.learn_command in ("compile", "evaluate", "curve"): return _program_command(args) + if args.learn_command == "experiment": + return _experiment(args) print(f"learn: unknown command {args.learn_command}", file=sys.stderr) return 2 +def _experiment(args: argparse.Namespace) -> int: + """register | units | run | analyse. Exit 2 on a refused manifest (unregistered, modified, uncommitted), 3 when a + ceiling stops the run (``unfinished``: a rerun resumes from the outcome rows).""" + from datetime import date + + from ..plane.budget import MeteredClient + from ..plane.memory import open_store + from . import experiments as ex + from .program import Caller + + store = open_store(args.memory) + try: + manifest = ex.load_manifest(args.manifest) + if args.experiment_command == "register": + row = ex.register(store, manifest, created=date.today().isoformat()) + shown = ("experiment", "manifest_digest", "units_status", "units_digest", "pin", "arms") + print(json.dumps({k: row[k] for k in shown}, indent=2)) + return 0 + if args.experiment_command == "analyse": + row = ex.check_registered(store, manifest, need_units=False) + provenance = {"code_commit": ex.committed_at_head(manifest.path)} + report = ex.analyse(manifest, ex.load_outcomes(store, manifest.id), resamples=args.resamples, provenance=provenance) + if args.record: + store.put_row(ex.result_row(manifest, str(row["pin"]), report)) + _write(report, args.out or (ex.RESULT_DIR / manifest.id / "result.json" if args.record else None)) + return 0 + from .embeddings import EmbeddingCache + from .examples import load_examples + + memory = load_examples(store, args.profile) + if args.experiment_command == "units": + units = ex.freeze_units(store, manifest, memory) + digest = ex.write_units(manifest.units_file, units) + paired = sum(1 for u in units if u.pair) + print(json.dumps({"units": len(units), "paired": paired, "groups": len({u.group for u in units}), "sha256": digest, + "path": str(manifest.units_file), "next": "commit the file, then register again"}, indent=2)) # fmt: skip + return 0 + cache = EmbeddingCache(store) + vectors = {sha: v for sha in cache.names() if (v := cache.get(sha)) is not None} or None + parameters = _model_parameters(args.manifest_model) + callers: dict[str, Caller] = {} + for name, arm in manifest.arms.items(): + if arm.replay_only: + callers[name] = Caller(None, _arm_model(manifest, arm, store), parameters, store=store) + continue + from ..config import load_config + from ..model.endpoint import resolve_chat_endpoint + + resolved = resolve_chat_endpoint(load_config(None)) + if resolved is None: + raise SystemExit("learn experiment run: no chat endpoint configured (DEEPSEEK_API_KEY); nothing was spent") + client = MeteredClient(resolved[0], prices=parameters, budget_usd=arm.budget_usd) + callers[name] = Caller(client, _arm_model(manifest, arm, store), parameters, store=store) + + def excerpt_text(sha: str) -> str | None: + data = store.get_blob("excerpts", sha) + return data.decode("utf-8") if data is not None else None + + early = not args.no_early_stop + summary = ex.run(store, manifest, memory, callers, excerpt_text, vectors, stop_early=early, resamples=args.resamples) + except ex.ExperimentError as exc: + print(f"learn experiment {args.experiment_command}: {exc}", file=sys.stderr) + return 2 + _write(summary.as_dict(), args.out) + if summary.status == "unfinished": + print(f"learn experiment run: unfinished -- {summary.reason}", file=sys.stderr) + return 3 + return 0 + + +def _arm_model(manifest: Any, arm: Any, store: Any) -> str: + from .compile import load_program + + if "model" in arm.spec: + return str(arm.spec["model"]) + pid = next(iter({**manifest.programs, **arm.programs}.values()), None) + return str(load_program(store, pid)["classify"]["model"]) if pid else "deepseek-flash" + + +def _write(report: dict[str, Any], out: Path | None) -> None: + text = json.dumps(report, indent=2, sort_keys=True, default=str) + "\n" + if out is not None: + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(text, encoding="utf-8") + print(text, end="") + + def _caller(args: argparse.Namespace, store: Any, model: str) -> Any: """The metered, cached caller; ``None`` client under --replay-only. Exits 2 without a ceiling or a key.""" from ..plane.budget import MeteredClient diff --git a/src/openultrasast/learn/experiments.py b/src/openultrasast/learn/experiments.py new file mode 100644 index 0000000..ec206a9 --- /dev/null +++ b/src/openultrasast/learn/experiments.py @@ -0,0 +1,898 @@ +"""A/B experiments over the decision engine and the plane (learned-decision-engine design section 6, Requirement 4). + +Nothing is adopted on one run's number. An experiment is a **manifest** (``plane/experiments/.yaml``) declared +before it runs: hypothesis, two or more **arms** (program variants: a compiled program id plus ``ProgramSpec`` +overrides; or plane Run configurations as ``task_env``), the unit of randomisation, pairing, the resampling cluster, +the metrics, the minimum detectable effect, the sequential rule, the budget ceiling in USD, the families and the +folds, and a frozen **units file**. :func:`register` records the sha256 of the manifest and of the units file as an +``experiment`` row; :func:`run` and :func:`analyse` refuse a manifest that is not registered or whose digest differs +from the registered one, so a question cannot be changed after looking at an answer. + +**Run.** Every unit runs in every arm (paired) in a seeded order of repositories; a seeded per-unit coin decides which +arm goes first, so provider drift over a run never aligns with arms. An arm is a :class:`.program.Program` over the +same memory and folds as the incumbent's evaluation; arm A replays the incumbent's cached responses at $0 when the +manifest says ``replay_only``, and a paid arm runs under its own ceiling. An arm with ``repeats: 2`` decides every +unit twice under different response-cache salts (the canary: do two runs agree?). Each decision is one +``arm_outcome`` row (unit, arm, replicate, score, verdict, usd); a stopped run resumes from the rows it has. + +**Analyse.** Per-unit paired outcomes of arms A and B; the estimate is the paired difference B - A of a metric +(AUC, within-pair AUC, ADVISORY recall and precision, canary agreement, flag rate, usd per candidate) with a 95% CI +from a repository-cluster bootstrap (resampling repositories with the units inside them, since units of one +repository are correlated), an exact McNemar test on the discordant units beside it for binary outcomes, two looks +(half and all of the repositories, in the pre-shuffled order) with O'Brien-Fleming boundaries (|z| >= 2.797, then +1.977), and the adoption rule: **adopt B** when the final CI of the first primary excludes 0 in B's favour and no +other primary's CI excludes 0 against B; **reject B** when any primary's CI excludes 0 against B; for a cost-only +change (``test.cost_only``) **equivalent** when the TOST 90% CI of the detection difference lies within the margin +and the cost CI shows a saving; otherwise **inconclusive**, and the incumbent stays. The result is an +``experiment_result`` row and ``benchmarks/experiments//result.json`` (counts and intervals only). +""" + +from __future__ import annotations + +import hashlib +import json +import math +import random +import subprocess +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import asdict, dataclass, field, fields, replace +from pathlib import Path +from typing import Any + +import yaml + +from ..plane.budget import BudgetExhausted +from ..plane.memory import MemoryStore, row_id +from .compile import load_program, spec_of +from .evaluate import candidate_from_example +from .examples import EXAMPLE_KIND, Example +from .folds import Fold, compile_split, folds_digest, outer_folds +from .program import Caller, Decision, Program, ProgramSpec + +MANIFEST_DIR = Path("plane/experiments") +RESULT_DIR = Path("benchmarks/experiments") +EXPERIMENT_KIND = "experiment" +OUTCOME_KIND = "arm_outcome" +RESULT_KIND = "experiment_result" +REGISTER_TASK = "register" +POPULATION = "experiments" # the envelope of every experiment row; `split` is the experiment id +RUN_TASK = "run" +ARM_KINDS = ("program", "plane") +UNITS = ("candidate", "file", "repository") +PAIRINGS = ("paired", "assigned") +# O'Brien-Fleming boundaries for two equally spaced looks at alpha 0.05, two-sided (design section 6). +OBF_TWO_LOOKS: Mapping[float, float] = {0.5: 2.797, 1.0: 1.977} +RESAMPLES = 10_000 +ALPHA = 0.05 +TOST_MARGIN = 0.05 +ADVISORY_VERDICT = "vulnerable" +# metric name -> (direction: +1 higher is better, -1 lower is better; binary per unit: McNemar applies) +METRICS: Mapping[str, tuple[int, bool]] = { + "auc": (1, False), + "within_pair_auc": (1, False), + "advisory_recall": (1, True), + "advisory_precision": (1, False), + "canary_agreement": (1, True), + "flag_rate": (1, True), + "usd_per_candidate": (-1, False), +} +VERDICTS = ("adopt", "reject", "equivalent", "inconclusive") + + +class ExperimentError(RuntimeError): + """A manifest that cannot be run as declared: malformed, unregistered, modified, or not a program experiment.""" + + +# --- the manifest ------------------------------------------------------------------------------------------------------ + + +@dataclass(frozen=True) +class Arm: + name: str + label: str = "" + programs: Mapping[str, str] = field(default_factory=dict) # family -> compiled program id (program arms) + spec: Mapping[str, Any] = field(default_factory=dict) # ProgramSpec overrides, the declared difference + task_env: Mapping[str, str] = field(default_factory=dict) # plane arms: the Run's task environment + replay_only: bool = False + repeats: int = 1 + budget_usd: float = 0.0 + + +@dataclass(frozen=True) +class Metric: + name: str + label: int | None = None # restrict to units of this label (``flag_rate[label=1]``) + target: float | None = None # a level B must reach, reported beside the test + + @classmethod + def parse(cls, text: str | Mapping[str, Any]) -> Metric: + if isinstance(text, Mapping): + return cls(str(text["name"]), text.get("label"), text.get("target")) + name, _, rest = str(text).partition("[") + label = None + if rest: + key, _, value = rest.rstrip("]").partition("=") + if key != "label": + raise ExperimentError(f"metric {text!r}: only a [label=0|1] selector is understood") + label = int(value) + return cls(name, label) + + @property + def direction(self) -> int: + return METRICS[self.name][0] + + @property + def binary(self) -> bool: + return METRICS[self.name][1] + + def as_dict(self) -> dict[str, Any]: + return {k: v for k, v in asdict(self).items() if v is not None} + + +@dataclass(frozen=True) +class Manifest: + id: str + hypothesis: str + kind: str + arms: Mapping[str, Arm] + unit: str + pairing: str + cluster: str + seed: int + primary: tuple[Metric, ...] + secondary: tuple[Metric, ...] + alpha: float + resamples: int + cost_only: bool + tost_margin: float + mde: float + looks: tuple[float, ...] + boundaries: tuple[float, ...] + budget_total_usd: float + families: tuple[str, ...] + folds: Mapping[str, Any] + units_file: Path + programs: Mapping[str, str] + path: Path + digest: str + raw: Mapping[str, Any] + + @property + def paid_arms(self) -> list[str]: + return [a for a, arm in self.arms.items() if not arm.replay_only] + + +def _digest(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _boundaries(looks: Sequence[float], boundary: str, explicit: Sequence[float] | None) -> tuple[float, ...]: + if explicit is not None: + if len(explicit) != len(looks): + raise ExperimentError("stopping.boundaries must give one z per look") + return tuple(float(z) for z in explicit) + if boundary != "obrien_fleming": + raise ExperimentError(f"stopping.boundary {boundary!r}: only obrien_fleming, or explicit boundaries") + if tuple(looks) != (0.5, 1.0): + raise ExperimentError("obrien_fleming boundaries are tabulated for looks [0.5, 1.0] only; give stopping.boundaries") + return tuple(OBF_TWO_LOOKS[look] for look in looks) + + +def load_manifest(path: Path) -> Manifest: + """Parse and validate a manifest; its ``digest`` is the sha256 of the file bytes (register and run compare it).""" + data = path.read_bytes() + raw = yaml.safe_load(data) + if not isinstance(raw, Mapping): + raise ExperimentError(f"{path}: a manifest is a mapping") + try: + kind = str(raw.get("kind", "program")) + if kind not in ARM_KINDS: + raise ExperimentError(f"{path}: kind {kind!r} (one of {', '.join(ARM_KINDS)})") + arms = { + str(name): Arm( + str(name), + str(spec.get("label", "")), + {str(k): str(v) for k, v in (spec.get("programs") or {}).items()}, + dict(spec.get("spec") or {}), + {str(k): str(v) for k, v in (spec.get("task_env") or {}).items()}, + bool(spec.get("replay_only", False)), + int(spec.get("repeats", 1)), + float(spec.get("budget_usd", 0.0)), + ) # fmt: skip + for name, spec in dict(raw["arms"]).items() + } + if len(arms) < 2 or "A" not in arms or "B" not in arms: + raise ExperimentError(f"{path}: arms A and B are required (A is the incumbent)") + unknown = {k for arm in arms.values() for k in arm.spec} - {f.name for f in fields(ProgramSpec)} + if unknown: + raise ExperimentError(f"{path}: arm spec overrides {sorted(unknown)} are not ProgramSpec fields") + metric = dict(raw["metric"]) + primary_raw = metric["primary"] + primary = tuple(Metric.parse(m) for m in (primary_raw if isinstance(primary_raw, list) else [primary_raw])) + secondary = tuple(Metric.parse(m) for m in metric.get("secondary", [])) + for m in (*primary, *secondary): + if m.name not in METRICS: + raise ExperimentError(f"{path}: metric {m.name!r} (one of {', '.join(METRICS)})") + test = dict(raw.get("test") or {}) + stopping = dict(raw.get("stopping") or {"looks": [1.0], "boundaries": [1.959964]}) + looks = tuple(float(x) for x in stopping.get("looks", [0.5, 1.0])) + if looks[-1] != 1.0 or any(b <= a for a, b in zip(looks, looks[1:], strict=False)) or looks[0] <= 0: + raise ExperimentError(f"{path}: stopping.looks must increase to 1.0") + budget = raw["budget_usd"] + total = float(budget["total"] if isinstance(budget, Mapping) else budget) + arm_total = sum(a.budget_usd for a in arms.values()) + if arm_total > total + 1e-9: + raise ExperimentError(f"{path}: the arms' ceilings ({arm_total}) exceed budget_usd.total ({total})") + unit = str(raw.get("unit", "candidate")) + pairing = str(raw.get("pairing", "paired")) + if unit not in UNITS or pairing not in PAIRINGS: + raise ExperimentError(f"{path}: unit in {UNITS} and pairing in {PAIRINGS}") + if "units_file" not in raw: + raise ExperimentError(f"{path}: units_file is required (the frozen list of units)") + manifest = Manifest( + id=str(raw["id"]), hypothesis=str(raw["hypothesis"]), kind=kind, arms=arms, unit=unit, pairing=pairing, + cluster=str(raw.get("cluster", "repository")), seed=int((raw.get("order") or {}).get("seed", 0)), + primary=primary, secondary=secondary, alpha=float(test.get("alpha", ALPHA)), resamples=int(test.get("resamples", RESAMPLES)), + cost_only=bool(test.get("cost_only", False)), tost_margin=float(test.get("tost_margin", TOST_MARGIN)), + mde=float(raw.get("mde", 0.0)), looks=looks, + boundaries=_boundaries(looks, str(stopping.get("boundary", "obrien_fleming")), stopping.get("boundaries")), + budget_total_usd=total, families=tuple(str(f) for f in raw.get("families", [])), folds=dict(raw.get("folds") or {}), + units_file=Path(str(raw["units_file"])), programs={str(k): str(v) for k, v in (raw.get("programs") or {}).items()}, + path=path, digest=_digest(data), raw=raw, + ) # fmt: skip + except KeyError as exc: + raise ExperimentError(f"{path}: missing {exc}") from exc + if manifest.id != path.stem: + raise ExperimentError(f"{path}: id {manifest.id!r} must equal the file stem") + if kind == "program" and not manifest.families: + raise ExperimentError(f"{path}: a program experiment names its families") + return manifest + + +# --- register ------------------------------------------------------------------------------------------------------------ + + +def committed_at_head(path: Path, cwd: Path | None = None) -> str | None: + """HEAD's commit when ``path`` is tracked and unchanged against HEAD, else ``None``.""" + try: + tracked = subprocess.run(["git", "ls-files", "--error-unmatch", str(path)], cwd=cwd, capture_output=True, text=True) + changed = subprocess.run(["git", "diff", "--quiet", "HEAD", "--", str(path)], cwd=cwd, capture_output=True, text=True) + head = subprocess.run(["git", "rev-parse", "HEAD"], cwd=cwd, capture_output=True, text=True) + except OSError: + return None + if tracked.returncode or changed.returncode or head.returncode: + return None + return head.stdout.strip() or None + + +def experiment_repo(experiment_id: str) -> str: + return f"experiments/{experiment_id}" + + +def registration_id(experiment_id: str) -> str: + return row_id(EXPERIMENT_KIND, f"experiment-{experiment_id}", REGISTER_TASK, experiment_id) + + +def registered(store: MemoryStore, experiment_id: str) -> dict[str, Any] | None: + rows = [r.row for r in store.rows(repo=experiment_repo(experiment_id), kind=EXPERIMENT_KIND) if r.row.get("task") == REGISTER_TASK] + return rows[0] if rows else None + + +def register( + store: MemoryStore, + manifest: Manifest, + *, + commit_of: Callable[[Path], str | None] | None = None, + created: str = "", +) -> dict[str, Any]: + """Record the manifest's and the units file's digests as the ``experiment`` row. Refuses a manifest that is not + committed at HEAD, a units file that exists but is not committed, and a second registration under a different + digest (a changed question is a new experiment id). A units file that does not exist yet is recorded as + ``pending``: the experiment cannot run until it is frozen and registered.""" + check = commit_of or committed_at_head + commit = check(manifest.path) + if commit is None: + raise ExperimentError(f"{manifest.path} is not committed at HEAD: commit the manifest before registering it") + units_digest = "" + if manifest.units_file.exists(): + if check(manifest.units_file) is None: + raise ExperimentError(f"{manifest.units_file} is not committed at HEAD: commit the frozen units before registering") + units_digest = _digest(manifest.units_file.read_bytes()) + previous = registered(store, manifest.id) + if previous is not None and previous.get("manifest_digest") != manifest.digest: + raise ExperimentError( + f"{manifest.id} is registered under digest {str(previous.get('manifest_digest'))[:12]}; the file now hashes to " + f"{manifest.digest[:12]}. A changed manifest is a new experiment: give it a new id" + ) + if previous is not None and previous.get("units_digest") and previous.get("units_digest") != units_digest: + raise ExperimentError(f"{manifest.id}: the units file changed after registration; a new units list is a new experiment") + row = { + "id": registration_id(manifest.id), "kind": EXPERIMENT_KIND, "repo": experiment_repo(manifest.id), "pin": commit, + "run": f"experiment-{manifest.id}", "task": REGISTER_TASK, "population": POPULATION, "split": manifest.id, "image": "host", + "experiment": manifest.id, "hypothesis": manifest.hypothesis, "manifest_path": str(manifest.path), + "manifest_digest": manifest.digest, "units_file": str(manifest.units_file), "units_digest": units_digest, + "units_status": "frozen" if units_digest else "pending", "arms": sorted(manifest.arms), "experiment_kind": manifest.kind, + "primary": [m.as_dict() for m in manifest.primary], "secondary": [m.as_dict() for m in manifest.secondary], + "mde": manifest.mde, "looks": list(manifest.looks), "boundaries": list(manifest.boundaries), + "budget_total_usd": manifest.budget_total_usd, "families": list(manifest.families), "created": created, + "status": "registered", + } # fmt: skip + store.put_row(row) + return row + + +def check_registered(store: MemoryStore, manifest: Manifest, *, need_units: bool = True) -> dict[str, Any]: + """The registration row, or a refusal: unregistered, modified manifest, missing or modified units.""" + row = registered(store, manifest.id) + if row is None: + raise ExperimentError(f"{manifest.id} is not registered: `ousast learn experiment register {manifest.path}` first") + if row.get("manifest_digest") != manifest.digest: + raise ExperimentError(f"{manifest.path} differs from the registered manifest ({str(row.get('manifest_digest'))[:12]}); refusing") + if need_units: + if not manifest.units_file.exists(): + raise ExperimentError(f"{manifest.units_file} does not exist: freeze the units, commit, register again") + if not row.get("units_digest"): + raise ExperimentError(f"{manifest.id}: the units file was not registered (pending); register again with it committed") + if _digest(manifest.units_file.read_bytes()) != row["units_digest"]: + raise ExperimentError(f"{manifest.units_file} differs from the registered units; refusing") + return row + + +# --- units --------------------------------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Unit: + unit: str # the example id (program experiments) or the plane's unit key + group: str # the repository: the resampling cluster + family: str + label: int + fold: str = "" + pair: str = "" # the two sides of one candidate (vulnerable and fixed pin) share it; "" when unpaired + + def as_dict(self) -> dict[str, Any]: + return asdict(self) + + +def read_units(path: Path) -> list[Unit]: + out = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + d = json.loads(line) + out.append( + Unit(str(d["unit"]), str(d["group"]), str(d["family"]), int(d["label"]), str(d.get("fold", "")), str(d.get("pair", ""))) + ) + return out + + +def write_units(path: Path, units: Sequence[Unit]) -> str: + path.parent.mkdir(parents=True, exist_ok=True) + text = "".join(json.dumps(u.as_dict(), sort_keys=True) + "\n" for u in units) + path.write_text(text, encoding="utf-8") + return _digest(text.encode("utf-8")) + + +def program_folds(memory: Sequence[Example], manifest: Manifest) -> tuple[list[Fold], str]: + """The incumbent evaluation's outer folds (compile split excluded) and their digest, from the manifest's ``folds``.""" + seed = int(manifest.folds.get("seed", 0)) + split = compile_split(memory, fraction=float(manifest.folds.get("compile_fraction", 0.25)), seed=seed) + folds = outer_folds(memory, k=int(manifest.folds.get("k", 5)), seed=seed, exclude=split.groups) + return folds, folds_digest(split, folds) + + +def freeze_units(store: MemoryStore, manifest: Manifest, memory: Sequence[Example]) -> list[Unit]: + """The evaluation candidates of the manifest's families under the incumbent's outer folds, paired by candidate + across the two pins where both sides are in memory (no model; the list is then committed and registered).""" + if manifest.kind != "program": + raise ExperimentError(f"{manifest.id}: units of a plane experiment come from candidate generation, not the memory") + folds, _ = program_folds(memory, manifest) + sides: dict[str, tuple[str, str, str]] = {} + for record in store.rows(kind=EXAMPLE_KIND): + row = record.row + sides[str(row["id"])] = (str(row.get("repo", "")), str(row.get("candidate", "")), str(row.get("family", ""))) + chosen = [(e, f) for f in folds if f.status == "ok" for e in memory if e.group in f.eval_groups and e.family in manifest.families] + by_side: dict[tuple[str, str, str], set[int]] = {} + for e, _ in chosen: + by_side.setdefault(sides.get(e.id, ("", e.id, e.family)), set()).add(e.label) + units = [] + for e, f in chosen: + side = sides.get(e.id, ("", e.id, e.family)) + pair = _digest("\x00".join(side).encode())[:16] if by_side[side] == {0, 1} and side[1] else "" + units.append(Unit(e.id, e.group, e.family, e.label, f.name, pair)) + return sorted(units, key=lambda u: (u.group, u.family, u.unit)) + + +# --- outcomes ------------------------------------------------------------------------------------------------------------ + + +@dataclass(frozen=True) +class Outcome: + """One decision: a unit under an arm, replicate ``r`` (0 is the run; 1 the canary re-run).""" + + unit: str + group: str + family: str + label: int + pair: str + arm: str + replicate: int + s: float | None + verdict: str + usd: float | None = None + parse_failed: int = 0 + order: int = 0 # which position the arm had for this unit (0 ran first) + + @property + def flag(self) -> bool: + return self.verdict == ADVISORY_VERDICT + + +def outcome_id(experiment_id: str, arm: str, unit: str, replicate: int) -> str: + return row_id(OUTCOME_KIND, f"experiment-{experiment_id}", f"arm-{arm}", f"{unit}:{replicate}") + + +def outcome_row(experiment_id: str, pin: str, outcome: Outcome) -> dict[str, Any]: + return { + "id": outcome_id(experiment_id, outcome.arm, outcome.unit, outcome.replicate), "kind": OUTCOME_KIND, + "repo": experiment_repo(experiment_id), "pin": pin, "run": f"experiment-{experiment_id}", "task": f"arm-{outcome.arm}", + "population": POPULATION, "split": experiment_id, "image": "host", "experiment": experiment_id, **asdict(outcome), + } # fmt: skip + + +def outcome_from_row(row: Mapping[str, Any]) -> Outcome: + return Outcome( + str(row["unit"]), str(row["group"]), str(row["family"]), int(row["label"]), str(row.get("pair") or ""), str(row["arm"]), + int(row.get("replicate", 0)), None if row.get("s") is None else float(row["s"]), str(row.get("verdict") or ""), + None if row.get("usd") is None else float(row["usd"]), int(row.get("parse_failed", 0)), int(row.get("order", 0)), + ) # fmt: skip + + +def load_outcomes(store: MemoryStore, experiment_id: str) -> list[Outcome]: + return [outcome_from_row(r.row) for r in store.rows(repo=experiment_repo(experiment_id), kind=OUTCOME_KIND)] + + +@dataclass(frozen=True) +class ArmUnit: + """A unit under one arm with its replicates folded in: the run's score and verdict, the mean cost, and whether + the canary replicate agreed with the run.""" + + unit: str + group: str + family: str + label: int + pair: str + s: float | None + verdict: str + usd: float | None + agree: bool | None + + @property + def flag(self) -> bool: + return self.verdict == ADVISORY_VERDICT + + +def arm_units(outcomes: Iterable[Outcome], arm: str) -> dict[str, ArmUnit]: + by_unit: dict[str, list[Outcome]] = {} + for o in outcomes: + if o.arm == arm: + by_unit.setdefault(o.unit, []).append(o) + out = {} + for unit, rows in by_unit.items(): + rows.sort(key=lambda o: o.replicate) + first = rows[0] + costs = [o.usd for o in rows if o.usd is not None] + agree = None if len(rows) < 2 else rows[0].verdict == rows[1].verdict + out[unit] = ArmUnit(unit, first.group, first.family, first.label, first.pair, first.s, first.verdict, + sum(costs) / len(costs) if costs else None, agree) # fmt: skip + return out + + +def paired_units(outcomes: Iterable[Outcome], a: str = "A", b: str = "B") -> list[tuple[ArmUnit, ArmUnit]]: + """Units decided in both arms (replicate 0), in unit order; a unit one arm skipped is not a pair.""" + rows = list(outcomes) + left, right = arm_units(rows, a), arm_units(rows, b) + return [(left[u], right[u]) for u in sorted(left) if u in right] + + +# --- metrics ------------------------------------------------------------------------------------------------------------- + + +def _mean(values: Sequence[float]) -> float | None: + return sum(values) / len(values) if values else None + + +def auc(rows: Sequence[ArmUnit]) -> float | None: + """Mann-Whitney AUC of the score against the label; ties count one half; ``None`` without both labels.""" + pos = sorted(r.s for r in rows if r.label == 1 and r.s is not None) + neg = sorted(r.s for r in rows if r.label == 0 and r.s is not None) + if not pos or not neg: + return None + wins = 0.0 + for p in pos: + for n in neg: + wins += 1.0 if p > n else 0.5 if p == n else 0.0 + return wins / (len(pos) * len(neg)) + + +def within_pair_auc(rows: Sequence[ArmUnit]) -> float | None: + """Of the candidate pairs (vulnerable and fixed side both present), the share where the vulnerable side scores + higher (ties one half).""" + pairs: dict[str, dict[int, float]] = {} + for r in rows: + if r.pair and r.s is not None: + pairs.setdefault(r.pair, {})[r.label] = r.s + done = [(p[1], p[0]) for p in pairs.values() if 0 in p and 1 in p] + if not done: + return None + return sum(1.0 if v > f else 0.5 if v == f else 0.0 for v, f in done) / len(done) + + +def metric_value(metric: Metric, rows: Sequence[ArmUnit]) -> float | None: + if metric.label is not None: + rows = [r for r in rows if r.label == metric.label] + name = metric.name + if name == "auc": + return auc(rows) + if name == "within_pair_auc": + return within_pair_auc(rows) + if name == "advisory_recall": + return _mean([float(r.flag) for r in rows if r.label == 1]) + if name == "advisory_precision": + return _mean([float(r.label == 1) for r in rows if r.flag]) + if name == "canary_agreement": + return _mean([float(r.agree) for r in rows if r.agree is not None]) + if name == "flag_rate": + return _mean([float(r.flag) for r in rows]) + if name == "usd_per_candidate": + return _mean([r.usd for r in rows if r.usd is not None]) + raise ExperimentError(f"unknown metric {name!r}") + + +def binary_value(metric: Metric, row: ArmUnit) -> bool | None: + """The per-unit binary outcome behind a binary metric (McNemar's cell), or ``None`` when it does not apply.""" + if metric.label is not None and row.label != metric.label: + return None + if metric.name == "advisory_recall": + return row.flag if row.label == 1 else None + if metric.name == "canary_agreement": + return row.agree + if metric.name == "flag_rate": + return row.flag + return None + + +# --- statistics ---------------------------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Estimate: + metric: str + a: float | None + b: float | None + diff: float | None + ci: tuple[float, float] | None # percentile bootstrap, two-sided at alpha + se: float | None + units: int + groups: int + resamples: int + + def as_dict(self) -> dict[str, Any]: + return asdict(self) + + +def paired_bootstrap( + pairs: Sequence[tuple[ArmUnit, ArmUnit]], metric: Metric, *, resamples: int = RESAMPLES, seed: int = 0, alpha: float = ALPHA +) -> Estimate: + """The paired difference B - A with a percentile CI from a cluster bootstrap over repositories: each resample draws + repositories with replacement and takes every paired unit inside them, in both arms.""" + by_group: dict[str, list[tuple[ArmUnit, ArmUnit]]] = {} + for pair in pairs: + by_group.setdefault(pair[0].group, []).append(pair) + names = sorted(by_group) + a_all = metric_value(metric, [p[0] for p in pairs]) + b_all = metric_value(metric, [p[1] for p in pairs]) + diff = None if a_all is None or b_all is None else b_all - a_all + rng = random.Random(seed) + values: list[float] = [] + for _ in range(resamples if names else 0): + sample = [p for g in (rng.choice(names) for _ in names) for p in by_group[g]] + a = metric_value(metric, [p[0] for p in sample]) + b = metric_value(metric, [p[1] for p in sample]) + if a is not None and b is not None: + values.append(b - a) + ci = se = None + if len(values) >= 2: + values.sort() + ci = (values[int(alpha / 2 * (len(values) - 1))], values[int((1 - alpha / 2) * (len(values) - 1))]) + mean = sum(values) / len(values) + se = math.sqrt(sum((v - mean) ** 2 for v in values) / (len(values) - 1)) + return Estimate(metric.name, a_all, b_all, diff, ci, se, len(pairs), len(names), len(values)) + + +@dataclass(frozen=True) +class McNemar: + """Exact McNemar: of the discordant pairs, ``b`` where only A has the outcome and ``c`` where only B has it; + the two-sided p-value of a binomial test with p = 0.5.""" + + b: int + c: int + concordant: int + p_value: float + + def as_dict(self) -> dict[str, Any]: + return asdict(self) + + +def exact_mcnemar(table: Sequence[tuple[bool, bool]]) -> McNemar: + b = sum(1 for x, y in table if x and not y) + c = sum(1 for x, y in table if y and not x) + n = b + c + if n == 0: + return McNemar(b, c, len(table), 1.0) + k = min(b, c) + tail = sum(math.comb(n, i) for i in range(k + 1)) / 2**n + return McNemar(b, c, len(table) - n, min(1.0, 2 * tail)) + + +def mcnemar_for(pairs: Sequence[tuple[ArmUnit, ArmUnit]], metric: Metric) -> McNemar | None: + if not metric.binary: + return None + table = [] + for a, b in pairs: + x, y = binary_value(metric, a), binary_value(metric, b) + if x is not None and y is not None: + table.append((x, y)) + return exact_mcnemar(table) if table else None + + +def look_order(groups: Iterable[str], seed: int) -> list[str]: + """The pre-shuffled repository order the looks (and a run) follow.""" + names = sorted(set(groups)) + random.Random(seed).shuffle(names) + return names + + +@dataclass(frozen=True) +class Look: + fraction: float + groups: int + boundary: float + estimate: Estimate + z: float | None + crossed: bool + + def as_dict(self) -> dict[str, Any]: + d = asdict(self) + d["estimate"] = self.estimate.as_dict() + return d + + +def sequential_looks( + pairs: Sequence[tuple[ArmUnit, ArmUnit]], metric: Metric, manifest: Manifest, *, resamples: int | None = None +) -> list[Look]: + """The estimate at each look (the first ceil(fraction * groups) repositories in the pre-shuffled order) against its + O'Brien-Fleming boundary. A crossed boundary stops the sequence: later looks are not taken.""" + order = look_order((p[0].group for p in pairs), manifest.seed) + out: list[Look] = [] + n = resamples if resamples is not None else manifest.resamples + for fraction, boundary in zip(manifest.looks, manifest.boundaries, strict=True): + taken = set(order[: math.ceil(fraction * len(order))]) + subset = [p for p in pairs if p[0].group in taken] + est = paired_bootstrap(subset, metric, resamples=n, seed=manifest.seed, alpha=manifest.alpha) + z = None if est.diff is None or not est.se else est.diff / est.se + # no bootstrap variance with a non-zero difference (every repository moved the same way) crosses any boundary + crossed = (z is not None and abs(z) >= boundary) or (est.se == 0 and bool(est.diff)) + out.append(Look(fraction, len(taken), boundary, est, z, crossed)) + if crossed: + break + return out + + +def excludes_zero(ci: tuple[float, float] | None, direction: int) -> str: + """``favour`` when the CI lies on B's good side, ``against`` on the bad side, else ``spans``.""" + if ci is None: + return "spans" + lo, hi = ci if direction > 0 else (-ci[1], -ci[0]) + return "favour" if lo > 0 else "against" if hi < 0 else "spans" + + +def decide(manifest: Manifest, finals: Mapping[str, Estimate], cost: Estimate | None, tost: Estimate | None) -> tuple[str, str]: + """The adoption rule (module docstring): verdict and the reason in one line.""" + sides = {m.name: excludes_zero(finals[m.name].ci, m.direction) for m in manifest.primary} + if any(side == "against" for side in sides.values()): + bad = [n for n, side in sides.items() if side == "against"] + return "reject", f"the CI of {', '.join(bad)} excludes 0 against B" + first = manifest.primary[0].name + if sides[first] == "favour": + return "adopt", f"the CI of {first} excludes 0 in B's favour and no primary CI excludes 0 against B" + if manifest.cost_only and cost is not None and tost is not None and tost.ci is not None and cost.ci is not None: + within = -manifest.tost_margin <= tost.ci[0] and tost.ci[1] <= manifest.tost_margin + saving = cost.ci[1] < 0 + if within and saving: + return "equivalent", f"the TOST 90% CI of {tost.metric} lies within +-{manifest.tost_margin} and the cost CI shows a saving" + return "inconclusive", f"the CI of {first} spans 0; the incumbent stays" + + +def analyse( + manifest: Manifest, outcomes: Sequence[Outcome], *, resamples: int | None = None, provenance: Mapping[str, Any] | None = None +) -> dict[str, Any]: + """Counts and intervals only: per-metric final estimates, McNemar where the outcome is binary, the looks of the + first primary, the adoption verdict.""" + pairs = paired_units(outcomes) + n = resamples if resamples is not None else manifest.resamples + finals: dict[str, Estimate] = {} + tests: dict[str, Any] = {} + for metric in (*manifest.primary, *manifest.secondary): + est = paired_bootstrap(pairs, metric, resamples=n, seed=manifest.seed, alpha=manifest.alpha) + finals[metric.name] = est + mc = mcnemar_for(pairs, metric) + tests[metric.name] = { + "estimate": est.as_dict(), "mcnemar": mc.as_dict() if mc else None, "direction": metric.direction, + "ci_vs_zero": excludes_zero(est.ci, metric.direction), + "target": None if metric.target is None else {"level": metric.target, "b_meets": est.b is not None and est.b >= metric.target}, + } # fmt: skip + looks = sequential_looks(pairs, manifest.primary[0], manifest, resamples=n) if pairs else [] + cost = finals.get("usd_per_candidate") + tost = None + if manifest.cost_only: + detection = next((m for m in (*manifest.primary, *manifest.secondary) if m.name != "usd_per_candidate"), None) + if detection is not None: + tost = paired_bootstrap(pairs, detection, resamples=n, seed=manifest.seed, alpha=0.10) + verdict, reason = decide(manifest, finals, cost, tost) if pairs else ("inconclusive", "no paired units") + arms = sorted({o.arm for o in outcomes}) + spend = {arm: round(sum(o.usd or 0.0 for o in outcomes if o.arm == arm), 6) for arm in arms} + return { + "experiment": manifest.id, "manifest_digest": manifest.digest, "hypothesis": manifest.hypothesis, + "arms": {a: {"label": arm.label, "replay_only": arm.replay_only, "repeats": arm.repeats, "budget_usd": arm.budget_usd} + for a, arm in manifest.arms.items()}, + "units": {"paired": len(pairs), "groups": len({p[0].group for p in pairs}), "outcomes": len(outcomes), + "by_label": {str(lab): sum(p[0].label == lab for p in pairs) for lab in (0, 1)}}, + "primary": [m.as_dict() for m in manifest.primary], "secondary": [m.as_dict() for m in manifest.secondary], + "tests": tests, "looks": [look.as_dict() for look in looks], + "tost": None if tost is None else tost.as_dict(), "mde": manifest.mde, "alpha": manifest.alpha, "resamples": n, + "decision": verdict, "reason": reason, "spend_usd": spend, **dict(provenance or {}), + } # fmt: skip + + +def result_row(manifest: Manifest, pin: str, result: Mapping[str, Any]) -> dict[str, Any]: + tests = result["tests"] + return { + "id": row_id(RESULT_KIND, f"experiment-{manifest.id}", "analyse", manifest.id), "kind": RESULT_KIND, + "repo": experiment_repo(manifest.id), "pin": pin, "run": f"experiment-{manifest.id}", "task": "analyse", "population": POPULATION, + "split": manifest.id, "image": "host", "experiment": manifest.id, "manifest_digest": manifest.digest, + "decision": result["decision"], "reason": result["reason"], "units": result["units"], "spend_usd": result["spend_usd"], + "tests": {n: {"diff": t["estimate"]["diff"], "ci": t["estimate"]["ci"], "ci_vs_zero": t["ci_vs_zero"]} for n, t in tests.items()}, + "looks": [{"fraction": look["fraction"], "z": look["z"], "crossed": look["crossed"]} for look in result["looks"]], + } # fmt: skip + + +# --- run ----------------------------------------------------------------------------------------------------------------- + + +def arm_spec(manifest: Manifest, arm: Arm, family: str, store: MemoryStore) -> ProgramSpec: + """The arm's program: the family's compiled incumbent, with the arm's declared overrides and nothing else.""" + pid = arm.programs.get(family) or manifest.programs.get(family) + if pid is None: + raise ExperimentError(f"{manifest.id}: arm {arm.name} names no program for family {family!r}") + return replace(spec_of(load_program(store, pid)), **dict(arm.spec)) + + +@dataclass +class RunSummary: + experiment: str + status: str # done | unfinished | stopped + units: int + decided: dict[str, int] + skipped: dict[str, int] + spend: dict[str, dict[str, Any]] + folds_digest: str + stopped_at_look: int | None = None + reason: str = "" + + def as_dict(self) -> dict[str, Any]: + return asdict(self) + + +def run( + store: MemoryStore, + manifest: Manifest, + memory: Sequence[Example], + callers: Mapping[str, Caller], + excerpt_text: Callable[[str], str | None], + vectors: Mapping[str, Sequence[float]] | None = None, + *, + stop_early: bool = True, + resamples: int | None = None, + on_decision: Callable[[Outcome, Decision], None] | None = None, +) -> RunSummary: + """Every unit in every arm, paired, repositories in the pre-shuffled look order, a seeded coin per unit for which + arm goes first; outcomes already in the store are not decided again. After the first look's repositories the + primary is tested against its boundary and, when crossed, the run stops. A ceiling ends the run ``unfinished``.""" + if manifest.kind != "program": + raise ExperimentError(f"{manifest.id}: plane arms run as plane Runs; their outcomes arrive as arm_outcome rows") + row = check_registered(store, manifest) + units = read_units(manifest.units_file) + folds, fdigest = program_folds(memory, manifest) + by_name = {f.name: f for f in folds} + index = {e.id: e for e in memory} + for unit in units: + fold = by_name.get(unit.fold) + if fold is None or unit.group not in fold.eval_groups or unit.unit not in index: + raise ExperimentError( + f"{manifest.id}: unit {unit.unit[:12]} is not under fold {unit.fold!r} of this memory; the memory changed" + ) + programs: dict[tuple[str, str], Program] = {} + for arm in manifest.arms.values(): + for family in manifest.families: + programs[arm.name, family] = Program(arm_spec(manifest, arm, family, store), memory, excerpt_text, vectors) + candidate_of = candidate_from_example(excerpt_text, vectors) + have = {outcome_id(manifest.id, o.arm, o.unit, o.replicate) for o in load_outcomes(store, manifest.id)} + order = look_order((u.group for u in units), manifest.seed) + first_look = set(order[: math.ceil(manifest.looks[0] * len(order))]) if len(manifest.looks) > 1 else set() + rank = {g: i for i, g in enumerate(order)} + decided = {a: 0 for a in manifest.arms} + skipped = {a: 0 for a in manifest.arms} + status, stopped_at, reason = "done", None, "" + rng = random.Random(manifest.seed) + arm_names = sorted(manifest.arms) + fold_seed = int(manifest.folds.get("seed", 0)) + try: + for unit in sorted(units, key=lambda u: (rank[u.group], u.family, u.unit)): + if stop_early and first_look and unit.group not in first_look: # the walk has passed every first-look repository + looks = sequential_looks( + paired_units(load_outcomes(store, manifest.id)), manifest.primary[0], manifest, resamples=resamples + ) + if looks and looks[0].crossed: + status, stopped_at, reason = "stopped", 1, f"the first look crossed its boundary (z = {looks[0].z:.3f})" + break + first_look = set() # tested once + first = rng.random() < 0.5 + ordered = arm_names if first else list(reversed(arm_names)) + for position, name in enumerate(ordered): + arm = manifest.arms[name] + for r in range(arm.repeats): + if outcome_id(manifest.id, name, unit.unit, r) in have: + skipped[name] += 1 + continue + candidate = candidate_of(index[unit.unit]) + if candidate is None: + raise ExperimentError(f"{manifest.id}: unit {unit.unit[:12]} has no excerpt in the store") + caller = callers[name] + caller.salt = f"r{r}:" if r else "" + before = caller.usd() or 0.0 + try: + decision = programs[name, unit.family].decide(candidate, by_name[unit.fold], caller, seed=fold_seed) + finally: + caller.salt = "" + usd = round((caller.usd() or 0.0) - before, 6) if caller.usd() is not None else None + outcome = Outcome(unit.unit, unit.group, unit.family, unit.label, unit.pair, name, r, decision.s, decision.verdict, usd, + decision.parse_failed, position) # fmt: skip + store.put_row(outcome_row(manifest.id, str(row["pin"]), outcome)) + have.add(outcome_id(manifest.id, name, unit.unit, r)) + decided[name] += 1 + if on_decision is not None: + on_decision(outcome, decision) + except BudgetExhausted as exc: + status, reason = "unfinished", str(exc) + spend = { + name: {"usd": c.usd(), "client_calls": c.calls, "replayed": c.replayed, "cache_digest": c.cache_digest()} + for name, c in callers.items() + } + summary = RunSummary(manifest.id, status, len(units), decided, skipped, spend, fdigest, stopped_at, reason) + store.put_row({ + "id": row_id(EXPERIMENT_KIND, f"experiment-{manifest.id}", RUN_TASK, manifest.id), "kind": EXPERIMENT_KIND, + "repo": experiment_repo(manifest.id), "pin": str(row["pin"]), "run": f"experiment-{manifest.id}", "task": RUN_TASK, + "population": POPULATION, "split": manifest.id, "image": "host", "experiment": manifest.id, + "manifest_digest": manifest.digest, **summary.as_dict(), + }) # fmt: skip + return summary + + +__all__ = [ + "ALPHA", "ARM_KINDS", "EXPERIMENT_KIND", "MANIFEST_DIR", "METRICS", "OBF_TWO_LOOKS", "OUTCOME_KIND", "RESAMPLES", "RESULT_DIR", + "RESULT_KIND", "TOST_MARGIN", "VERDICTS", "Arm", "ArmUnit", "Estimate", "ExperimentError", "Look", "Manifest", "McNemar", "Metric", + "Outcome", "RunSummary", "Unit", "analyse", "arm_spec", "arm_units", "auc", "check_registered", "committed_at_head", "decide", + "exact_mcnemar", "excludes_zero", "freeze_units", "load_manifest", "load_outcomes", "look_order", "mcnemar_for", "metric_value", + "outcome_row", "paired_bootstrap", "paired_units", "program_folds", "read_units", "register", "registered", "result_row", "run", + "sequential_looks", "within_pair_auc", "write_units", +] # fmt: skip diff --git a/src/openultrasast/learn/program.py b/src/openultrasast/learn/program.py index f03a9b7..436cc68 100644 --- a/src/openultrasast/learn/program.py +++ b/src/openultrasast/learn/program.py @@ -28,7 +28,7 @@ from __future__ import annotations import hashlib import json import re -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Collection, Mapping, Sequence from dataclasses import asdict, dataclass, field, replace from typing import Any @@ -42,6 +42,11 @@ from .schema import DEFAULT_INPUTS, features_for, instruments_for, validate_x, w VERDICTS = ("vulnerable", "not_vulnerable", "unsure") UNSURE = "unsure" TEMPERATURES = (0.0, 0.7) +# How the k samples differ: by temperature (the default: sample 1 at 0, samples 2..k at 0.7 over one retrieved +# example set), or by the retrieved examples (every sample at temperature 0, each over a disjoint example set -- +# the variance-reduction arm of experiment exp-002; design section 6, engine-level experiments). +SAMPLING = ("temperature", "retrieval_ensemble") +RETRIEVAL_ENSEMBLE = "retrieval_ensemble" MAX_OUTPUT_TOKENS = 300 CHARS_PER_TOKEN = 4.0 STATE_WORDS = {"none": "no coverage", "failed": "failed", "not_applicable": "not applicable"} @@ -199,6 +204,7 @@ class ProgramSpec: priors: str = "off" inputs: str = DEFAULT_INPUTS # the input profile: which instruments the classifier never sees (v1: the engine) max_output_tokens: int = MAX_OUTPUT_TOKENS + sampling: str = SAMPLING[0] # temperature | retrieval_ensemble (an experiment arm's setting; never compiled) @dataclass(frozen=True) @@ -363,6 +369,7 @@ class Caller: replayed: int = 0 keys: list[str] = field(default_factory=list) usage: dict[str, int] = field(default_factory=lambda: {"prompt_tokens": 0, "prompt_cache_hit_tokens": 0, "completion_tokens": 0}) + salt: str = "" # prefixed to every sample label: a replicate of a run that must not replay the first run's responses _memory: dict[str, dict[str, Any]] = field(default_factory=dict) @property @@ -379,7 +386,7 @@ class Caller: return None def ask(self, messages: list[dict[str, object]], *, temperature: float, sample: str, json_object: bool = True) -> str: - key = request_key(self.model, self.params_digest, messages, temperature, sample, json_object) + key = request_key(self.model, self.params_digest, messages, temperature, f"{self.salt}{sample}", json_object) self.keys.append(key) entry = self._cached(key) if entry is not None: @@ -445,16 +452,21 @@ def aggregate(answers: Sequence[Answer]) -> Decision: ) # fmt: skip +def sample_answer(prompt: Prompt, caller: Caller, *, temperature: float, sample: str, families: Sequence[str] = ()) -> Answer: + """One sample; one retry at temperature 0 when the answer does not parse, then ``PARSE_FAILED``.""" + answer = parse_answer(caller.ask(prompt.messages, temperature=temperature, sample=sample), prompt.valid_lines, families) + if answer is None: + retry = caller.ask(prompt.messages, temperature=TEMPERATURES[0], sample=f"{sample}:retry") + answer = parse_answer(retry, prompt.valid_lines, families) or PARSE_FAILED + return answer + + def classify(prompt: Prompt, caller: Caller, k: int = 1, *, families: Sequence[str] = ()) -> Decision: """``k`` samples (the first at temperature 0, the rest at 0.7); one retry at temperature 0 per unparsable answer.""" - answers: list[Answer] = [] - for j in range(k): - temperature = TEMPERATURES[0] if j == 0 else TEMPERATURES[1] - answer = parse_answer(caller.ask(prompt.messages, temperature=temperature, sample=str(j)), prompt.valid_lines, families) - if answer is None: - retry = caller.ask(prompt.messages, temperature=TEMPERATURES[0], sample=f"{j}:retry") - answer = parse_answer(retry, prompt.valid_lines, families) or PARSE_FAILED - answers.append(answer) + answers = [ + sample_answer(prompt, caller, temperature=TEMPERATURES[0] if j == 0 else TEMPERATURES[1], sample=str(j), families=families) + for j in range(k) + ] return aggregate(answers) @@ -471,25 +483,48 @@ class Program: def __post_init__(self) -> None: self.index = {e.id: e for e in self.memory} - def prepare(self, candidate: Candidate, fold: Fold, *, seed: int = 0) -> tuple[Prompt, Retrieval]: + def prepare( + self, candidate: Candidate, fold: Fold, *, seed: int = 0, exclude: Collection[str] = frozenset() + ) -> tuple[Prompt, Retrieval]: target = candidate.target got = retrieve(target, self.memory, fold, vectors=self.vectors, n1=self.spec.n1, k=self.spec.k_ret, lam=self.spec.lam, seed=seed, - priors=self.spec.priors, inputs=self.spec.inputs) # fmt: skip + priors=self.spec.priors, inputs=self.spec.inputs, exclude=exclude) # fmt: skip demos = demonstrations_for([d.as_dict() for d in self.spec.demos], self.index, target, fold, vectors=self.vectors) usable = [Demo.from_dict(d) for d in demos] prompt = render(self.spec, candidate, got.examples, usable, self.index, self.excerpt_text) assert_boundary([*prompt.example_ids, *prompt.demo_ids], self.index, target, fold, vectors=self.vectors) return prompt, got + def prepare_ensemble(self, candidate: Candidate, fold: Fold, *, seed: int = 0) -> list[tuple[Prompt, Retrieval]]: + """``k`` prompts over disjoint retrieved example sets (the same demonstrations and prefix); a later set is + ``short`` when the eligible pool runs out, never padded with an example an earlier set used.""" + out: list[tuple[Prompt, Retrieval]] = [] + used: set[str] = set() + for _ in range(self.spec.k): + prompt, got = self.prepare(candidate, fold, seed=seed, exclude=frozenset(used)) + used.update(prompt.example_ids) + out.append((prompt, got)) + return out + def decide(self, candidate: Candidate, fold: Fold, caller: Caller, *, seed: int = 0) -> Decision: + if self.spec.sampling == RETRIEVAL_ENSEMBLE: + prepared = self.prepare_ensemble(candidate, fold, seed=seed) + answers = [ + sample_answer(prompt, caller, temperature=TEMPERATURES[0], sample=str(j), families=self.families) + for j, (prompt, _) in enumerate(prepared) + ] + got = prepared[0][1] + return replace(aggregate(answers), retrieval=got.mode, neighbours=got.neighbours) + if self.spec.sampling != SAMPLING[0]: + raise ValueError(f"unknown sampling {self.spec.sampling!r} (one of {', '.join(SAMPLING)})") prompt, got = self.prepare(candidate, fold, seed=seed) decision = classify(prompt, caller, self.spec.k, families=self.families) return replace(decision, retrieval=got.mode, neighbours=got.neighbours) __all__ = [ - "ANSWER_FORMAT", "BASELINE_INSTRUCTION", "PARSE_FAILED", "SIGNATURE", "UNSURE", "VERDICTS", "Answer", "Caller", "Candidate", - "Decision", "Demo", + "ANSWER_FORMAT", "BASELINE_INSTRUCTION", "PARSE_FAILED", "RETRIEVAL_ENSEMBLE", "SAMPLING", "SIGNATURE", "UNSURE", "VERDICTS", + "Answer", "Caller", "Candidate", "Decision", "Demo", "Program", "ProgramSpec", "Prompt", "ReplayMiss", "aggregate", "classify", "label_word", "line_numbers", - "params_digest", "parse_answer", "render", "render_prefix", "render_roles", "render_signals", "request_key", + "params_digest", "parse_answer", "render", "render_prefix", "render_roles", "render_signals", "request_key", "sample_answer", ] # fmt: skip diff --git a/src/openultrasast/learn/retrieve.py b/src/openultrasast/learn/retrieve.py index 7693438..ae67417 100644 --- a/src/openultrasast/learn/retrieve.py +++ b/src/openultrasast/learn/retrieve.py @@ -28,7 +28,7 @@ repeats the check on the ids it actually rendered (:func:`assert_boundary`). from __future__ import annotations import hashlib -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Collection, Iterable, Mapping, Sequence from dataclasses import dataclass from typing import Any @@ -192,9 +192,11 @@ def retrieve( seed: int = 0, priors: str = "off", inputs: str = DEFAULT_INPUTS, + exclude: Collection[str] = frozenset(), ) -> Retrieval: - """The ``k`` balanced examples for ``target`` from the eligible part of ``pool`` (module docstring).""" - candidates = [e for e in pool if e.profile == target.profile and eligible(e, target, fold)] + """The ``k`` balanced examples for ``target`` from the eligible part of ``pool`` (module docstring). ``exclude`` + holds example ids a caller has already used: the retrieval-ensemble sampling draws ``k`` disjoint sets with it.""" + candidates = [e for e in pool if e.id not in exclude and e.profile == target.profile and eligible(e, target, fold)] ranked_all = sorted( ((gower(target.x, target.instruments, e.x, e.instruments, target.profile, priors=priors, inputs=inputs), e) for e in candidates), key=lambda pair: (pair[0], _tie(pair[1].id, seed)), diff --git a/src/openultrasast/plane/tasks/verify.py b/src/openultrasast/plane/tasks/verify.py index 67bc5c8..2c126fd 100644 --- a/src/openultrasast/plane/tasks/verify.py +++ b/src/openultrasast/plane/tasks/verify.py @@ -159,9 +159,10 @@ def callers_line(rows: Sequence[Mapping[str, Any]], cap: int = CALLERS_CAP) -> s def hunt_prompt(path: str, group: Sequence[Candidate], family: str, facts: Mapping[str, Any] | None) -> str: lines = [] + callers_on = os.environ.get("OUSAST_VERIFY_CALLERS", "on") != "off" # arm B of exp-001 (design section 6) turns the line off for _, function, line in group: lines.append(f"- function `{function}` around line {line}") - known = callers_line(callers_of(facts, path, function)) + known = callers_line(callers_of(facts, path, function)) if callers_on else "" if known: lines.append(f" {known}") listing = "\n".join(lines) diff --git a/tests/test_learn_experiments.py b/tests/test_learn_experiments.py new file mode 100644 index 0000000..c7d35f2 --- /dev/null +++ b/tests/test_learn_experiments.py @@ -0,0 +1,485 @@ +"""A/B experiments (learned-decision-engine design section 6, task 7): manifest digest refusal, frozen units and +pairing, the paired repository bootstrap, exact McNemar on a known table, the two-look boundary, the adoption rule, +arm A replaying at zero client calls, and the retrieval-ensemble sampling variant. Every client is scripted; stores +are ``FileStore`` under ``tmp_path``.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path + +import pytest +from learn_fixtures import ScriptedChat, corpus, x_record + +from openultrasast.cli import main +from openultrasast.learn import experiments as ex +from openultrasast.learn.compile import canonical, program_id +from openultrasast.learn.program import TEMPERATURES, Caller, Candidate, Program, ProgramSpec +from openultrasast.learn.retrieve import Target +from openultrasast.plane.budget import MeteredClient +from openultrasast.plane.memory import FileStore + +ROOT = Path(__file__).resolve().parents[1] +PRICES = {"cache_hit_per_m": 0.014, "input_per_m": 0.44, "output_per_m": 1.32} +EXAMPLES = corpus(groups=12) +CODE = { + e.excerpt_sha: f" 10 def f{i}(q):\n 11 return {'execute(q)' if e.label else 'escape(q)'}\n" for i, e in enumerate(EXAMPLES) +} +COMMIT = "c" * 40 +PIN = "a" * 40 + + +def excerpt_text(sha: str) -> str | None: + return CODE.get(sha) + + +def artifact(**classify: object) -> dict[str, object]: + return { + "kind": "llm-program", "instruction": "Judge the candidate.", "demos": [], "profile": "static", + "retrieval": {"k": 4, "lam": 0.5, "n1": 40, "priors": "off"}, + "classify": {"model": "deepseek-flash", "k": 2, "inputs": "v1", "max_output_tokens": 300, **classify}, + } # fmt: skip + + +def put_program(store: FileStore, art: dict[str, object]) -> str: + pid = program_id(art) + store.put_blob("programs", canonical(art), name=pid) + return pid + + +def example_rows() -> list[dict[str, object]]: + """Example rows as the memory build writes them (repo, candidate): the two sides of g0-e0/g0-e1 share a candidate.""" + rows = [] + for i, e in enumerate(EXAMPLES): + candidate = "a.py::f" if e.group == "org0/repo0" and i < 2 else f"a.py::f{i}" + rows.append({ + "id": e.id, "kind": "example", "repo": e.group, "pin": PIN, "run": "learn-memory", "task": "build", "population": "pairs", + "split": "pairs", "image": "host", "candidate": candidate, "family": e.family, "label": e.label, "profile": "static", + }) # fmt: skip + return rows + + +MANIFEST = """ +id: {id} +hypothesis: the variant changes the score +kind: program +programs: {{injection: {pid}, path: {pid}}} +arms: + A: {{spec: {{}}, replay_only: {replay}, repeats: {repeats}, budget_usd: 1.0}} + B: {{spec: {{sampling: retrieval_ensemble}}, repeats: {repeats}, budget_usd: 2.0}} +order: {{seed: 3}} +metric: + primary: [{{name: canary_agreement, target: 0.9}}, within_pair_auc] + secondary: [usd_per_candidate, auc, advisory_recall] +test: {{alpha: 0.05, resamples: 200}} +mde: 0.1 +stopping: {{looks: [0.5, 1.0], boundary: obrien_fleming}} +budget_usd: {{total: 3}} +families: [injection, path] +folds: {{k: 5, seed: 0, compile_fraction: 0.25}} +units_file: {units} +""" + + +def write_manifest(tmp_path: Path, pid: str, *, name: str = "exp-900-test", replay: str = "true", repeats: int = 1) -> Path: + path = tmp_path / f"{name}.yaml" + path.write_text(MANIFEST.format(id=name, pid=pid, replay=replay, repeats=repeats, units=tmp_path / f"{name}.units.jsonl")) + return path + + +def committed(_: Path) -> str | None: + return COMMIT + + +def unit(name: str, group: str, label: int, pair: str = "") -> ex.Unit: + return ex.Unit(name, group, "injection", label, "outer-0", pair) + + +def outcome(u: ex.Unit, arm: str, s: float, verdict: str, *, replicate: int = 0, usd: float = 0.001) -> ex.Outcome: + return ex.Outcome(u.unit, u.group, u.family, u.label, u.pair, arm, replicate, s, verdict, usd) + + +def crafted(groups: int = 8, per_group: int = 4, *, b_better: bool, a_period: int = 2) -> list[ex.Outcome]: + """Paired outcomes over ``groups`` repositories: A flags the positives of every ``a_period``-th repository only, + B flags every positive when ``b_better`` (else it is A's copy); neither flags a negative.""" + out = [] + for g in range(groups): + for i in range(per_group): + u = unit(f"u{g}-{i}", f"org/repo{g}", i % 2, pair=f"p{g}-{i // 2}") + a_flag = u.label == 1 and g % a_period == 0 + out.append(outcome(u, "A", 0.8 if a_flag else 0.2, "vulnerable" if a_flag else "not_vulnerable")) + b_flag = bool(u.label) if b_better else a_flag + out.append(outcome(u, "B", 0.9 if b_flag else 0.1, "vulnerable" if b_flag else "not_vulnerable")) + return out + + +# --- manifests and registration ------------------------------------------------------------------------------------------- + + +def test_the_registered_manifests_load_as_declared() -> None: + exp1 = ex.load_manifest(ROOT / "plane/experiments/exp-001-known-callers.yaml") + exp2 = ex.load_manifest(ROOT / "plane/experiments/exp-002-retrieval-ensemble.yaml") + assert exp1.kind == "plane" and exp1.arms["B"].task_env == {"OUSAST_VERIFY_CALLERS": "off"} and exp1.arms["A"].task_env == {} + assert exp1.primary[0].as_dict() == {"name": "flag_rate", "label": 1} and exp1.budget_total_usd == 15 and exp1.mde == 0.15 + assert exp2.kind == "program" and exp2.arms["B"].spec == {"sampling": "retrieval_ensemble"} and exp2.arms["A"].spec == {} + assert [m.name for m in exp2.primary] == ["canary_agreement", "within_pair_auc"] and exp2.primary[0].target == 0.9 + assert exp2.budget_total_usd == 4 and exp2.families == ("injection", "access_control") and exp2.arms["B"].repeats == 2 + assert exp1.boundaries == exp2.boundaries == (2.797, 1.977) and exp2.looks == (0.5, 1.0) + assert set(exp2.programs) == {"injection", "access_control"} + + +def test_manifest_refusals(tmp_path: Path) -> None: + path = write_manifest(tmp_path, "0" * 64) + text = path.read_text() + for bad, why in ( + (text.replace("sampling: retrieval_ensemble", "temperature: 0.1"), "ProgramSpec fields"), + (text.replace("looks: [0.5, 1.0]", "looks: [0.3, 1.0]"), "tabulated"), + (text.replace("budget_usd: {total: 3}", "budget_usd: {total: 2}"), "exceed"), + (text.replace(" B: {spec", " C: {spec"), "arms A and B"), + (text.replace("advisory_recall", "f1"), "metric"), + ): + path.write_text(bad) + with pytest.raises(ex.ExperimentError, match=why): + ex.load_manifest(path) + (tmp_path / "other.yaml").write_text(text) + with pytest.raises(ex.ExperimentError, match="file stem"): + ex.load_manifest(tmp_path / "other.yaml") + + +def test_register_refuses_uncommitted_and_modified_manifests(tmp_path: Path) -> None: + store = FileStore(tmp_path / "memory") + path = write_manifest(tmp_path, "0" * 64) + manifest = ex.load_manifest(path) + with pytest.raises(ex.ExperimentError, match="not committed"): + ex.register(store, manifest, commit_of=lambda _: None) + with pytest.raises(ex.ExperimentError, match="not registered"): + ex.check_registered(store, manifest, need_units=False) + row = ex.register(store, manifest, commit_of=committed, created="2026-10-02") + assert row["pin"] == COMMIT and row["units_status"] == "pending" and row["manifest_digest"] == manifest.digest + assert ex.check_registered(store, manifest, need_units=False)["id"] == row["id"] + with pytest.raises(ex.ExperimentError, match="freeze the units"): + ex.check_registered(store, manifest) # the units are not frozen: it cannot run + ex.register(store, manifest, commit_of=committed) # the same file again is idempotent + path.write_text(path.read_text().replace("mde: 0.1", "mde: 0.2")) + changed = ex.load_manifest(path) + with pytest.raises(ex.ExperimentError, match="new experiment"): + ex.register(store, changed, commit_of=committed) + with pytest.raises(ex.ExperimentError, match="differs from the registered"): + ex.check_registered(store, changed, need_units=False) + assert len(store.rows(kind="experiment")) == 1 + + +def test_frozen_units_cover_the_evaluated_groups_and_pair_the_sides(tmp_path: Path) -> None: + store = FileStore(tmp_path / "memory") + store.put_rows(example_rows()) + manifest = ex.load_manifest(write_manifest(tmp_path, "0" * 64)) + units = ex.freeze_units(store, manifest, EXAMPLES) + folds, _ = ex.program_folds(EXAMPLES, manifest) + evaluated = {g for f in folds for g in f.eval_groups} + assert units and {u.group for u in units} == evaluated and {u.family for u in units} <= {"injection", "path"} + assert all(u.fold.startswith("outer-") for u in units) + paired = [u for u in units if u.pair] + assert {u.unit for u in paired} == {EXAMPLES[0].id, EXAMPLES[1].id} or not paired # the one shared candidate, if evaluated + digest = ex.write_units(manifest.units_file, units) + assert ex.read_units(manifest.units_file) == units + row = ex.register(store, manifest, commit_of=committed) + assert row["units_status"] == "frozen" and row["units_digest"] == digest + manifest.units_file.write_text(manifest.units_file.read_text() + json.dumps(unit("x", "org9/repo9", 1).as_dict()) + "\n") + with pytest.raises(ex.ExperimentError, match="units"): + ex.check_registered(store, manifest) + + +# --- statistics ------------------------------------------------------------------------------------------------------------- + + +def test_metrics_on_crafted_rows() -> None: + rows = [ + ex.ArmUnit("a", "g", "injection", 1, "p1", 0.9, "vulnerable", 0.002, True), + ex.ArmUnit("b", "g", "injection", 0, "p1", 0.4, "not_vulnerable", 0.004, False), + ex.ArmUnit("c", "h", "injection", 1, "p2", 0.5, "unsure", None, None), + ex.ArmUnit("d", "h", "injection", 0, "p2", 0.5, "vulnerable", None, True), + ] + assert ex.metric_value(ex.Metric("auc"), rows) == pytest.approx((1 + 1 + 1 + 0.5) / 4) + assert ex.metric_value(ex.Metric("within_pair_auc"), rows) == pytest.approx(0.75) + assert ex.metric_value(ex.Metric("advisory_recall"), rows) == 0.5 and ex.metric_value(ex.Metric("advisory_precision"), rows) == 0.5 + assert ex.metric_value(ex.Metric("canary_agreement"), rows) == pytest.approx(2 / 3) + assert ex.metric_value(ex.Metric("usd_per_candidate"), rows) == pytest.approx(0.003) + assert ex.metric_value(ex.Metric("flag_rate", label=0), rows) == 0.5 + assert ex.metric_value(ex.Metric("auc"), rows[:1]) is None and ex.metric_value(ex.Metric("within_pair_auc"), rows[:1]) is None + assert ex.Metric.parse("flag_rate[label=1]") == ex.Metric("flag_rate", 1) and ex.Metric("usd_per_candidate").direction == -1 + + +def test_paired_bootstrap_over_repositories() -> None: + better = ex.paired_units(crafted(b_better=True)) + est = ex.paired_bootstrap(better, ex.Metric("advisory_recall"), resamples=300, seed=1) + assert est.a == 0.5 and est.b == 1.0 and est.diff == 0.5 and est.groups == 8 and est.units == 32 + assert est.ci is not None and est.ci[0] > 0 and ex.excludes_zero(est.ci, 1) == "favour" and ex.excludes_zero(est.ci, -1) == "against" + same = ex.paired_units(crafted(b_better=False)) + est = ex.paired_bootstrap(same, ex.Metric("flag_rate"), resamples=100, seed=1) + assert est.diff == 0 and est.ci == (0.0, 0.0) and est.se == 0 and ex.excludes_zero(est.ci, 1) == "spans" + assert ex.paired_bootstrap([], ex.Metric("auc"), resamples=10).ci is None + # the bootstrap resamples repositories, not units: a repository's units move together + only_a = ex.paired_bootstrap(better[:4], ex.Metric("advisory_recall"), resamples=50, seed=0) + assert only_a.groups == 1 and only_a.ci == (only_a.diff, only_a.diff) + + +def test_exact_mcnemar_on_a_known_table() -> None: + table = [(True, False)] * 2 + [(False, True)] * 10 + [(True, True)] * 5 + [(False, False)] * 3 + mc = ex.exact_mcnemar(table) + assert (mc.b, mc.c, mc.concordant) == (2, 10, 8) + assert mc.p_value == pytest.approx(2 * (1 + 12 + 66) / 2**12) # 0.0386 + assert ex.exact_mcnemar([(True, False)] * 3 + [(False, True)] * 3).p_value == 1.0 + assert ex.exact_mcnemar([(True, True)]).p_value == 1.0 + pairs = ex.paired_units(crafted(b_better=True)) + assert ex.mcnemar_for(pairs, ex.Metric("auc")) is None # not a per-unit binary outcome + recall = ex.mcnemar_for(pairs, ex.Metric("advisory_recall")) + assert recall is not None and recall.b == 0 and recall.c == 8 and recall.p_value == pytest.approx(2 / 256) + + +def test_two_looks_with_obrien_fleming_boundaries(tmp_path: Path) -> None: + manifest = ex.load_manifest(write_manifest(tmp_path, "0" * 64)) + recall = ex.Metric("advisory_recall") + strong = ex.sequential_looks(ex.paired_units(crafted(groups=12, b_better=True, a_period=6)), recall, manifest, resamples=200) + assert len(strong) == 1 and strong[0].crossed and strong[0].fraction == 0.5 and strong[0].boundary == 2.797 and strong[0].groups == 6 + assert strong[0].z is None or strong[0].z >= 2.797 # None: every first-look repository moved the same way (no variance) + late = ex.sequential_looks(ex.paired_units(crafted(groups=12, b_better=True, a_period=2)), recall, manifest, resamples=200) + assert [look.crossed for look in late] == [False, True] and late[1].z is not None and late[1].z >= 1.977 + none = ex.sequential_looks(ex.paired_units(crafted(groups=12, b_better=False)), recall, manifest, resamples=50) + assert [look.fraction for look in none] == [0.5, 1.0] and not any(look.crossed for look in none) + assert [look.boundary for look in none] == [2.797, 1.977] and none[1].groups == 12 + assert ex.look_order(["b", "a", "c"], 3) == ex.look_order(["c", "a", "b"], 3) # the pre-shuffled order is seeded + + +def estimate(name: str, lo: float, hi: float) -> ex.Estimate: + return ex.Estimate(name, 0.5, 0.5 + (lo + hi) / 2, (lo + hi) / 2, (lo, hi), 0.1, 10, 5, 100) + + +def test_the_adoption_rule(tmp_path: Path) -> None: + manifest = ex.load_manifest(write_manifest(tmp_path, "0" * 64)) # primaries: canary_agreement, within_pair_auc + favour, against, spans = ( + estimate("canary_agreement", 0.02, 0.2), + estimate("within_pair_auc", -0.3, -0.1), + estimate("within_pair_auc", -0.1, 0.1), + ) + assert ex.decide(manifest, {"canary_agreement": favour, "within_pair_auc": spans}, None, None)[0] == "adopt" + assert ex.decide(manifest, {"canary_agreement": favour, "within_pair_auc": against}, None, None)[0] == "reject" + assert ex.decide(manifest, {"canary_agreement": spans, "within_pair_auc": spans}, None, None)[0] == "inconclusive" + assert ( + ex.decide(manifest, {"canary_agreement": estimate("canary_agreement", -0.2, -0.02), "within_pair_auc": spans}, None, None)[0] + == "reject" + ) + cost_only = ex.Manifest(**{**manifest.__dict__, "cost_only": True}) + saving, no_saving = estimate("usd_per_candidate", -0.002, -0.001), estimate("usd_per_candidate", -0.001, 0.001) + tost_in, tost_out = estimate("canary_agreement", -0.03, 0.02), estimate("canary_agreement", -0.08, 0.02) + finals = {"canary_agreement": spans, "within_pair_auc": spans} + assert ex.decide(cost_only, finals, saving, tost_in)[0] == "equivalent" + assert ex.decide(cost_only, finals, no_saving, tost_in)[0] == "inconclusive" + assert ex.decide(cost_only, finals, saving, tost_out)[0] == "inconclusive" + assert ex.decide(manifest, finals, saving, tost_in)[0] == "inconclusive" # not declared cost-only: no equivalence + + +def test_analyse_reports_counts_and_intervals_only(tmp_path: Path) -> None: + store = FileStore(tmp_path / "memory") + manifest = ex.load_manifest(write_manifest(tmp_path, "0" * 64)) + outcomes = crafted(groups=10, b_better=True) + for o in outcomes: + store.put_row(ex.outcome_row(manifest.id, COMMIT, o)) + loaded = ex.load_outcomes(store, manifest.id) + assert sorted(loaded, key=lambda o: (o.unit, o.arm)) == sorted(outcomes, key=lambda o: (o.unit, o.arm)) + report = ex.analyse(manifest, loaded, resamples=100, provenance={"code_commit": COMMIT}) + assert report["units"] == {"paired": 40, "groups": 10, "outcomes": 80, "by_label": {"0": 20, "1": 20}} + assert report["tests"]["canary_agreement"]["estimate"]["diff"] is None # one replicate: no canary + assert report["tests"]["advisory_recall"]["estimate"]["diff"] == 0.5 and report["tests"]["advisory_recall"]["mcnemar"]["c"] == 10 + assert report["tests"]["within_pair_auc"]["ci_vs_zero"] == "favour" and report["decision"] == "inconclusive" # the first primary spans + assert report["code_commit"] == COMMIT and report["manifest_digest"] == manifest.digest + text = json.dumps(report) + assert "org/repo" not in text and "u3-1" not in text # no identities + store.put_row(ex.result_row(manifest, COMMIT, report)) + assert store.rows(kind="experiment_result")[0].row["decision"] == "inconclusive" + + +# --- run -------------------------------------------------------------------------------------------------------------------- + + +def paid(answer: ScriptedChat | None = None, budget: float = 5.0) -> tuple[ScriptedChat, Caller]: + chat = answer or ScriptedChat() + return chat, Caller(MeteredClient(chat, prices=PRICES, budget_usd=budget), "deepseek-flash", PRICES) + + +def prepared(tmp_path: Path, *, replay: str = "true", repeats: int = 1) -> tuple[FileStore, ex.Manifest, list[ex.Unit]]: + store = FileStore(tmp_path / "memory") + store.put_rows(example_rows()) + pid = put_program(store, artifact()) + manifest = ex.load_manifest(write_manifest(tmp_path, pid, replay=replay, repeats=repeats)) + units = ex.freeze_units(store, manifest, EXAMPLES)[:12] + ex.write_units(manifest.units_file, units) + ex.register(store, manifest, commit_of=committed) + return store, manifest, units + + +def test_arm_a_replays_the_evaluation_at_zero_client_calls(tmp_path: Path) -> None: + store, manifest, units = prepared(tmp_path) + # the incumbent's evaluation: every unit decided once, responses in the cache + chat, caller = paid() + caller.store = store + program = Program(ex.arm_spec(manifest, manifest.arms["A"], "injection", store), EXAMPLES, excerpt_text) + folds = {f.name: f for f in ex.program_folds(EXAMPLES, manifest)[0]} + index = {e.id: e for e in EXAMPLES} + for u in units: + e = index[u.unit] + program.decide(Candidate(Target.of(e), CODE[e.excerpt_sha], "python"), folds[u.fold], caller, seed=0) + evaluation_calls = len(chat.calls) + assert evaluation_calls == 2 * len(units) + b_chat, b_caller = paid() + b_caller.store = store + seen: list[tuple[str, str]] = [] + summary = ex.run(store, manifest, EXAMPLES, {"A": Caller(None, "deepseek-flash", PRICES, store=store), "B": b_caller}, excerpt_text, + on_decision=lambda o, d: seen.append((o.arm, o.verdict)), resamples=20) # fmt: skip + assert summary.status == "done" and summary.decided == {"A": len(units), "B": len(units)} and summary.skipped == {"A": 0, "B": 0} + assert len(chat.calls) == evaluation_calls # arm A made no client call + assert summary.spend["A"]["client_calls"] == 0 and summary.spend["A"]["replayed"] == 2 * len(units) + # B's first sample renders the same prompt as A's temperature-0 sample: the cache answers it; the second set is new + assert summary.spend["B"]["client_calls"] == len(units) and summary.spend["B"]["replayed"] == len(units) + assert len(b_chat.calls) == len(units) + assert all(call["temperature"] == 0.0 for call in b_chat.calls) # the ensemble arm samples at temperature 0 + rows = ex.load_outcomes(store, manifest.id) + assert len(rows) == 2 * len(units) and {o.arm for o in rows} == {"A", "B"} and all(o.usd is not None and o.usd > 0 for o in rows) + assert {o.order for o in rows if o.arm == "A"} == {0, 1} # the seeded coin put A first for some units and second for others + # a rerun decides nothing again and spends nothing + again = ex.run(store, manifest, EXAMPLES, {"A": Caller(None, "deepseek-flash", PRICES, store=store), "B": paid()[1]}, excerpt_text) + assert again.decided == {"A": 0, "B": 0} and again.skipped == {"A": len(units), "B": len(units)} + report = ex.analyse(manifest, ex.load_outcomes(store, manifest.id), resamples=50) + assert report["units"]["paired"] == len(units) and report["decision"] in ex.VERDICTS + assert store.rows(kind="experiment", where=None)[-1].row["task"] in ("register", "run") + + +def test_run_refuses_an_unregistered_or_modified_manifest_and_a_missing_cache(tmp_path: Path) -> None: + store, manifest, _ = prepared(tmp_path) + manifest.path.write_text(manifest.path.read_text().replace("mde: 0.1", "mde: 0.15")) + changed = ex.load_manifest(manifest.path) + callers = {"A": Caller(None, "deepseek-flash", PRICES, store=store), "B": paid()[1]} + with pytest.raises(ex.ExperimentError, match="differs"): + ex.run(store, changed, EXAMPLES, callers, excerpt_text) + with pytest.raises(ex.ExperimentError, match="not registered"): + ex.run(FileStore(tmp_path / "other"), manifest, EXAMPLES, callers, excerpt_text) + from openultrasast.learn.program import ReplayMiss + + with pytest.raises(ReplayMiss): # arm A replay-only without the evaluation's cache: loud, never a silent zero + ex.run(store, manifest, EXAMPLES, callers, excerpt_text) + assert ex.load_outcomes(store, manifest.id) == [] + + +def test_canary_replicates_use_a_fresh_cache_key_and_a_ceiling_leaves_the_run_resumable(tmp_path: Path) -> None: + store, manifest, units = prepared(tmp_path, replay="false", repeats=2) + a_chat, a_caller = paid(budget=0.0005) + a_caller.store = store + b_chat, b_caller = paid() + b_caller.store = store + summary = ex.run(store, manifest, EXAMPLES, {"A": a_caller, "B": b_caller}, excerpt_text, stop_early=False) + assert summary.status == "unfinished" and "ceiling" in summary.reason.lower() or "budget" in summary.reason.lower() + partial = ex.load_outcomes(store, manifest.id) + assert partial and len(partial) < 4 * len(units) + assert a_caller.salt == "" and b_caller.salt == "" + a2_chat, a2_caller = paid() + a2_caller.store = store + done = ex.run(store, manifest, EXAMPLES, {"A": a2_caller, "B": b_caller}, excerpt_text, stop_early=False) + assert done.status == "done" and done.skipped["A"] + done.decided["A"] == 2 * len(units) + rows = ex.load_outcomes(store, manifest.id) + assert len(rows) == 4 * len(units) and {o.replicate for o in rows} == {0, 1} + by_arm = ex.arm_units(rows, "A") + assert all(u.agree is not None for u in by_arm.values()) + report = ex.analyse(manifest, rows, resamples=50) + assert ( + report["tests"]["canary_agreement"]["estimate"]["a"] is not None and report["tests"]["canary_agreement"]["target"]["level"] == 0.9 + ) + # replicate 1 asked the client again (a fresh key), it did not replay replicate 0 + assert len(a_chat.calls) + len(a2_chat.calls) >= 2 * len(units) + 1 + + +def test_plane_experiments_do_not_run_here(tmp_path: Path) -> None: + manifest = ex.load_manifest(ROOT / "plane/experiments/exp-001-known-callers.yaml") + store = FileStore(tmp_path / "memory") + with pytest.raises(ex.ExperimentError, match="plane"): + ex.freeze_units(store, manifest, EXAMPLES) + with pytest.raises(ex.ExperimentError, match="plane"): + ex.run(store, manifest, EXAMPLES, {}, excerpt_text) + + +def test_cli_register_and_analyse(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + pid = "0" * 64 + path = write_manifest(tmp_path, pid) + memory = f"file://{tmp_path / 'memory'}" + monkeypatch.setattr(ex, "committed_at_head", lambda _p, cwd=None: None) + assert main(["learn", "experiment", "register", str(path), "--memory", memory]) == 2 # not committed + monkeypatch.setattr(ex, "committed_at_head", lambda _p, cwd=None: COMMIT) + assert main(["learn", "experiment", "register", str(path), "--memory", memory]) == 0 + assert ( + main(["learn", "experiment", "analyse", str(path), "--memory", memory, "--resamples", "10", "--out", str(tmp_path / "r.json")]) == 0 + ) + report = json.loads((tmp_path / "r.json").read_text()) + assert report["decision"] == "inconclusive" and report["reason"] == "no paired units" + path.write_text(path.read_text().replace("mde: 0.1", "mde: 0.3")) + assert main(["learn", "experiment", "analyse", str(path), "--memory", memory]) == 2 + + +# --- the retrieval-ensemble variant -------------------------------------------------------------------------------------- + + +def candidate() -> Candidate: + rec = x_record(hits=2) + return Candidate( + Target("cand/repo", "injection", "static", rec["x"], rec["instruments"]), + " 40 def view(q):\n 41 return execute(q)\n", + "python", + ) + + +def test_retrieval_ensemble_draws_disjoint_example_sets_and_the_default_is_unchanged() -> None: + from openultrasast.learn.folds import outer_folds + + fold = outer_folds(EXAMPLES, seed=0)[0] + ensemble = Program(ProgramSpec(k=3, k_ret=4, sampling="retrieval_ensemble"), EXAMPLES, excerpt_text) + prepared_sets = ensemble.prepare_ensemble(candidate(), fold) + ids = [set(p.example_ids) for p, _ in prepared_sets] + assert len(ids) == 3 and all(ids) and not (ids[0] & ids[1]) and not (ids[0] & ids[2]) and not (ids[1] & ids[2]) + assert len({p.prefix for p, _ in prepared_sets}) == 1 # the same instruction and demonstrations: the prefix cache still hits + chat, caller = paid() + decision = ensemble.decide(candidate(), fold, caller) + assert len(chat.calls) == 3 and all(c["temperature"] == 0.0 for c in chat.calls) + assert ( + len({json.dumps(c["messages"]) for c in chat.calls}) == 3 and decision.verdict == "vulnerable" and sum(decision.votes.values()) == 3 + ) + # the default: one example set, temperatures 0 then 0.7, byte-identical messages + chat, caller = paid() + default = Program(ProgramSpec(k=3, k_ret=4), EXAMPLES, excerpt_text) + default.decide(candidate(), fold, caller) + assert [c["temperature"] for c in chat.calls] == [TEMPERATURES[0], TEMPERATURES[1], TEMPERATURES[1]] + assert len({json.dumps(c["messages"]) for c in chat.calls}) == 1 + assert ProgramSpec().sampling == "temperature" + with pytest.raises(ValueError, match="sampling"): + Program(ProgramSpec(sampling="dice"), EXAMPLES, excerpt_text).decide(candidate(), fold, paid()[1]) + + +def test_a_salted_caller_never_replays_the_unsalted_response(tmp_path: Path) -> None: + store = FileStore(tmp_path / "memory") + chat, caller = paid(ScriptedChat(lambda m, t: json.dumps({"verdict": "vulnerable", "confidence": 0.9, "cited_lines": [1]}))) + caller.store = store + messages: list[dict[str, object]] = [{"role": "user", "content": " 1 execute(q)"}] + caller.ask(messages, temperature=0.0, sample="0") + caller.ask(messages, temperature=0.0, sample="0") + assert caller.calls == 1 and caller.replayed == 1 + caller.salt = "r1:" + caller.ask(messages, temperature=0.0, sample="0") + assert caller.calls == 2 and len(store.blob_names("responses")) == 2 + assert math.isclose(caller.usd() or 0, caller.usd() or 0) + + +def test_verify_callers_line_switch(monkeypatch: pytest.MonkeyPatch) -> None: + from openultrasast.plane.tasks import verify + + facts = {"callers": {"app.py::run": [{"path": "web.py", "line": 12, "enclosing": "handler"}]}} + group = [("app.py::run", "run", 2)] + assert "Known callers: web.py:12 in handler()" in verify.hunt_prompt("app.py", group, "injection", facts) + monkeypatch.setenv("OUSAST_VERIFY_CALLERS", "off") + assert "Known callers" not in verify.hunt_prompt("app.py", group, "injection", facts) diff --git a/tests/test_plane_increment.py b/tests/test_plane_increment.py index 1393f6e..b667e18 100644 --- a/tests/test_plane_increment.py +++ b/tests/test_plane_increment.py @@ -36,7 +36,7 @@ def _mebibytes(quantity: str) -> float: def test_every_manifest_under_plane_loads() -> None: - paths = sorted(PLANE.rglob("*.yaml")) + paths = sorted(p for p in PLANE.rglob("*.yaml") if "experiments" not in p.parts) # experiment manifests are not ax manifests assert len(paths) >= 36, paths manifests = load_manifests(paths) assert set(manifests.models) == {"deepseek-flash", "openrouter-embedding"} # the second: embeddings only (decision engine) From 3b91a5278853316396a404b7ec54e0bbc574785d Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:08:36 +0200 Subject: [PATCH 2/2] bench(learn): exp-001 and exp-002 registered in the memory store, units pending (decision-engine task 7.2) Registration rows at pin fbc3220 in s3://sast-memory; in-repo records under benchmarks/experiments//registration.json. exp-002's units could not be frozen: every example row in the store carries null plane-only instruments and learn.examples.load_examples raises on them (unchanged code since a90e570) -- a store-side change to diagnose before any paid run. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo --- .kiro/specs/learned-decision-engine/tasks.md | 5 ++++- .../exp-001-known-callers/registration.json | 11 +++++++++++ .../exp-002-retrieval-ensemble/registration.json | 11 +++++++++++ 3 files changed, 26 insertions(+), 1 deletion(-) create mode 100644 benchmarks/experiments/exp-001-known-callers/registration.json create mode 100644 benchmarks/experiments/exp-002-retrieval-ensemble/registration.json diff --git a/.kiro/specs/learned-decision-engine/tasks.md b/.kiro/specs/learned-decision-engine/tasks.md index 941324c..1d756d1 100644 --- a/.kiro/specs/learned-decision-engine/tasks.md +++ b/.kiro/specs/learned-decision-engine/tasks.md @@ -133,7 +133,10 @@ a model or embedding call. annotation checks / the `experiment-outcome` task are not written) and `exp-002-retrieval-ensemble` (arm A today's k = 5 temperature sampling, arm B the retrieval ensemble; primary canary agreement (target 0.90) and within-pair AUC, secondary usd per candidate, AUC, ADVISORY recall/precision; families injection and access_control on the - slice-2 candidates; ceiling $4: A $1, B $3). + slice-2 candidates; ceiling $4: A $1, B $3). Both registered in `s3://sast-memory` at fbc3220 + (`benchmarks/experiments//registration.json`). exp-002's units are `pending` too: on 2026-10-02 every example + row in the store carried null plane-only instruments (`agree`, `model_sinks`, `verify`) and `load_examples` raised + on them, so `ousast learn experiment units` could not freeze the candidates. - [ ] 7.3 exp-001 run and recorded (~$7.7 expected, ceiling $15) -- an inconclusive result is pre-registered as likely; needs the frozen units file first. - [ ] 7.4 exp-002 run and recorded **(budget go-ahead: ~$2 expected, ceiling $4)**; further program variants diff --git a/benchmarks/experiments/exp-001-known-callers/registration.json b/benchmarks/experiments/exp-001-known-callers/registration.json new file mode 100644 index 0000000..444a260 --- /dev/null +++ b/benchmarks/experiments/exp-001-known-callers/registration.json @@ -0,0 +1,11 @@ +{ + "experiment": "exp-001-known-callers", + "registered": "2026-10-02", + "store": "s3://sast-memory", + "manifest": "plane/experiments/exp-001-known-callers.yaml", + "manifest_digest": "76b54aa5ae908a4cd7797f3b4654acf03b90e5ad5d068216036e722104313929", + "pin": "fbc3220bbe1c8a4f8afa0a724fd41a51b390dd22", + "units_status": "pending", + "units_note": "model-free candidate generation on the 41 development repositories is not built; the plane's experiment-outcome task and Run annotation checks are not written", + "status": "registered, not run" +} diff --git a/benchmarks/experiments/exp-002-retrieval-ensemble/registration.json b/benchmarks/experiments/exp-002-retrieval-ensemble/registration.json new file mode 100644 index 0000000..c454261 --- /dev/null +++ b/benchmarks/experiments/exp-002-retrieval-ensemble/registration.json @@ -0,0 +1,11 @@ +{ + "experiment": "exp-002-retrieval-ensemble", + "registered": "2026-10-02", + "store": "s3://sast-memory", + "manifest": "plane/experiments/exp-002-retrieval-ensemble.yaml", + "manifest_digest": "f5004fb8618b1f7cc7dcf298d42b4ac1b6c0c1aa72a97299c77f937bd6175148", + "pin": "fbc3220bbe1c8a4f8afa0a724fd41a51b390dd22", + "units_status": "pending", + "units_note": "`ousast learn experiment units` could not freeze the slice-2 candidates on 2026-10-02: all 904 example rows in the store carry null for the plane-only instruments (agree, model_sinks, verify) and learn.examples.load_examples (unchanged since a90e570) raises on them; freeze, commit and register again once the rows read", + "status": "registered, not run" +}