From 91d5fd2d2af34c4ea95850ce4b27169162886bd8 Mon Sep 17 00:00:00 2001 From: ankitmishralive <=ankitmishra.letter@gmail.com> Date: Sun, 12 Jul 2026 15:11:22 +0530 Subject: [PATCH] perf(cli): parallelize extraction stages --- graphify/cli.py | 299 +++++++++++++-------------- tests/test_parallel_orchestration.py | 169 +++++++++++++++ 2 files changed, 311 insertions(+), 157 deletions(-) create mode 100644 tests/test_parallel_orchestration.py diff --git a/graphify/cli.py b/graphify/cli.py index 25bf5499b..63d1d28e1 100644 --- a/graphify/cli.py +++ b/graphify/cli.py @@ -11,7 +11,7 @@ import sys from graphify.paths import GRAPHIFY_OUT as _GRAPHIFY_OUT from pathlib import Path - +import concurrent.futures _SEARCH_NUDGE = json.dumps({ "hookSpecificOutput": { @@ -2249,170 +2249,155 @@ def _parse_float(name: str, raw: str) -> float: ) sys.exit(1) - # AST extraction on code files. Empty code list (docs-only corpus) is - # the issue #698 case — skip cleanly instead of crashing inside extract(). - ast_result: dict = {"nodes": [], "edges": [], "input_tokens": 0, "output_tokens": 0} - if code_files: - from graphify.extract import extract as _ast_extract - # Anchor the cache at the output root, not the scanned project: - # with --out, a /graphify-out/cache/ would leak a - # graphify-out/ dir into a project that asked for external output. - ast_kwargs: dict = {"cache_root": out_root} - if cli_max_workers is not None: - ast_kwargs["max_workers"] = cli_max_workers - print(f"[graphify extract] AST extraction on {len(code_files)} code files...") - try: - ast_result = _ast_extract(code_files, **ast_kwargs) - except Exception as exc: - print(f"[graphify extract] AST extraction failed: {exc}", file=sys.stderr) - ast_result = {"nodes": [], "edges": [], "input_tokens": 0, "output_tokens": 0} - stages.mark("AST extract") - - # Semantic extraction on docs/papers/images. Check cache first. - from graphify.cache import ( - check_semantic_cache as _check_semantic_cache, - prune_semantic_cache as _prune_semantic_cache, - save_semantic_cache as _save_semantic_cache, - ) - sem_result: dict = { - "nodes": [], "edges": [], "hyperedges": [], - "input_tokens": 0, "output_tokens": 0, - } - sem_cache_hits = 0 - sem_cache_misses = 0 - if semantic_files: - sem_paths_str = [str(p) for p in semantic_files] - cached_nodes, cached_edges, cached_hyperedges, uncached_paths = ( - _check_semantic_cache(sem_paths_str, root=out_root) - ) - sem_cache_hits = len(semantic_files) - len(uncached_paths) - sem_cache_misses = len(uncached_paths) - sem_result["nodes"].extend(cached_nodes) - sem_result["edges"].extend(cached_edges) - sem_result["hyperedges"].extend(cached_hyperedges) - if sem_cache_hits: - print(f"[graphify extract] semantic cache: {sem_cache_hits} hit / {sem_cache_misses} miss") - - if uncached_paths: - print(f"[graphify extract] semantic extraction on {len(uncached_paths)} files via {backend}...") - corpus_kwargs: dict = { - "backend": backend, - "model": model, - "root": target, - } - if deep_mode: - corpus_kwargs["deep_mode"] = True - if cli_token_budget is not None: - corpus_kwargs["token_budget"] = cli_token_budget - if cli_max_concurrency is not None: - corpus_kwargs["max_concurrency"] = cli_max_concurrency - - # Minimal progress callback so the CLI is no longer silent - # during long local-inference runs (issue #792 addendum). - # Also track per-chunk success so we can fail loudly when - # every chunk errors (e.g. missing backend SDK package). - _chunk_stats = {"total": 0, "succeeded": 0} - def _progress(idx: int, total: int, _result: dict) -> None: - _chunk_stats["total"] = total - _chunk_stats["succeeded"] += 1 - print( - f"[graphify extract] chunk {idx + 1}/{total} done", - flush=True, - ) - corpus_kwargs["on_chunk_done"] = _progress + + def run_ast(): + res = {"nodes": [], "edges": [], "input_tokens": 0, "output_tokens": 0} + if code_files: + from graphify.extract import extract as _ast_extract + ast_kwargs: dict = {"cache_root": out_root} + if cli_max_workers is not None: + ast_kwargs["max_workers"] = cli_max_workers + print(f"[graphify extract] AST extraction on {len(code_files)} code files...") try: - fresh = _extract_corpus_parallel( - [Path(p) for p in uncached_paths], - **corpus_kwargs, - ) - except ImportError as exc: - print(f"error: {exc}", file=sys.stderr) - sys.exit(1) + res = _ast_extract(code_files, **ast_kwargs) except Exception as exc: - print( - f"[graphify extract] semantic extraction failed: {exc}", - file=sys.stderr, - ) - fresh = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0} + print(f"[graphify extract] AST extraction failed: {exc}", file=sys.stderr) + return res + + def run_semantic(): + res = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0} + sem_cache_hits = 0 + sem_cache_misses = 0 + if semantic_files: + from graphify.cache import ( + check_semantic_cache as _check_semantic_cache, + save_semantic_cache as _save_semantic_cache, + ) + sem_paths_str = [str(p) for p in semantic_files] + cached_nodes, cached_edges, cached_hyperedges, uncached_paths = ( + _check_semantic_cache(sem_paths_str, root=out_root) + ) + sem_cache_hits = len(semantic_files) - len(uncached_paths) + sem_cache_misses = len(uncached_paths) + res["nodes"].extend(cached_nodes) + res["edges"].extend(cached_edges) + res["hyperedges"].extend(cached_hyperedges) + if sem_cache_hits: + print(f"[graphify extract] semantic cache: {sem_cache_hits} hit / {sem_cache_misses} miss") + + if uncached_paths: + print(f"[graphify extract] semantic extraction on {len(uncached_paths)} files via {backend}...") + corpus_kwargs: dict = { + "backend": backend, + "model": model, + "root": target, + } + if deep_mode: + corpus_kwargs["deep_mode"] = True + if cli_token_budget is not None: + corpus_kwargs["token_budget"] = cli_token_budget + if cli_max_concurrency is not None: + corpus_kwargs["max_concurrency"] = cli_max_concurrency + + _chunk_stats = {"total": 0, "succeeded": 0} + def _progress(idx: int, total: int, _result: dict) -> None: + _chunk_stats["total"] = total + _chunk_stats["succeeded"] += 1 + print(f"[graphify extract] chunk {idx + 1}/{total} done", flush=True) + corpus_kwargs["on_chunk_done"] = _progress - # on_chunk_done only fires after a chunk succeeds. If fresh - # semantic extraction was requested and no chunks completed, - # fail instead of writing an AST-only graph with exit 0. - if uncached_paths and _chunk_stats["succeeded"] == 0: - print( - f"[graphify extract] error: all semantic chunks failed " - f"for backend '{backend}' ({len(uncached_paths)} uncached files) - " - f"see per-chunk errors above. If you see 'requires the X package', " - f"run `pip install X` and retry.", - file=sys.stderr, - ) - sys.exit(1) - try: - _save_semantic_cache( - fresh.get("nodes", []), - fresh.get("edges", []), - fresh.get("hyperedges", []), - root=out_root, - ) - except Exception as exc: - print(f"[graphify extract] warning: could not write semantic cache: {exc}", file=sys.stderr) - sem_result["nodes"].extend(fresh.get("nodes", [])) - sem_result["edges"].extend(fresh.get("edges", [])) - sem_result["hyperedges"].extend(fresh.get("hyperedges", [])) - sem_result["input_tokens"] += fresh.get("input_tokens", 0) - sem_result["output_tokens"] += fresh.get("output_tokens", 0) - - # Prune orphaned semantic cache entries. The semantic cache is - # content-hash-keyed and unversioned, so it is never swept by the AST - # version-cleanup: every content change or file deletion leaves a - # permanent orphan that accumulates unbounded (#1527). Sweep it against - # the FULL live document set (``files_by_type`` — present in both the - # incremental and full branches), NOT the incremental ``semantic_files`` - # changed-subset, which would delete every unchanged doc's valid entry. - # Best-effort: a prune failure must never break extraction. - try: - from graphify.cache import file_hash as _file_hash - _live_hashes: set[str] = set() - for _kind in ("document", "paper", "image"): - for _fp in files_by_type.get(_kind, []): - _abs = Path(_fp) - if not _abs.is_absolute(): - _abs = Path(out_root) / _abs - if not _abs.is_file(): - continue # deleted/missing — leave out so its entry is pruned try: - _live_hashes.add(_file_hash(_abs, out_root)) - except OSError: - pass - _prune_semantic_cache(out_root, _live_hashes) - except Exception as exc: - print(f"[graphify extract] warning: could not prune semantic cache: {exc}", file=sys.stderr) - stages.mark("semantic extract") + fresh = _extract_corpus_parallel([Path(p) for p in uncached_paths], **corpus_kwargs) + except ImportError as exc: + print(f"error: {exc}", file=sys.stderr) + sys.exit(1) + except Exception as exc: + print(f"[graphify extract] semantic extraction failed: {exc}", file=sys.stderr) + fresh = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0} - pg_result: dict = {"nodes": [], "edges": []} - if cli_postgres_dsn is not None: - from graphify.pg_introspect import introspect_postgres - print(f"[graphify extract] introspecting PostgreSQL schema...") + if uncached_paths and _chunk_stats["succeeded"] == 0: + print( + f"[graphify extract] error: all semantic chunks failed " + f"for backend '{backend}' ({len(uncached_paths)} uncached files) - " + f"see per-chunk errors above. If you see 'requires the X package', " + f"run `pip install X` and retry.", + file=sys.stderr, + ) + sys.exit(1) + try: + _save_semantic_cache( + fresh.get("nodes", []), + fresh.get("edges", []), + fresh.get("hyperedges", []), + root=out_root, + ) + except Exception as exc: + print(f"[graphify extract] warning: could not write semantic cache: {exc}", file=sys.stderr) + res["nodes"].extend(fresh.get("nodes", [])) + res["edges"].extend(fresh.get("edges", [])) + res["hyperedges"].extend(fresh.get("hyperedges", [])) + res["input_tokens"] += fresh.get("input_tokens", 0) + res["output_tokens"] += fresh.get("output_tokens", 0) + return res, sem_cache_hits, sem_cache_misses + + def run_pg(): + res = {"nodes": [], "edges": []} + if cli_postgres_dsn is not None: + from graphify.pg_introspect import introspect_postgres + print(f"[graphify extract] introspecting PostgreSQL schema...") + try: + res = introspect_postgres(cli_postgres_dsn) + except (ConnectionError, ImportError) as exc: + print(f"error: {exc}", file=sys.stderr) + sys.exit(1) + print(f"[graphify extract] PostgreSQL: {len(res['nodes'])} nodes, {len(res['edges'])} edges") + return res + + def run_cargo(): + res = {"nodes": [], "edges": []} + if cli_cargo: + from graphify.cargo_introspect import introspect_cargo + print("[graphify extract] introspecting Cargo workspace...") + try: + res = introspect_cargo(target) + except (ConnectionError, ImportError, OSError) as exc: + print(f"error: {exc}", file=sys.stderr) + sys.exit(1) + print(f"[graphify extract] Cargo: {len(res['nodes'])} nodes, {len(res['edges'])} edges") + return res + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + fut_ast = executor.submit(run_ast) + fut_sem = executor.submit(run_semantic) + fut_pg = executor.submit(run_pg) + fut_cargo = executor.submit(run_cargo) + + ast_result = fut_ast.result() + stages.mark("AST extract") + + sem_result, sem_cache_hits, sem_cache_misses = fut_sem.result() try: - pg_result = introspect_postgres(cli_postgres_dsn) - except (ConnectionError, ImportError) as exc: - print(f"error: {exc}", file=sys.stderr) - sys.exit(1) - print(f"[graphify extract] PostgreSQL: {len(pg_result['nodes'])} nodes, " - f"{len(pg_result['edges'])} edges") + from graphify.cache import prune_semantic_cache as _prune_semantic_cache, file_hash as _file_hash + _live_hashes: set[str] = set() + for _kind in ("document", "paper", "image"): + for _fp in files_by_type.get(_kind, []): + _abs = Path(_fp) + if not _abs.is_absolute(): + _abs = Path(out_root) / _abs + if not _abs.is_file(): + continue + try: + _live_hashes.add(_file_hash(_abs, out_root)) + except OSError: + pass + _prune_semantic_cache(out_root, _live_hashes) + except Exception as exc: + print(f"[graphify extract] warning: could not prune semantic cache: {exc}", file=sys.stderr) + stages.mark("semantic extract") - cargo_result: dict = {"nodes": [], "edges": []} - if cli_cargo: - from graphify.cargo_introspect import introspect_cargo - print("[graphify extract] introspecting Cargo workspace...") - try: - cargo_result = introspect_cargo(target) - except (ConnectionError, ImportError, OSError) as exc: - print(f"error: {exc}", file=sys.stderr) - sys.exit(1) - print(f"[graphify extract] Cargo: {len(cargo_result['nodes'])} nodes, " - f"{len(cargo_result['edges'])} edges") + pg_result = fut_pg.result() + cargo_result = fut_cargo.result() # Merge AST + semantic + pg_result + cargo_result. Order matters for deduplication: passing AST # first means semantic node attributes win on collision (richer labels diff --git a/tests/test_parallel_orchestration.py b/tests/test_parallel_orchestration.py new file mode 100644 index 000000000..e943efef1 --- /dev/null +++ b/tests/test_parallel_orchestration.py @@ -0,0 +1,169 @@ +import time +import pytest +import json +import graphify.__main__ as mainmod + +def test_pipeline_runs_in_parallel(monkeypatch, tmp_path): + """ + Verify that AST and Semantic extraction run concurrently. + If they run sequentially, total time will be >= 1.0s. + If they run in parallel, total time should be ~0.5s. + """ + # 1. Setup a dummy project with one code file and one markdown file + (tmp_path / "main.py").write_text("def test(): pass") + (tmp_path / "doc.md").write_text("# Documentation") + out_dir = tmp_path / "out" + + # Bypass API key checks + monkeypatch.setenv("GEMINI_API_KEY", "fake") + monkeypatch.setattr(mainmod, "_check_skill_version", lambda _: None) + + # 2. Mock the AST extractor to take 0.5 seconds + def mock_ast(*args, **kwargs): + time.sleep(0.5) + return {"nodes": [{"id": "ast_node"}], "edges": [], "input_tokens": 0, "output_tokens": 0} + + # 3. Mock the Semantic extractor to take 0.5 seconds + def mock_semantic(*args, **kwargs): + time.sleep(0.5) + if "on_chunk_done" in kwargs: + kwargs["on_chunk_done"](0, 1, {}) + return {"nodes": [{"id": "sem_node"}], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0} + + # Apply the mocks + import graphify.extract + monkeypatch.setattr(graphify.extract, "extract", mock_ast) + monkeypatch.setattr("graphify.llm.extract_corpus_parallel", mock_semantic) + + # Setup the CLI arguments + monkeypatch.setattr(mainmod.sys, "argv", [ + "graphify", "extract", str(tmp_path), "--out", str(out_dir), "--no-cluster" + ]) + + # 4. Time the execution + start_time = time.time() + try: + mainmod.main() + except SystemExit: + pass + elapsed = time.time() - start_time + + # 5. Assert Parallel Execution Time + assert elapsed < 1.2, f"Tasks ran sequentially! Took {elapsed:.2f}s instead of parallelized ~0.9s" + + # 6. Assert Data Integrity (both nodes must exist in the final graph) + graph_path = out_dir / "graphify-out" / "graph.json" + graph_data = json.loads(graph_path.read_text()) + node_ids = {n["id"] for n in graph_data["nodes"]} + + assert "ast_node" in node_ids, "AST node was lost during parallel merge!" + assert "sem_node" in node_ids, "Semantic node was lost during parallel merge!" + + +def test_parallel_failure_propagates(monkeypatch, tmp_path): + """ + Verify that if the semantic extractor fails completely (all chunks fail), + the CLI exits with code 1 and doesn't silently generate an incomplete graph. + """ + (tmp_path / "main.py").write_text("def test(): pass") + (tmp_path / "doc.md").write_text("# Documentation") + out_dir = tmp_path / "out" + + monkeypatch.setenv("GEMINI_API_KEY", "fake") + monkeypatch.setattr(mainmod, "_check_skill_version", lambda _: None) + + import graphify.extract + monkeypatch.setattr(graphify.extract, "extract", lambda *a, **k: {"nodes": [], "edges": []}) + + # Do NOT call on_chunk_done to simulate total failure + def mock_semantic_fail(*args, **kwargs): + return {"nodes": [], "edges": [], "hyperedges": []} + + monkeypatch.setattr("graphify.llm.extract_corpus_parallel", mock_semantic_fail) + + monkeypatch.setattr(mainmod.sys, "argv", [ + "graphify", "extract", str(tmp_path), "--out", str(out_dir), "--no-cluster" + ]) + + with pytest.raises(SystemExit) as exc_info: + mainmod.main() + + assert exc_info.value.code == 1, "CLI should exit 1 when semantic thread fails completely" + + +def test_parallel_code_only_corpus(monkeypatch, tmp_path): + """ + Verify that when there are no documents, the parallel runner + doesn't crash on the empty semantic thread. + """ + # Only code files + (tmp_path / "main.py").write_text("def test(): pass") + out_dir = tmp_path / "out" + + monkeypatch.setenv("GEMINI_API_KEY", "fake") + monkeypatch.setattr(mainmod, "_check_skill_version", lambda _: None) + + import graphify.extract + monkeypatch.setattr(graphify.extract, "extract", lambda *a, **k: {"nodes": [{"id": "ast_only"}], "edges": []}) + + # We don't even need to mock semantic extraction since it shouldn't be called + + monkeypatch.setattr(mainmod.sys, "argv", [ + "graphify", "extract", str(tmp_path), "--out", str(out_dir), "--no-cluster" + ]) + + try: + mainmod.main() + except SystemExit as exc: + assert exc.code in (None, 0) + + graph_path = out_dir / "graphify-out" / "graph.json" + graph_data = json.loads(graph_path.read_text()) + node_ids = {n["id"] for n in graph_data["nodes"]} + + assert "ast_only" in node_ids + assert len(node_ids) == 1 + + +def test_postgres_and_cargo_merge(monkeypatch, tmp_path): + """ + Verify that PostgreSQL and Cargo nodes are correctly merged + by the thread pool orchestrator. + """ + (tmp_path / "main.py").write_text("def test(): pass") + out_dir = tmp_path / "out" + + monkeypatch.setenv("GEMINI_API_KEY", "fake") + monkeypatch.setattr(mainmod, "_check_skill_version", lambda _: None) + + import graphify.extract + monkeypatch.setattr(graphify.extract, "extract", lambda *a, **k: {"nodes": [], "edges": []}) + + def mock_pg(*args, **kwargs): + return {"nodes": [{"id": "pg_node"}], "edges": []} + + def mock_cargo(*args, **kwargs): + return {"nodes": [{"id": "cargo_node"}], "edges": []} + + import graphify.pg_introspect + monkeypatch.setattr(graphify.pg_introspect, "introspect_postgres", mock_pg) + + import graphify.cargo_introspect + monkeypatch.setattr(graphify.cargo_introspect, "introspect_cargo", mock_cargo) + + monkeypatch.setattr(mainmod.sys, "argv", [ + "graphify", "extract", str(tmp_path), "--out", str(out_dir), + "--no-cluster", "--postgres", "postgresql://user:pass@localhost/db", "--cargo" + ]) + + try: + mainmod.main() + except SystemExit as exc: + assert exc.code in (None, 0) + + graph_path = out_dir / "graphify-out" / "graph.json" + graph_data = json.loads(graph_path.read_text()) + node_ids = {n["id"] for n in graph_data["nodes"]} + + assert "pg_node" in node_ids + assert "cargo_node" in node_ids