#!/usr/bin/env python from __future__ import annotations import argparse import json import sys from pathlib import Path from typing import Any REQUIRED_SPLITS = ("train", "val", "test") def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description="Audit CIL chart split indexes for leakage.") parser.add_argument("--chart-root", type=Path, default=Path("data/cil_charts")) parser.add_argument("--run-root", type=Path, default=Path("runs")) parser.add_argument("--out-dir", type=Path, default=Path("runs/leakage_audit")) parser.add_argument( "--no-markdown-report", action="store_true", help="Do not write report.md; persistent prose is consolidated in README.md.", ) args = parser.parse_args(argv) indexes = _load_indexes(args.chart_root) violations: list[str] = [] warnings: list[str] = [] for split in REQUIRED_SPLITS: if split not in indexes: violations.append(f"missing index for split {split}: {args.chart_root / split / 'index.json'}") if not violations: _audit_split_flags(indexes, violations) _audit_overlap(indexes, violations) _audit_shards(indexes, violations) _audit_run_hashes(args.run_root, warnings) report = { "status": "pass" if not violations else "fail", "chart_root": str(args.chart_root), "indexes": { split: { "path": str(path), "num_groups_exported": payload.get("num_groups_exported"), "num_rows": payload.get("num_rows"), "include_outcomes": payload.get("include_outcomes"), "retrieval_index_allowed": payload.get("retrieval_index_allowed"), "audience": payload.get("audience"), "content_hash": payload.get("content_hash"), "split_hash": payload.get("split_hash"), } for split, (path, payload) in sorted(indexes.items()) }, "violations": violations, "warnings": warnings, } out_dir = args.out_dir out_dir.mkdir(parents=True, exist_ok=True) (out_dir / "report.json").write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") _write_markdown_report(out_dir, report, no_markdown_report=args.no_markdown_report) print(json.dumps({"status": report["status"], "violations": len(violations)}, indent=2)) return 1 if violations else 0 def _write_markdown_report( out_dir: Path, report: dict[str, Any], *, no_markdown_report: bool, ) -> None: report_path = out_dir / "report.md" if no_markdown_report: report_path.unlink(missing_ok=True) return report_path.write_text(_markdown(report) + "\n") def _load_indexes(chart_root: Path) -> dict[str, tuple[Path, dict[str, Any]]]: output: dict[str, tuple[Path, dict[str, Any]]] = {} for split in REQUIRED_SPLITS: path = chart_root / split / "index.json" if path.exists(): output[split] = (path, json.loads(path.read_text())) return output def _audit_split_flags(indexes: dict[str, tuple[Path, dict[str, Any]]], violations: list[str]) -> None: for split, (path, payload) in indexes.items(): include_outcomes = bool(payload.get("include_outcomes")) retrieval = bool(payload.get("retrieval_index_allowed")) audience = str(payload.get("audience")) if split == "train": if not retrieval: violations.append(f"{path}: train split is not retrieval_index_allowed") if not include_outcomes: violations.append(f"{path}: train split should include outcomes for training") else: if retrieval: violations.append(f"{path}: non-train split is retrieval_index_allowed") if audience != "evaluator_only": violations.append(f"{path}: non-train split audience must be evaluator_only") contract = payload.get("leakage_contract", {}) if not bool(contract.get("deployment_must_not_read_outcomes", False)): violations.append( f"{path}: non-train split must declare deployment_must_not_read_outcomes=true" ) if split != str(payload.get("split")): violations.append(f"{path}: index split field does not match directory") if not payload.get("deployment_candidate_excludes_expert", False): violations.append(f"{path}: missing deployment_candidate_excludes_expert=true") def _audit_overlap(indexes: dict[str, tuple[Path, dict[str, Any]]], violations: list[str]) -> None: for left in REQUIRED_SPLITS: for right in REQUIRED_SPLITS: if left >= right or left not in indexes or right not in indexes: continue left_payload = indexes[left][1] right_payload = indexes[right][1] group_overlap = set(left_payload.get("group_hashes", [])) & set( right_payload.get("group_hashes", []) ) state_overlap = set(left_payload.get("state_hashes", [])) & set( right_payload.get("state_hashes", []) ) if group_overlap: violations.append(f"{left}/{right}: {len(group_overlap)} overlapping chart_ids") if state_overlap: violations.append(f"{left}/{right}: {len(state_overlap)} overlapping state_hashes") def _audit_shards(indexes: dict[str, tuple[Path, dict[str, Any]]], violations: list[str]) -> None: for split, (path, payload) in indexes.items(): for shard in payload.get("shards", []): shard_path = path.parent / str(shard.get("path", "")) if not shard_path.exists(): violations.append(f"{path}: missing shard {shard_path}") def _audit_run_hashes(run_root: Path, warnings: list[str]) -> None: metrics_files = sorted(run_root.glob("*/metrics.json")) if not metrics_files: warnings.append(f"{run_root}: no run metrics.json files found for hash audit") return for path in metrics_files: try: payload = json.loads(path.read_text()) except json.JSONDecodeError: warnings.append(f"{path}: invalid JSON; skipped hash audit") continue if "data_hash" not in payload and not _sidecar_has_content(path.parent / "data_hash.txt"): warnings.append(f"{path}: missing data_hash") if "split_hash" not in payload and not _sidecar_has_content(path.parent / "split_hash.txt"): warnings.append(f"{path}: missing split_hash") def _sidecar_has_content(path: Path) -> bool: return path.exists() and bool(path.read_text().strip()) def _markdown(report: dict[str, Any]) -> str: lines = [ "# CIL Chart Leakage Audit", "", f"Status: `{report['status']}`", "", "| Split | Rows | Charts | Outcomes in DB | Retrieval allowed | Audience |", "| --- | ---: | ---: | --- | --- | --- |", ] for split, payload in sorted(report["indexes"].items()): lines.append( f"| {split} | {payload.get('num_rows')} | {payload.get('num_groups_exported')} | " f"{payload.get('include_outcomes')} | {payload.get('retrieval_index_allowed')} | " f"{payload.get('audience')} |" ) lines.extend(["", "## Violations", ""]) lines.extend(f"- {item}" for item in report["violations"]) if report["violations"] else lines.append("- None") lines.extend(["", "## Warnings", ""]) lines.extend(f"- {item}" for item in report["warnings"]) if report["warnings"] else lines.append("- None") return "\n".join(lines) if __name__ == "__main__": raise SystemExit(main())