Add polyglot parity check script

check_polyglot_parity.py

@@ -0,0 +1,250 @@
+"""
+Parity check: does the trajectory-eval shallow clone produce the same
+polyglot-parsed graph + BERT features as the pre-existing big-machine
+clone (data_multilang), for the same commit?
+
+Runs on big where both clones exist. For each common (repo, commit) pair
+it encounters, it snapshots the working tree from *both* clones, canonical-
+hashes the graph structure + feature tensors, and reports match/mismatch.
+
+A match confirms:
+  - git clone --filter=blob:none + checkout fetches the same file content
+    as the original full clone
+  - parse_repo_polyglot is deterministic w.r.t. the file tree (modulo
+    rglob ordering — we sort before hashing)
+  - BertTokenEmbedder is deterministic
+
+Usage (on big):
+    python -m graphjepa.check_polyglot_parity \\
+        --traj-repos  ./outputs/traj_real/repos \\
+        --multi-repos /raid/train/datasets/code-graph-v7/data_multilang \\
+        --n-pairs 4
+
+If it picks a commit not present in the shallow clone's blobless ref set
+(some base_commits may need lazy blob fetch), the script does the fetch
+automatically via checkout.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import subprocess
+import sys
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+
+def run(cmd: List[str], cwd: Optional[Path] = None, check: bool = True):
+    r = subprocess.run(cmd, cwd=str(cwd) if cwd else None,
+                        capture_output=True, text=True)
+    if check and r.returncode != 0:
+        raise RuntimeError(f'{" ".join(cmd)} failed: {r.stderr[-400:]}')
+    return r
+
+
+def list_commits(repo_dir: Path, n: int = 20) -> List[str]:
+    r = run(['git', 'log', '--format=%H', '-n', str(n)], cwd=repo_dir)
+    return r.stdout.split()
+
+
+def checkout(repo_dir: Path, sha: str) -> bool:
+    run(['git', 'reset', '--hard', '-q'], cwd=repo_dir, check=False)
+    run(['git', 'clean', '-fdx', '-q'], cwd=repo_dir, check=False)
+    r = run(['git', 'checkout', '-q', '--detach', sha], cwd=repo_dir, check=False)
+    if r.returncode != 0:
+        # Try fetching the ref
+        run(['git', 'fetch', '-q', 'origin', sha], cwd=repo_dir, check=False)
+        r = run(['git', 'checkout', '-q', '--detach', sha], cwd=repo_dir,
+                 check=False)
+    return r.returncode == 0
+
+
+def canonical_hash(graph, features) -> Tuple[str, str, dict]:
+    """Deterministic hash of (graph structure, feature tensors).
+    Sorts node IDs so walk order doesn't matter.
+    Returns (graph_hash, feature_hash, stats_dict).
+    """
+    import torch
+
+    # Nodes: sort by id, hash (id, kind, content, type_description).
+    h_nodes = hashlib.sha256()
+    node_items = sorted(graph.nodes.items())
+    for nid, n in node_items:
+        h_nodes.update(nid.encode())
+        h_nodes.update(b'\x00')
+        h_nodes.update(getattr(n.kind, 'value', str(n.kind)).encode())
+        h_nodes.update(b'\x00')
+        h_nodes.update((n.content or '').encode())
+        h_nodes.update(b'\x00')
+        h_nodes.update((n.type_description or '').encode())
+        h_nodes.update(b'\x01')
+
+    # Edges: sort by (src, dst, kind).
+    edge_keys = sorted(
+        (e.src, e.dst, getattr(e.kind, 'value', str(e.kind)))
+        for e in graph.edges.values()
+    )
+    for src, dst, k in edge_keys:
+        h_nodes.update(f'E|{src}|{dst}|{k}|'.encode())
+
+    graph_hash = h_nodes.hexdigest()
+
+    # Feature tensors: for each kind in deterministic order, hash
+    # (sorted_ids, content_sum, type_sum, content_first_vec, type_first_vec).
+    h_feats = hashlib.sha256()
+    for kind, d in sorted((k, v) for k, v in features.items() if v is not None):
+        kind_str = getattr(kind, 'value', str(kind))
+        h_feats.update(kind_str.encode())
+        h_feats.update(b'\x00')
+        ids = list(d['ids'])
+        sort_idx = sorted(range(len(ids)), key=lambda i: ids[i])
+        content = d['content'][sort_idx] if sort_idx else d['content']
+        typev   = d['type'][sort_idx] if sort_idx else d['type']
+        sorted_ids = [ids[i] for i in sort_idx]
+        for sid in sorted_ids:
+            h_feats.update(sid.encode()); h_feats.update(b'\x00')
+        # Digest feature tensors numerically with fixed precision so
+        # hashes match across float ops that might differ in trailing ULP.
+        content_q = (content * 1e5).round().to(torch.int64)
+        typev_q   = (typev   * 1e5).round().to(torch.int64)
+        h_feats.update(content_q.cpu().numpy().tobytes())
+        h_feats.update(typev_q.cpu().numpy().tobytes())
+
+    feat_hash = h_feats.hexdigest()
+
+    stats = {
+        'n_nodes': len(graph.nodes),
+        'n_edges': len(graph.edges),
+        'n_feat_kinds': sum(1 for v in features.values() if v is not None),
+        'feat_dim': next((v['content'].shape[1] for v in features.values()
+                          if v is not None), None),
+    }
+    return graph_hash, feat_hash, stats
+
+
+def snapshot(repo_dir: Path, embedder) -> Tuple[str, str, dict]:
+    from graphjepa.trajectory_pipeline import snapshot_working_tree
+    g, feats = snapshot_working_tree(repo_dir, embedder, verbose=False)
+    return canonical_hash(g, feats)
+
+
+# Mapping from trajectory-eval repo dirname → data_multilang subpath.
+# traj repos: django__django; data_multilang: python/django
+_REPO_DIR_MAP = {
+    'django__django': ('python', 'django'),
+    'sympy__sympy': ('python', 'sympy'),
+    'sphinx-doc__sphinx': ('python', 'sphinx'),
+    'matplotlib__matplotlib': ('python', 'matplotlib'),
+    'scikit-learn__scikit-learn': ('python', 'scikit-learn'),
+    'astropy__astropy': ('python', 'astropy'),
+    'pydata__xarray': ('python', 'xarray'),
+    'pytest-dev__pytest': ('python', 'pytest'),
+    'pylint-dev__pylint': ('python', 'pylint'),
+    'psf__requests': ('python', 'requests'),
+    'mwaskom__seaborn': ('python', 'seaborn'),
+    'pallets__flask': ('python', 'flask'),
+}
+
+
+def find_pairs(traj_root: Path, multi_root: Path) -> List[Tuple[str, Path, Path]]:
+    pairs = []
+    if not traj_root.is_dir():
+        return pairs
+    for name, (lang, mname) in _REPO_DIR_MAP.items():
+        tpath = traj_root / name
+        mpath = multi_root / lang / mname
+        if tpath.is_dir() and mpath.is_dir():
+            pairs.append((name, tpath, mpath))
+    return pairs
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--traj-repos', required=True,
+                   help='outputs/traj_real/repos dir from the transfer bundle')
+    p.add_argument('--multi-repos', required=True,
+                   help='data_multilang dir used to build cache_v7')
+    p.add_argument('--n-pairs', type=int, default=3,
+                   help='Number of (repo, commit) pairs to test')
+    p.add_argument('--output', default=None,
+                   help='Write a JSON report here')
+    args = p.parse_args()
+
+    traj_root = Path(args.traj_repos)
+    multi_root = Path(args.multi_repos)
+
+    pairs = find_pairs(traj_root, multi_root)
+    if not pairs:
+        print(f'[parity] no common repos found under {traj_root} and '
+              f'{multi_root}'); sys.exit(1)
+    print(f'[parity] {len(pairs)} repo pairs available:')
+    for n, t, m in pairs:
+        print(f'  {n:30s}  traj={t}  multi={m}')
+
+    # For each pair, pick a commit that exists in both. HEAD of the
+    # multi clone is a safe default since that clone has full history.
+    tests = []
+    for name, tpath, mpath in pairs[:args.n_pairs]:
+        mcommits = list_commits(mpath, n=5)
+        if not mcommits:
+            print(f'[parity] {name}: no commits in multi clone, skip')
+            continue
+        tests.append((name, tpath, mpath, mcommits[0]))
+
+    # Import embedder once — BERT load is slow.
+    from graphjepa.features import BertTokenEmbedder
+    print('\n[parity] loading BERT embedder ...')
+    embedder = BertTokenEmbedder(device='cpu')
+
+    results = []
+    for name, tpath, mpath, sha in tests:
+        print(f'\n[parity] === {name} @ {sha[:10]} ===')
+
+        print(f'  checkout traj clone ...')
+        if not checkout(tpath, sha):
+            print(f'  [parity] traj clone cannot reach {sha[:10]}; skip')
+            results.append({'repo': name, 'sha': sha, 'error': 'traj_checkout_failed'})
+            continue
+        print(f'  checkout multi clone ...')
+        if not checkout(mpath, sha):
+            print(f'  [parity] multi clone cannot reach {sha[:10]}; skip')
+            results.append({'repo': name, 'sha': sha, 'error': 'multi_checkout_failed'})
+            continue
+
+        print(f'  snapshotting traj clone ...')
+        tg, tf, tstats = snapshot(tpath, embedder)
+        print(f'  snapshotting multi clone ...')
+        mg, mf, mstats = snapshot(mpath, embedder)
+
+        match_g = tg == mg
+        match_f = tf == mf
+        print(f'  graph hash   traj={tg[:12]}  multi={mg[:12]}  '
+              f'{"MATCH" if match_g else "MISMATCH"}')
+        print(f'  feature hash traj={tf[:12]}  multi={mf[:12]}  '
+              f'{"MATCH" if match_f else "MISMATCH"}')
+        print(f'  stats traj={tstats}  multi={mstats}')
+        results.append({
+            'repo': name, 'sha': sha,
+            'graph_match': match_g, 'feature_match': match_f,
+            'traj_stats': tstats, 'multi_stats': mstats,
+        })
+
+    print('\n' + '=' * 60)
+    n_g = sum(1 for r in results if r.get('graph_match'))
+    n_f = sum(1 for r in results if r.get('feature_match'))
+    print(f'graph parity:   {n_g}/{len(results)} matched')
+    print(f'feature parity: {n_f}/{len(results)} matched')
+    print('=' * 60)
+
+    if args.output:
+        with open(args.output, 'w') as f:
+            json.dump(results, f, indent=2)
+        print(f'[parity] report saved: {args.output}')
+
+    sys.exit(0 if (n_g == n_f == len(results) and results) else 1)
+
+
+if __name__ == '__main__':
+    main()