| """Phase 2 CLI: runs annotate_sample over the corpus, writing one JSONL |
| record per sample. Resumable (spec §4: "resumable... per-sample artifact |
| files, skip if exists") via a line-count check against an existing output |
| file combined with --resume; CPU-parallelizable via --workers (spec: "target |
| 16 workers") using a process pool, since tree-sitter parsing is CPU-bound and |
| releases the GIL across process boundaries trivially via multiprocessing. |
| |
| Does not invoke Joern (see src/cpg/annotate.py's module docstring) — this is |
| the tree-sitter-backed first pass. A real Joern batch pass + calibration |
| subset (spec §4.1's 200-sample edge-recall comparison) is tracked as |
| follow-up, not run yet. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import signal |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
| from src.cpg.annotate import annotate_sample, annotation_to_dict |
|
|
| ROW_TIMEOUT_S = 15 |
|
|
|
|
| class _RowTimeout(Exception): |
| pass |
|
|
|
|
| def _alarm_handler(signum, frame): |
| raise _RowTimeout() |
|
|
|
|
| def _annotate_row(row: dict) -> dict: |
| """Never lets one pathological sample kill — or hang — the whole pool |
| (spec §1.6: "fail loudly... never silently drop >2% of samples without |
| logging a warning with counts and reasons"). Two failure modes guarded: |
| exceptions (caught directly) and hangs/pathological slowness (bounded by |
| a SIGALRM timeout). The latter was found empirically, not hypothetically: |
| the real corpus contains at least one PrimeVul row whose "function" is a |
| 482KB/23,938-line blob (clearly a misclassified data table, not a real |
| C function) that made build_defuse_graph's O(n^2) edge-building step |
| (src/cpg/defuse.py) take so long it looked indistinguishable from a |
| hang. A size threshold would have been a narrower, more brittle fix — |
| the timeout catches this case AND any future pathological input we |
| haven't seen yet, without needing to characterize what "too large" means |
| in advance.""" |
| signal.signal(signal.SIGALRM, _alarm_handler) |
| signal.alarm(ROW_TIMEOUT_S) |
| try: |
| ann = annotate_sample( |
| sample_id=row["sample_id"], |
| func_vuln=row["func_vuln"], |
| patch_lines_vuln=list(row["patch_lines_vuln"]), |
| cwe=row["cwe"], |
| ) |
| return annotation_to_dict(ann) |
| except _RowTimeout: |
| return { |
| "sample_id": row["sample_id"], |
| "error": f"timeout after {ROW_TIMEOUT_S}s (func_len={len(row['func_vuln'])} chars)", |
| "l1_lines": [], "l2_lines": [], "l3_lines": [], |
| "trace_gt": [], "trace_unverifiable": True, "trace_reason": "annotation_timeout", |
| "annotation_backend": "tree_sitter", |
| } |
| except Exception as e: |
| return { |
| "sample_id": row["sample_id"], |
| "error": f"{type(e).__name__}: {e}", |
| "l1_lines": [], "l2_lines": [], "l3_lines": [], |
| "trace_gt": [], "trace_unverifiable": True, "trace_reason": "annotation_error", |
| "annotation_backend": "tree_sitter", |
| } |
| finally: |
| signal.alarm(0) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--corpus", default="data/processed/corpus.parquet") |
| parser.add_argument("--out", default="data/processed/annotations.jsonl") |
| parser.add_argument("--workers", type=int, default=16) |
| parser.add_argument("--limit", type=int, default=None, help="for smoke testing on a subset") |
| parser.add_argument("--resume", action="store_true", default=False) |
| args = parser.parse_args() |
|
|
| df = pd.read_parquet(args.corpus) |
| if args.limit: |
| df = df.head(args.limit) |
|
|
| out_path = Path(args.out) |
| already_done: set[str] = set() |
| if args.resume and out_path.exists(): |
| with open(out_path) as f: |
| for line in f: |
| line = line.strip() |
| if line: |
| already_done.add(json.loads(line)["sample_id"]) |
| print(f"resuming: {len(already_done)} samples already annotated") |
|
|
| rows = [r for r in df.to_dict("records") if r["sample_id"] not in already_done] |
| print(f"annotating {len(rows)} samples with {args.workers} workers") |
|
|
| out_path.parent.mkdir(parents=True, exist_ok=True) |
| n_errors = 0 |
| n_trace_unverifiable = 0 |
| mode = "a" if args.resume else "w" |
| with open(out_path, mode) as f, Pool(args.workers) as pool: |
| for i, result in enumerate(pool.imap_unordered(_annotate_row, rows, chunksize=10)): |
| f.write(json.dumps(result) + "\n") |
| if result.get("error"): |
| n_errors += 1 |
| elif result["trace_unverifiable"]: |
| n_trace_unverifiable += 1 |
| if (i + 1) % 1000 == 0: |
| f.flush() |
| print(f"{i + 1}/{len(rows)} annotated ({n_errors} errors so far)") |
|
|
| error_rate = n_errors / len(rows) if rows else 0.0 |
| print(f"done. {len(rows)} annotated, {n_trace_unverifiable} trace_unverifiable, " |
| f"{n_errors} errors ({error_rate:.1%})") |
| if error_rate > 0.02: |
| print(f"WARNING: error rate {error_rate:.1%} exceeds the 2% threshold spec §1.6 " |
| f"flags as needing investigation, not silent acceptance.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
