File size: 5,333 Bytes
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """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: # noqa: BLE001 - deliberately broad, see docstring
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()
|