Checkpoint experiments through SWE pilot v2

scripts/freeze_swe_population_plan.py

@@ -0,0 +1,110 @@
+"""Freeze the all-task ImpossibleBench SWE population schedule; no model calls."""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+import subprocess
+
+from messageboardbench.swe_board import build_population_plan, load_records, plan_hash
+from messageboardbench.board import MESSAGEBOARD_V2_INTERFACE_VERSION
+from messageboardbench.swe_validation import DATASET
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--revision", required=True)
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--out", type=Path, required=True)
+    parser.add_argument("--teams", type=int, default=12)
+    parser.add_argument("--cohorts", type=int, default=3)
+    parser.add_argument("--seed", type=int, default=910)
+    parser.add_argument(
+        "--sample-size", type=int,
+        help="freeze a deterministic pilot subset of this size (default: full split)",
+    )
+    parser.add_argument(
+        "--exclude-plan", type=Path,
+        help="exclude the explicitly selected IDs in an earlier frozen pilot plan",
+    )
+    parser.add_argument(
+        "--messageboard-v2", action="store_true",
+        help="freeze the send_message/read_messages plus organizer-feedback interface",
+    )
+    args = parser.parse_args()
+    if args.messageboard_v2 and args.sample_size is None:
+        parser.error("--messageboard-v2 requires --sample-size")
+    if args.exclude_plan is not None and args.sample_size is None:
+        parser.error("--exclude-plan requires --sample-size")
+    upstream = (ROOT.parent / "impossiblebench").resolve()
+    commit = subprocess.run(
+        ["git", "rev-parse", "HEAD"], cwd=upstream, check=True,
+        capture_output=True, text=True,
+    ).stdout.strip()
+    records = load_records(args.revision, "conflicting")
+    selected = None
+    if args.sample_size is not None:
+        if not 1 <= args.sample_size <= len(records):
+            parser.error("--sample-size must be between 1 and the split size")
+        ranked = sorted(
+            records,
+            key=lambda instance_id: hashlib.sha256(
+                f"swe-pilot-selection-v1:{args.seed}:{instance_id}".encode()
+            ).digest(),
+        )
+        excluded: set[str] = set()
+        excluded_plan_bytes = None
+        if args.exclude_plan is not None:
+            excluded_plan_bytes = args.exclude_plan.read_bytes()
+            earlier = json.loads(excluded_plan_bytes)
+            excluded = set(earlier.get("selection", {}).get("instance_ids", []))
+            if (earlier.get("status") != "frozen"
+                    or earlier.get("plan_sha256") != plan_hash(earlier)
+                    or earlier.get("dataset") != {
+                        "path": DATASET,
+                        "revision": args.revision,
+                        "split": "conflicting",
+                    }
+                    or earlier.get("seed") != args.seed
+                    or not excluded or not excluded <= set(records)):
+                parser.error("--exclude-plan does not contain a valid pinned subset")
+        selected = [instance_id for instance_id in ranked if instance_id not in excluded][
+            :args.sample_size
+        ]
+        if len(selected) != args.sample_size:
+            parser.error("not enough unexcluded records for --sample-size")
+    plan = build_population_plan(
+        records, revision=args.revision, model=args.model,
+        upstream_git_commit=commit, teams=args.teams, cohorts=args.cohorts,
+        seed=args.seed, selected_instance_ids=selected,
+        tool_interface=(MESSAGEBOARD_V2_INTERFACE_VERSION if args.messageboard_v2 else None),
+    )
+    if selected is not None:
+        plan["selection"].update({
+            "ranking_namespace": "swe-pilot-selection-v1",
+            "ranking_seed": args.seed,
+            "excluded_instance_ids": sorted(excluded),
+            "excluded_plan": ({
+                "path": str(args.exclude_plan),
+                "file_sha256": hashlib.sha256(excluded_plan_bytes).hexdigest(),
+                "plan_sha256": earlier["plan_sha256"],
+            } if args.exclude_plan is not None else None),
+        })
+        plan["plan_sha256"] = plan_hash(plan)
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    with args.out.open("x") as handle:
+        json.dump(plan, handle, indent=2)
+        handle.write("\n")
+    print(json.dumps({
+        "plan": str(args.out), "instances": len(records),
+        "episodes": plan["planned_episodes"], "plan_sha256": plan["plan_sha256"],
+    }, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())