| """Build compressed, per-sample DFG caches for chronological-train R3 scoring.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import gzip |
| import json |
| import signal |
| from collections import Counter |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
| from src.rewards.checker import trace_validity |
| from src.rewards.cache import cache_path |
| from src.rewards.line_graph import build_line_graph |
|
|
| ROW_TIMEOUT_S = 15 |
|
|
|
|
| class _RowTimeout(Exception): |
| pass |
|
|
|
|
| def _alarm_handler(signum, frame): |
| raise _RowTimeout() |
|
|
|
|
| def _build_one(task: tuple[dict, dict, str]) -> dict: |
| row, annotation, root_name = task |
| sample_id = row["sample_id"] |
| signal.signal(signal.SIGALRM, _alarm_handler) |
| signal.alarm(ROW_TIMEOUT_S) |
| try: |
| graph = build_line_graph(row["func_vuln"]) |
| sinks = set() |
| for path in annotation["trace_gt"]: |
| lines = [graph.node_lines[tuple(map(int, node.split("-", 1)))] for node in path] |
| valid, _ = trace_validity(lines, graph, sink=lines[-1]) |
| if not valid: |
| return {"sample_id": sample_id, "status": "trace_mismatch"} |
| sinks.add(lines[-1]) |
| value = {"sample_id": sample_id, "annotation_backend": annotation["annotation_backend"], |
| "sink_lines": sorted(sinks), "graph": graph.to_dict()} |
| path = cache_path(Path(root_name), sample_id) |
| tmp = path.with_suffix(".tmp") |
| with gzip.open(tmp, "wt") as f: |
| json.dump(value, f) |
| tmp.replace(path) |
| return {"sample_id": sample_id, "status": "cached", "n_paths": len(annotation["trace_gt"])} |
| except _RowTimeout: |
| return {"sample_id": sample_id, "status": "timeout"} |
| except Exception as exc: |
| return {"sample_id": sample_id, "status": "error", "error": f"{type(exc).__name__}: {exc}"} |
| finally: |
| signal.alarm(0) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--corpus", default="data/processed/corpus.parquet") |
| parser.add_argument("--split", default="data/splits/chronological/train.json") |
| parser.add_argument("--annotations", default="data/processed/annotations.jsonl") |
| parser.add_argument("--out-dir", default="data/cpg_cache/reward_graphs") |
| parser.add_argument("--report", default="results/phase5/reward_graph_cache.json") |
| parser.add_argument("--limit", type=int, default=None) |
| parser.add_argument("--workers", type=int, default=4) |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
|
|
| train_pairs = set(json.loads(Path(args.split).read_text())) |
| annotations = {a["sample_id"]: a for line in Path(args.annotations).read_text().splitlines() |
| if line.strip() and (a := json.loads(line)).get("trace_gt")} |
| df = pd.read_parquet(args.corpus, columns=["sample_id", "pair_id", "target", "func_vuln"]) |
| rows = [row for row in df.to_dict("records") if row["pair_id"] in train_pairs |
| and int(row["target"]) == 1 and row["sample_id"] in annotations] |
| if args.limit is not None: |
| rows = rows[:args.limit] |
| root = Path(args.out_dir) |
| root.mkdir(parents=True, exist_ok=True) |
| existing = sum(cache_path(root, row["sample_id"]).exists() for row in rows) |
| tasks = [(row, annotations[row["sample_id"]], str(root)) for row in rows |
| if not cache_path(root, row["sample_id"]).exists()] |
| print(f"trace-covered rows for split {args.split}={len(rows)}, existing={existing}, pending={len(tasks)}", flush=True) |
|
|
| counts: Counter = Counter({"already_cached": existing}) |
| failures = [] |
| with Pool(max(1, args.workers)) as pool: |
| for i, result in enumerate(pool.imap_unordered(_build_one, tasks, chunksize=1), 1): |
| counts[result["status"]] += 1 |
| if result["status"] not in ("cached", "already_cached"): |
| failures.append(result) |
| if i % 500 == 0: |
| print(f"processed {i}/{len(tasks)}; cached={counts['cached']}; errors={len(failures)}", flush=True) |
|
|
| report = {"seed": args.seed, "split": args.split, "n_trace_covered_rows": len(rows), |
| "counts": dict(counts), "n_cached_total": counts["already_cached"] + counts["cached"], |
| "error_rate": len(failures) / len(rows) if rows else 0.0, |
| "failures": failures[:50], "cache_dir": str(root)} |
| report_path = Path(args.report) |
| report_path.parent.mkdir(parents=True, exist_ok=True) |
| report_path.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps({k: v for k, v in report.items() if k != "failures"}, indent=2)) |
| if report["error_rate"] > 0.02: |
| raise ValueError("reward graph cache error rate exceeds 2%") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
