mirror of
https://github.com/safishamsi/graphify.git
synced 2026-07-13 10:57:13 +00:00
merge PR #663: parallel AST extraction with ProcessPoolExecutor (1.66x speedup on 84 files)
Co-Authored-By: hanzala-sohrab <hanzala-sohrab@users.noreply.github.com> Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Safi
+189
-59
@@ -3650,7 +3650,167 @@ def _check_tree_sitter_version() -> None:
|
||||
)
|
||||
|
||||
|
||||
def extract(paths: list[Path], cache_root: Path | None = None) -> dict:
|
||||
_DISPATCH: dict[str, Any] = {
|
||||
".py": extract_python,
|
||||
".js": extract_js,
|
||||
".jsx": extract_js,
|
||||
".mjs": extract_js,
|
||||
".ts": extract_js,
|
||||
".tsx": extract_js,
|
||||
".go": extract_go,
|
||||
".rs": extract_rust,
|
||||
".java": extract_java,
|
||||
".c": extract_c,
|
||||
".h": extract_c,
|
||||
".cpp": extract_cpp,
|
||||
".cc": extract_cpp,
|
||||
".cxx": extract_cpp,
|
||||
".hpp": extract_cpp,
|
||||
".rb": extract_ruby,
|
||||
".cs": extract_csharp,
|
||||
".kt": extract_kotlin,
|
||||
".kts": extract_kotlin,
|
||||
".scala": extract_scala,
|
||||
".php": extract_php,
|
||||
".swift": extract_swift,
|
||||
".lua": extract_lua,
|
||||
".toc": extract_lua,
|
||||
".zig": extract_zig,
|
||||
".ps1": extract_powershell,
|
||||
".ex": extract_elixir,
|
||||
".exs": extract_elixir,
|
||||
".m": extract_objc,
|
||||
".mm": extract_objc,
|
||||
".jl": extract_julia,
|
||||
".vue": extract_js,
|
||||
".svelte": extract_js,
|
||||
".dart": extract_dart,
|
||||
".v": extract_verilog,
|
||||
".sv": extract_verilog,
|
||||
".sql": extract_sql,
|
||||
}
|
||||
|
||||
|
||||
def _get_extractor(path: Path) -> Any | None:
|
||||
"""Return the correct extractor function for a file, or None if unsupported."""
|
||||
if path.name.endswith(".blade.php"):
|
||||
return extract_blade
|
||||
return _DISPATCH.get(path.suffix)
|
||||
|
||||
|
||||
def _extract_single_file(args: tuple) -> tuple[int, dict]:
|
||||
"""Worker function for parallel extraction. Runs in a subprocess.
|
||||
|
||||
Must be at module level (not a closure) so it can be pickled by
|
||||
ProcessPoolExecutor.
|
||||
|
||||
Args:
|
||||
args: (index, path_str, cache_root_str) tuple
|
||||
|
||||
Returns:
|
||||
(index, result_dict) so results can be placed back in order.
|
||||
"""
|
||||
idx, path_str, cache_root_str = args
|
||||
path = Path(path_str)
|
||||
cache_root = Path(cache_root_str)
|
||||
|
||||
# Check cache first (avoid re-extraction)
|
||||
cached = load_cached(path, cache_root)
|
||||
if cached is not None:
|
||||
return idx, cached
|
||||
|
||||
extractor = _get_extractor(path)
|
||||
if extractor is None:
|
||||
return idx, {"nodes": [], "edges": []}
|
||||
|
||||
result = extractor(path)
|
||||
if "error" not in result:
|
||||
save_cached(path, result, cache_root)
|
||||
return idx, result
|
||||
|
||||
|
||||
def _extract_parallel(
|
||||
uncached_work: list[tuple[int, Path]],
|
||||
per_file: list[dict | None],
|
||||
effective_root: Path,
|
||||
max_workers: int | None,
|
||||
total_files: int,
|
||||
) -> None:
|
||||
"""Extract uncached files in parallel using ProcessPoolExecutor."""
|
||||
import concurrent.futures
|
||||
|
||||
if max_workers is None:
|
||||
max_workers = min(os.cpu_count() or 4, len(uncached_work), 8)
|
||||
|
||||
root_str = str(effective_root)
|
||||
work_items = [(idx, str(path), root_str) for idx, path in uncached_work]
|
||||
|
||||
done_count = 0
|
||||
_PROGRESS_INTERVAL = 100
|
||||
with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as pool:
|
||||
futures = {
|
||||
pool.submit(_extract_single_file, item): item[0] for item in work_items
|
||||
}
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
idx, result = future.result()
|
||||
per_file[idx] = result
|
||||
done_count += 1
|
||||
if (
|
||||
total_files >= _PROGRESS_INTERVAL
|
||||
and done_count % _PROGRESS_INTERVAL == 0
|
||||
):
|
||||
print(
|
||||
f" AST extraction: {done_count}/{len(uncached_work)} uncached files "
|
||||
f"({done_count * 100 // len(uncached_work)}%) [{max_workers} workers]",
|
||||
flush=True,
|
||||
)
|
||||
if total_files >= _PROGRESS_INTERVAL:
|
||||
print(
|
||||
f" AST extraction: {total_files}/{total_files} files (100%) [{max_workers} workers]",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def _extract_sequential(
|
||||
uncached_work: list[tuple[int, Path]],
|
||||
per_file: list[dict | None],
|
||||
effective_root: Path,
|
||||
total_files: int,
|
||||
) -> None:
|
||||
"""Extract uncached files sequentially (fallback for small batches)."""
|
||||
_PROGRESS_INTERVAL = 100
|
||||
for work_idx, (idx, path) in enumerate(uncached_work):
|
||||
if (
|
||||
total_files >= _PROGRESS_INTERVAL
|
||||
and work_idx % _PROGRESS_INTERVAL == 0
|
||||
and work_idx > 0
|
||||
):
|
||||
print(
|
||||
f" AST extraction: {work_idx}/{len(uncached_work)} uncached files ({work_idx * 100 // len(uncached_work)}%)",
|
||||
flush=True,
|
||||
)
|
||||
extractor = _get_extractor(path)
|
||||
if extractor is None:
|
||||
per_file[idx] = {"nodes": [], "edges": []}
|
||||
continue
|
||||
result = extractor(path)
|
||||
if "error" not in result:
|
||||
save_cached(path, result, effective_root)
|
||||
per_file[idx] = result
|
||||
if total_files >= _PROGRESS_INTERVAL:
|
||||
print(f" AST extraction: {total_files}/{total_files} files (100%)", flush=True)
|
||||
|
||||
|
||||
_PARALLEL_THRESHOLD = 20
|
||||
|
||||
|
||||
def extract(
|
||||
paths: list[Path],
|
||||
cache_root: Path | None = None,
|
||||
*,
|
||||
parallel: bool = True,
|
||||
max_workers: int | None = None,
|
||||
) -> dict:
|
||||
"""Extract AST nodes and edges from a list of code files.
|
||||
|
||||
Two-pass process:
|
||||
@@ -3663,9 +3823,11 @@ def extract(paths: list[Path], cache_root: Path | None = None) -> dict:
|
||||
cache_root: explicit root for graphify-out/cache/ (overrides the
|
||||
inferred common path prefix). Pass Path('.') when running on a
|
||||
subdirectory so the cache stays at ./graphify-out/cache/.
|
||||
parallel: if True and there are >= _PARALLEL_THRESHOLD uncached files,
|
||||
use ProcessPoolExecutor for multi-core extraction.
|
||||
max_workers: max subprocess count. Defaults to min(cpu_count, 8).
|
||||
"""
|
||||
_check_tree_sitter_version()
|
||||
per_file: list[dict] = []
|
||||
|
||||
# Infer a common root for cache keys (use first diverging segment, not sum of all matches)
|
||||
try:
|
||||
@@ -3686,68 +3848,36 @@ def extract(paths: list[Path], cache_root: Path | None = None) -> dict:
|
||||
root = Path(".")
|
||||
root = root.resolve()
|
||||
|
||||
_DISPATCH: dict[str, Any] = {
|
||||
".py": extract_python,
|
||||
".js": extract_js,
|
||||
".jsx": extract_js,
|
||||
".mjs": extract_js,
|
||||
".ts": extract_js,
|
||||
".tsx": extract_js,
|
||||
".go": extract_go,
|
||||
".rs": extract_rust,
|
||||
".java": extract_java,
|
||||
".c": extract_c,
|
||||
".h": extract_c,
|
||||
".cpp": extract_cpp,
|
||||
".cc": extract_cpp,
|
||||
".cxx": extract_cpp,
|
||||
".hpp": extract_cpp,
|
||||
".rb": extract_ruby,
|
||||
".cs": extract_csharp,
|
||||
".kt": extract_kotlin,
|
||||
".kts": extract_kotlin,
|
||||
".scala": extract_scala,
|
||||
".php": extract_php,
|
||||
".swift": extract_swift,
|
||||
".lua": extract_lua,
|
||||
".toc": extract_lua,
|
||||
".zig": extract_zig,
|
||||
".ps1": extract_powershell,
|
||||
".ex": extract_elixir,
|
||||
".exs": extract_elixir,
|
||||
".m": extract_objc,
|
||||
".mm": extract_objc,
|
||||
".jl": extract_julia,
|
||||
".vue": extract_js,
|
||||
".svelte": extract_js,
|
||||
".dart": extract_dart,
|
||||
".v": extract_verilog,
|
||||
".sv": extract_verilog,
|
||||
".sql": extract_sql,
|
||||
}
|
||||
|
||||
effective_root = cache_root or root
|
||||
total = len(paths)
|
||||
_PROGRESS_INTERVAL = 100
|
||||
|
||||
# Phase 1: separate cached hits from uncached work
|
||||
per_file: list[dict | None] = [None] * total
|
||||
uncached_work: list[tuple[int, Path]] = []
|
||||
|
||||
for i, path in enumerate(paths):
|
||||
if total >= _PROGRESS_INTERVAL and i % _PROGRESS_INTERVAL == 0 and i > 0:
|
||||
print(f" AST extraction: {i}/{total} files ({i * 100 // total}%)", flush=True)
|
||||
# .blade.php must be checked before suffix lookup since Path.suffix returns .php
|
||||
if path.name.endswith(".blade.php"):
|
||||
extractor = extract_blade
|
||||
else:
|
||||
extractor = _DISPATCH.get(path.suffix)
|
||||
if extractor is None:
|
||||
if _get_extractor(path) is None:
|
||||
per_file[i] = {"nodes": [], "edges": []}
|
||||
continue
|
||||
cached = load_cached(path, cache_root or root)
|
||||
cached = load_cached(path, effective_root)
|
||||
if cached is not None:
|
||||
per_file.append(cached)
|
||||
per_file[i] = cached
|
||||
continue
|
||||
result = extractor(path)
|
||||
if "error" not in result:
|
||||
save_cached(path, result, cache_root or root)
|
||||
per_file.append(result)
|
||||
if total >= _PROGRESS_INTERVAL:
|
||||
print(f" AST extraction: {total}/{total} files (100%)", flush=True)
|
||||
uncached_work.append((i, path))
|
||||
|
||||
# Phase 2: extract uncached files (parallel or sequential)
|
||||
if uncached_work:
|
||||
if parallel and len(uncached_work) >= _PARALLEL_THRESHOLD:
|
||||
_extract_parallel(
|
||||
uncached_work, per_file, effective_root, max_workers, total
|
||||
)
|
||||
else:
|
||||
_extract_sequential(uncached_work, per_file, effective_root, total)
|
||||
|
||||
# Fill any remaining None slots (shouldn't happen, but defensive)
|
||||
for i in range(total):
|
||||
if per_file[i] is None:
|
||||
per_file[i] = {"nodes": [], "edges": []}
|
||||
|
||||
all_nodes: list[dict] = []
|
||||
all_edges: list[dict] = []
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Benchmark: sequential vs parallel AST extraction.
|
||||
|
||||
Usage:
|
||||
python tests/bench_extract.py [path-to-repo]
|
||||
|
||||
Defaults to the current directory if no path is given.
|
||||
Clears the AST cache between runs so every file is re-extracted.
|
||||
|
||||
Example output:
|
||||
=== Graphify AST Extraction Benchmark ===
|
||||
Files: 1,247
|
||||
Languages: Python (412), TypeScript (389), Go (201), ...
|
||||
|
||||
Sequential: 4.32s (8,934 nodes, 12,456 edges)
|
||||
Parallel (8): 1.28s (8,934 nodes, 12,456 edges)
|
||||
|
||||
Speedup: 3.38x
|
||||
Results: ✓ identical
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
# Ensure the project root is importable
|
||||
_project_root = Path(__file__).resolve().parent.parent
|
||||
if str(_project_root) not in sys.path:
|
||||
sys.path.insert(0, str(_project_root))
|
||||
|
||||
from graphify.extract import extract, collect_files
|
||||
from graphify.cache import clear_cache
|
||||
|
||||
|
||||
def _count_by_ext(paths: list[Path]) -> dict[str, int]:
|
||||
"""Count files by extension."""
|
||||
counter: Counter[str] = Counter()
|
||||
for p in paths:
|
||||
ext = p.suffix.lower()
|
||||
counter[ext] += 1
|
||||
return dict(counter.most_common())
|
||||
|
||||
|
||||
_EXT_NAMES: dict[str, str] = {
|
||||
".py": "Python",
|
||||
".js": "JavaScript",
|
||||
".jsx": "JSX",
|
||||
".mjs": "MJS",
|
||||
".ts": "TypeScript",
|
||||
".tsx": "TSX",
|
||||
".go": "Go",
|
||||
".rs": "Rust",
|
||||
".java": "Java",
|
||||
".c": "C",
|
||||
".h": "C Header",
|
||||
".cpp": "C++",
|
||||
".cc": "C++",
|
||||
".cxx": "C++",
|
||||
".hpp": "C++ Header",
|
||||
".rb": "Ruby",
|
||||
".cs": "C#",
|
||||
".kt": "Kotlin",
|
||||
".kts": "Kotlin Script",
|
||||
".scala": "Scala",
|
||||
".php": "PHP",
|
||||
".swift": "Swift",
|
||||
".lua": "Lua",
|
||||
".toc": "Lua TOC",
|
||||
".zig": "Zig",
|
||||
".ps1": "PowerShell",
|
||||
".ex": "Elixir",
|
||||
".exs": "Elixir Script",
|
||||
".m": "Obj-C",
|
||||
".mm": "Obj-C++",
|
||||
".jl": "Julia",
|
||||
".vue": "Vue",
|
||||
".svelte": "Svelte",
|
||||
".dart": "Dart",
|
||||
".v": "Verilog",
|
||||
".sv": "SystemVerilog",
|
||||
".sql": "SQL",
|
||||
}
|
||||
|
||||
|
||||
def _format_languages(ext_counts: dict[str, int]) -> str:
|
||||
parts = []
|
||||
for ext, count in ext_counts.items():
|
||||
name = _EXT_NAMES.get(ext, ext)
|
||||
parts.append(f"{name} ({count})")
|
||||
return ", ".join(parts)
|
||||
|
||||
|
||||
def _run_extraction(
|
||||
paths: list[Path],
|
||||
cache_root: Path,
|
||||
parallel: bool,
|
||||
max_workers: int | None = None,
|
||||
) -> tuple[float, int, int]:
|
||||
"""Run extraction, return (elapsed_seconds, node_count, edge_count)."""
|
||||
clear_cache(cache_root)
|
||||
t0 = time.perf_counter()
|
||||
result = extract(
|
||||
paths, cache_root=cache_root, parallel=parallel, max_workers=max_workers
|
||||
)
|
||||
elapsed = time.perf_counter() - t0
|
||||
nodes = len(result.get("nodes", []))
|
||||
edges = len(result.get("edges", []))
|
||||
return elapsed, nodes, edges
|
||||
|
||||
|
||||
def main() -> None:
|
||||
target = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(".")
|
||||
target = target.resolve()
|
||||
|
||||
if not target.exists():
|
||||
print(f"Error: {target} does not exist", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print("=== Graphify AST Extraction Benchmark ===\n")
|
||||
print(f"Scanning {target} ...", flush=True)
|
||||
|
||||
paths = collect_files(target)
|
||||
if not paths:
|
||||
print("No extractable files found.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
ext_counts = _count_by_ext(paths)
|
||||
print(f"Files: {len(paths):,}")
|
||||
print(f"Languages: {_format_languages(ext_counts)}")
|
||||
print()
|
||||
|
||||
cache_root = target if target.is_dir() else target.parent
|
||||
|
||||
# Workers count (same logic as _extract_parallel)
|
||||
import os
|
||||
|
||||
workers = min(os.cpu_count() or 4, len(paths), 8)
|
||||
|
||||
# Run sequential
|
||||
print("Running sequential extraction...", flush=True)
|
||||
seq_time, seq_nodes, seq_edges = _run_extraction(paths, cache_root, parallel=False)
|
||||
print(f"Sequential: {seq_time:.2f}s ({seq_nodes:,} nodes, {seq_edges:,} edges)")
|
||||
|
||||
# Run parallel
|
||||
print(f"\nRunning parallel extraction ({workers} workers)...", flush=True)
|
||||
par_time, par_nodes, par_edges = _run_extraction(
|
||||
paths, cache_root, parallel=True, max_workers=workers
|
||||
)
|
||||
print(
|
||||
f"Parallel ({workers}): {par_time:.2f}s ({par_nodes:,} nodes, {par_edges:,} edges)"
|
||||
)
|
||||
|
||||
# Results
|
||||
print()
|
||||
if seq_time > 0:
|
||||
speedup = seq_time / par_time if par_time > 0 else float("inf")
|
||||
print(f"Speedup: {speedup:.2f}x")
|
||||
print(f"Workers: {workers} (auto-detected)")
|
||||
|
||||
# Validate correctness
|
||||
if seq_nodes == par_nodes and seq_edges == par_edges:
|
||||
print("Results: ✓ identical (node count, edge count match)")
|
||||
else:
|
||||
print("Results: ✗ MISMATCH!")
|
||||
print(f" Sequential: {seq_nodes} nodes, {seq_edges} edges")
|
||||
print(f" Parallel: {par_nodes} nodes, {par_edges} edges")
|
||||
sys.exit(1)
|
||||
|
||||
# Clean up cache after benchmark
|
||||
clear_cache(cache_root)
|
||||
print("\nCache cleared after benchmark.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user