File size: 4,769 Bytes
867597e 646ccfc 867597e 646ccfc 867597e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """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()
|