From 69168842279e4fe44d8ed3efa0bed56a5cfb203c Mon Sep 17 00:00:00 2001 From: Rahul Krishna Date: Fri, 26 Jun 2026 23:37:15 -0400 Subject: [PATCH 1/5] feat(pycg): coupling-aware shard planning from the Jedi module graph Sharding lets PyCG (level 2) scale past its ~500-file ceiling by analysing the project in independent pieces. The existing scheme shards one-per-package with a flat file-count ceiling, which is blind to call coupling: it severs heavily-interacting modules (their cross-shard edges become ghost nodes PyCG never resolves) and drops oversized packages wholesale. Add a coupling-aware planner that partitions the module-dependency graph *derived from the Jedi call graph already computed at level 1*: 1. project Jedi callable->callable edges onto a weighted module DiGraph; 2. condense strongly-connected components (import cycles become atomic and are never split across shards); 3. cluster with Louvain so tightly-coupled modules co-compute; 4. enforce the per-shard file budget (re-partition oversized communities, then merge/first-fit-pack the remainder to recover edges and cut count). The reported cut_ratio (fraction of Jedi edge weight crossing shard boundaries) is an upper bound on PyCG edges lost to sharding; on a synthetic worst case it drops from 0.55 (per-package) to 0.03. Wire it into PyCG behind --pycg-shard-strategy {jedi,package} (default jedi). Because planner shards are arbitrary file sets rather than directories, each runs through a temporary symlink mini-project (_shard_symlink_root) so PyCG's own package-root bound confines analysis to the shard and emits project-relative edge names with no prefix rewrite. Thread the level-1 Jedi edges through core -> _get_pycg_call_graph -> build_call_graph_edges to feed the planner. Ray parallelism falls back to sequential under the jedi strategy for now. Add test/test_shard_planner.py (graph projection, SCC atomicity, budget, single-assignment, cut-ratio vs naive, determinism). --- codeanalyzer/__main__.py | 17 +- codeanalyzer/core.py | 13 +- codeanalyzer/options/__init__.py | 4 +- codeanalyzer/options/options.py | 15 + .../semantic_analysis/pycg/pycg_analysis.py | 189 ++++++++- .../semantic_analysis/pycg/shard_planner.py | 394 ++++++++++++++++++ test/test_shard_planner.py | 156 +++++++ 7 files changed, 774 insertions(+), 14 deletions(-) create mode 100644 codeanalyzer/semantic_analysis/pycg/shard_planner.py create mode 100644 test/test_shard_planner.py diff --git a/codeanalyzer/__main__.py b/codeanalyzer/__main__.py index 6693cba..8c6bf72 100644 --- a/codeanalyzer/__main__.py +++ b/codeanalyzer/__main__.py @@ -7,7 +7,7 @@ from codeanalyzer.utils import _set_log_level, logger from codeanalyzer.config import OutputFormat from codeanalyzer.schema import model_dump_json -from codeanalyzer.options import AnalysisOptions, EmitTarget +from codeanalyzer.options import AnalysisOptions, EmitTarget, ShardStrategy def main( @@ -186,6 +186,20 @@ def main( min=0, ), ] = 120, + pycg_shard_strategy: Annotated[ + ShardStrategy, + typer.Option( + "--pycg-shard-strategy", + help=( + "How --pycg-shard groups files (level 2 only). 'jedi' (default) " + "partitions the Jedi module-dependency graph (SCC + Louvain) so " + "tightly-coupled modules co-compute and few call edges are " + "severed between shards; import cycles are never split. " + "'package' uses the legacy one-shard-per-package-directory " + "grouping." + ), + ), + ] = ShardStrategy.JEDI, ): options = AnalysisOptions( input=input, @@ -209,6 +223,7 @@ def main( pycg_shard=pycg_shard, pycg_shard_ceiling=pycg_shard_ceiling, pycg_shard_timeout=pycg_shard_timeout, + pycg_shard_strategy=pycg_shard_strategy, ) _set_log_level(options.verbosity) diff --git a/codeanalyzer/core.py b/codeanalyzer/core.py index 5fc4003..fc87f95 100644 --- a/codeanalyzer/core.py +++ b/codeanalyzer/core.py @@ -433,8 +433,9 @@ def analyze(self) -> PyApplication: logger.info("✅ Jedi: %d edges in %.1fs", len(call_graph), time.perf_counter() - t0_jedi) if self.analysis_level >= 2: - # Level 2: also add PyCG edges. - pycg_edges = self._get_pycg_call_graph(symbol_table) + # Level 2: also add PyCG edges. The Jedi edges double as the + # coupling graph that drives coupling-aware PyCG sharding. + pycg_edges = self._get_pycg_call_graph(symbol_table, jedi_edges) call_graph = merge_edges(call_graph, pycg_edges) call_graph = filter_external_edges(call_graph, symbol_table) @@ -661,6 +662,7 @@ def _build_symbol_table(self, cached_symbol_table: Optional[Dict[str, PyModule]] def _get_pycg_call_graph( self, symbol_table: Dict[str, PyModule], + jedi_edges: List[PyCallEdge], ) -> List[PyCallEdge]: """Build PyCG-resolved call edges. @@ -668,6 +670,10 @@ def _get_pycg_call_graph( and returns edges with ``provenance=["pycg"]``. Falls back to an empty list and logs a warning on any failure so the caller can continue with Jedi-only edges. + + *jedi_edges* are the level-1 call edges; under the ``jedi`` shard + strategy they drive coupling-aware partitioning (see + :func:`shard_planner.plan_shards`). """ try: pycg = PyCG( @@ -676,9 +682,10 @@ def _get_pycg_call_graph( shard=self.options.pycg_shard, shard_ceiling=self.options.pycg_shard_ceiling, shard_timeout=self.options.pycg_shard_timeout, + shard_strategy=self.options.pycg_shard_strategy, using_ray=self.using_ray, ) - return pycg.build_call_graph_edges(symbol_table) + return pycg.build_call_graph_edges(symbol_table, jedi_edges=jedi_edges) except PyCGExceptions.PyCGImportError as exc: logger.warning(f"PyCG not installed — level 2 edges will be Jedi-only: {exc}") return [] diff --git a/codeanalyzer/options/__init__.py b/codeanalyzer/options/__init__.py index 127a183..317d42b 100644 --- a/codeanalyzer/options/__init__.py +++ b/codeanalyzer/options/__init__.py @@ -1,3 +1,3 @@ -from .options import AnalysisOptions, EmitTarget, OutputFormat +from .options import AnalysisOptions, EmitTarget, OutputFormat, ShardStrategy -__all__ = ["AnalysisOptions", "EmitTarget", "OutputFormat"] \ No newline at end of file +__all__ = ["AnalysisOptions", "EmitTarget", "OutputFormat", "ShardStrategy"] \ No newline at end of file diff --git a/codeanalyzer/options/options.py b/codeanalyzer/options/options.py index 41eb404..5e3b56b 100644 --- a/codeanalyzer/options/options.py +++ b/codeanalyzer/options/options.py @@ -23,6 +23,20 @@ class EmitTarget(str, Enum): SCHEMA = "schema" +class ShardStrategy(str, Enum): + """How ``--pycg-shard`` groups files into shards (level 2 only). + + - ``jedi`` : partition the Jedi module-dependency graph (strongly- + connected-component condensation + Louvain) so tightly- + coupled modules co-compute and few call edges are severed + between shards. Import cycles are never split. + - ``package`` : legacy one-shard-per-package-directory grouping. + """ + + JEDI = "jedi" + PACKAGE = "package" + + @dataclass class AnalysisOptions: input: Path @@ -46,3 +60,4 @@ class AnalysisOptions: pycg_shard: bool = False pycg_shard_ceiling: int = 100 pycg_shard_timeout: int = 120 + pycg_shard_strategy: ShardStrategy = ShardStrategy.JEDI diff --git a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py index 9a8ece9..bb69a03 100644 --- a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py +++ b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py @@ -43,12 +43,14 @@ import importlib.util # noqa: F401 import contextlib import json # noqa: F401 +import shutil import signal +import tempfile import time from collections import Counter, defaultdict from pathlib import Path -from typing import Any, Dict, Generator, List, Optional, Set, Union +from typing import Any, Dict, Generator, List, Optional, Set, Tuple, Union @contextlib.contextmanager @@ -79,9 +81,64 @@ def _handler(signum: int, frame: object) -> None: from codeanalyzer.schema.py_schema import PyCallEdge, PyModule from codeanalyzer.semantic_analysis.call_graph import iter_callables_in_symbol_table from codeanalyzer.semantic_analysis.pycg.pycg_exceptions import PyCGExceptions +from codeanalyzer.semantic_analysis.pycg.shard_planner import plan_shards from codeanalyzer.utils import ProgressBar, logger +@contextlib.contextmanager +def _shard_symlink_root( + files: List[str], + project_dir: Path, +) -> Generator[Tuple[Path, List[str]], None, None]: + """Materialise a shard's files as a temporary mini-project. + + PyCG bounds its import-following to the ``package`` directory — only + modules whose resolved file lives under that root are followed; everything + else becomes a ghost node (``ImportManager``: ``if self.mod_dir not in + mod.__file__: return``). A coupling-derived shard is an arbitrary set of + files that need not form a directory, so we mirror the project layout into + a temp dir holding symlinks to exactly the shard's files plus the + ``__init__.py`` chain each needs for package resolution. Running PyCG with + this mirror as the package root confines analysis to the shard while + emitting project-relative edge names (so ``prefix=""`` — no rename needed). + + Yields ``(root, entry_points)`` where *entry_points* are the symlinked + paths inside *root*. The temp tree is removed on exit. + """ + root = Path(tempfile.mkdtemp(prefix="canpy_pycg_shard_")) + entry_points: List[str] = [] + linked_inits: Set[Path] = set() + try: + for f in files: + src = Path(f).resolve() + try: + rel = src.relative_to(project_dir) + except ValueError: + continue # defensively skip files outside the project + dst = root / rel + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(src) + entry_points.append(str(dst)) + + # Symlink the __init__.py chain from project root down to this + # file's package so PyCG/importlib can resolve the dotted module + # name. These add ~0 analysis cost (usually empty) and keep + # out-of-shard siblings unresolved → ghost nodes. + for i in range(len(rel.parent.parts) + 1): + pkg_rel = Path(*rel.parent.parts[:i]) + real_init = project_dir / pkg_rel / "__init__.py" + link_init = root / pkg_rel / "__init__.py" + if real_init.exists() and link_init not in linked_inits: + link_init.parent.mkdir(parents=True, exist_ok=True) + if not link_init.exists(): + link_init.symlink_to(real_init.resolve()) + linked_inits.add(link_init) + yield root, entry_points + finally: + shutil.rmtree(root, ignore_errors=True) + + def _pycg_shard_worker( entry_points: List[str], package_dir: str, @@ -313,6 +370,7 @@ def __init__( shard: bool = False, shard_ceiling: Optional[int] = None, shard_timeout: Optional[int] = None, + shard_strategy: str = "jedi", using_ray: bool = False, ) -> None: self.project_dir = Path(project_dir).resolve() @@ -324,9 +382,31 @@ def __init__( self.shard_timeout = ( shard_timeout if shard_timeout is not None else self._PYCG_SHARD_TIMEOUT ) + # "jedi": partition the Jedi module graph (SCC + Louvain) so coupled + # modules co-compute and few edges are severed (see shard_planner). + # "package": legacy one-shard-per-package-directory grouping. + self.shard_strategy = shard_strategy self.using_ray = using_ray self._CallGraphGenerator: Optional[Any] = None + @staticmethod + def _coalesce_edges(edges: List[PyCallEdge]) -> List[PyCallEdge]: + """Sum weights of duplicate ``(source, target)`` pairs across shards.""" + merged: Dict[tuple, PyCallEdge] = {} + for edge in edges: + key = (edge.source, edge.target) + if key in merged: + existing = merged[key] + merged[key] = PyCallEdge( + source=existing.source, + target=existing.target, + weight=existing.weight + edge.weight, + provenance=existing.provenance, + ) + else: + merged[key] = edge + return list(merged.values()) + # ------------------------------------------------------------------ # Entry-point collection # ------------------------------------------------------------------ @@ -455,6 +535,88 @@ def _run_pycg_batch( # Sharded analysis # ------------------------------------------------------------------ + def _build_sharded_planned( + self, + jedi_edges: List[PyCallEdge], + symbol_table: Dict[str, PyModule], + resolver: "_PyCGCallableResolver", + ) -> List[PyCallEdge]: + """Coupling-aware sharding driven by the Jedi module graph. + + Unlike :meth:`_build_sharded` (one shard per package directory), the + shards here are chosen to *minimise the call edges severed between + shards*: :func:`shard_planner.plan_shards` condenses the Jedi call + graph by strongly-connected component (so import cycles never split) + and clusters it with Louvain so tightly-coupled modules land together. + Each shard — an arbitrary set of files — is run through PyCG via a + symlinked mini-project (:func:`_shard_symlink_root`) that bounds PyCG + to exactly those files. + + Reported ``cut_ratio`` is the fraction of Jedi edge weight crossing + shard boundaries — an upper bound on the PyCG edges lost to sharding. + """ + plan = plan_shards( + symbol_table, jedi_edges, budget=self.shard_ceiling, merge_small=True + ) + m = plan.metrics + logger.info( + "PyCG: planned %d shard(s) from Jedi module graph " + "(cut_ratio=%.3f, max_shard=%d files, %d modules)", + int(m["num_shards"]), m["cut_ratio"], + int(m["max_shard_files"]), int(m["modules"]), + ) + if m["oversized_shards"]: + logger.warning( + "PyCG: %d shard(s) exceed the %d-file ceiling — skipped " + "(atomic import cycles larger than the budget)", + int(m["oversized_shards"]), self.shard_ceiling, + ) + + if self.using_ray: + logger.info( + "PyCG: Ray parallelism is not yet wired for the 'jedi' shard " + "strategy — running shards sequentially." + ) + + all_edges: List[PyCallEdge] = [] + skipped = 0 + with ProgressBar(len(plan.shards), "Building call graph shards", item_label="shards") as progress: + for idx, files in enumerate(plan.shards): + n = len(files) + if n > self.shard_ceiling: + skipped += 1 + progress.advance() + continue + try: + with _shard_symlink_root(files, self.project_dir) as (root, eps): + with _shard_timeout(self.shard_timeout): + edges = self._run_pycg_batch(eps, root, resolver, prefix="") + all_edges.extend(edges) + logger.debug("PyCG shard %d: %d edges from %d files", idx, len(edges), n) + except TimeoutError: + logger.warning( + "PyCG shard %d timed out after %ds — skipped", + idx, self.shard_timeout, + ) + skipped += 1 + except PyCGExceptions.PyCGAnalysisError as exc: + logger.warning("PyCG shard %d failed — skipped: %s", idx, exc) + skipped += 1 + progress.advance() + + if skipped: + logger.warning( + "PyCG: %d/%d shard(s) skipped (ceiling, %ds timeout, or failure)", + skipped, len(plan.shards), self.shard_timeout, + ) + + result = self._coalesce_edges(all_edges) + logger.info( + "PyCG: %d edges from %d/%d shard(s) (%d before dedup, Jedi-planned)", + len(result), len(plan.shards) - skipped, len(plan.shards), len(all_edges), + ) + return result + def _build_sharded( self, entry_points: List[str], @@ -668,7 +830,9 @@ def _build_sharded_ray(self, shards: Dict[Path, List[str]]) -> List[PyCallEdge]: # ------------------------------------------------------------------ def build_call_graph_edges( - self, symbol_table: Dict[str, PyModule] + self, + symbol_table: Dict[str, PyModule], + jedi_edges: Optional[List[PyCallEdge]] = None, ) -> List[PyCallEdge]: """Run PyCG and return ``PyCallEdge`` entries with ``provenance=["pycg"]``. @@ -701,12 +865,21 @@ def build_call_graph_edges( if n_files > self._PYCG_FILE_CEILING: if self.shard: - mode = "Ray-parallel" if self.using_ray else "sequential" - logger.info( - "PyCG: starting sharded call graph analysis (%d files, %s)", - n_files, mode, - ) - edges = self._build_sharded(entry_points, resolver) + if self.shard_strategy == "jedi" and jedi_edges is not None: + logger.info( + "PyCG: starting Jedi-planned sharded analysis (%d files)", + n_files, + ) + edges = self._build_sharded_planned( + jedi_edges, symbol_table, resolver + ) + else: + mode = "Ray-parallel" if self.using_ray else "sequential" + logger.info( + "PyCG: starting per-package sharded analysis (%d files, %s)", + n_files, mode, + ) + edges = self._build_sharded(entry_points, resolver) else: logger.warning( "PyCG: %d entry points exceeds ceiling of %d — " diff --git a/codeanalyzer/semantic_analysis/pycg/shard_planner.py b/codeanalyzer/semantic_analysis/pycg/shard_planner.py new file mode 100644 index 0000000..43bf648 --- /dev/null +++ b/codeanalyzer/semantic_analysis/pycg/shard_planner.py @@ -0,0 +1,394 @@ +################################################################################ +# Copyright IBM Corporation 2025 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ + +"""Coupling-aware shard planning for PyCG call-graph construction. + +Sharding splits a project so PyCG can be run on each part independently; any +call edge whose caller and callee land in *different* shards becomes a ghost +node in both, so PyCG never resolves it. The directory + file-count heuristic +is blind to coupling and can sever heavily-interacting modules. This planner +instead partitions the **module dependency graph derived from the Jedi call +graph** (already computed at analysis level 1), so the cuts fall on the weakest +seams. + +Pipeline: + +1. **Module graph** — project the Jedi ``PyCallEdge`` list (callable→callable) + down to a weighted directed graph over *modules*. ``weight(A, B)`` is the + number of Jedi call sites from a callable in module ``A`` to one in ``B``. +2. **SCC condensation** — modules in an import/call cycle must be co-computed + (splitting them breaks both shards), so each strongly-connected component is + collapsed into one indivisible unit via Tarjan's algorithm. +3. **Community detection** — Louvain modularity over the undirected weighted + condensation groups tightly-coupled units; each community is a candidate + shard. +4. **Budget enforcement** — communities over the per-shard file budget are + re-partitioned (higher Louvain resolution, then a greedy first-fit fallback + that guarantees termination); communities under budget are agglomeratively + merged by inter-shard weight to recover cut edges and reduce shard count. + +The result is a list of shards (each a list of file paths) plus a metrics +report whose headline figure is ``cut_ratio`` — the fraction of inter-module +Jedi edge weight severed by the partition. Lower is better; it is the +estimated upper bound on PyCG edges lost to sharding. +""" + +from __future__ import annotations + +import logging +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Dict, Iterator, List, Optional, Set, Tuple + +import networkx as nx + +from codeanalyzer.schema.py_schema import PyCallable, PyCallEdge, PyClass, PyModule + +logger = logging.getLogger(__name__) + + +# ---------------------------------------------------------------------------- +# Symbol-table walking: callable / class signature -> owning module +# ---------------------------------------------------------------------------- + +def _walk_callable_sigs(c: PyCallable) -> Iterator[str]: + yield c.signature + for inner in c.inner_callables.values(): + yield from _walk_callable_sigs(inner) + for inner_cls in c.inner_classes.values(): + yield from _walk_class_sigs(inner_cls) + + +def _walk_class_sigs(cls: PyClass) -> Iterator[str]: + yield cls.signature + for method in cls.methods.values(): + yield from _walk_callable_sigs(method) + for inner in cls.inner_classes.values(): + yield from _walk_class_sigs(inner) + + +def _signature_to_module(symbol_table: Dict[str, PyModule]) -> Dict[str, str]: + """Map every callable/class signature in the project to its module name. + + Built by walking each module's full nesting tree, so the mapping is exact + rather than relying on longest-prefix matching against module names (which + is ambiguous when one module name is a prefix of another). + """ + sig_to_mod: Dict[str, str] = {} + for module in symbol_table.values(): + for fn in module.functions.values(): + for sig in _walk_callable_sigs(fn): + sig_to_mod[sig] = module.module_name + for cls in module.classes.values(): + for sig in _walk_class_sigs(cls): + sig_to_mod[sig] = module.module_name + return sig_to_mod + + +# ---------------------------------------------------------------------------- +# Result type +# ---------------------------------------------------------------------------- + +@dataclass +class ShardPlan: + """Output of :func:`plan_shards`. + + ``shards`` is what the PyCG executor consumes: each inner list is the set of + project file paths to analyse together as one shard. ``module_shards`` is + the same partition expressed in module names (handy for logging/tests). + """ + + shards: List[List[str]] = field(default_factory=list) + module_shards: List[List[str]] = field(default_factory=list) + metrics: Dict[str, float] = field(default_factory=dict) + + def __str__(self) -> str: + m = self.metrics + return ( + f"ShardPlan({len(self.shards)} shards, " + f"cut_ratio={m.get('cut_ratio', 0):.3f}, " + f"max_shard_files={int(m.get('max_shard_files', 0))}, " + f"oversized={int(m.get('oversized_shards', 0))})" + ) + + +# ---------------------------------------------------------------------------- +# Module dependency graph +# ---------------------------------------------------------------------------- + +def build_module_graph( + symbol_table: Dict[str, PyModule], + jedi_edges: List[PyCallEdge], +) -> nx.DiGraph: + """Project Jedi callable→callable edges onto a weighted module DiGraph. + + Every project module is a node (isolated modules included). Edge weight is + the summed Jedi weight of cross-module call sites; intra-module edges and + edges touching external/library symbols (no symbol-table entry) are + dropped — they cannot influence how the project is partitioned. + """ + sig_to_mod = _signature_to_module(symbol_table) + + g = nx.DiGraph() + for module in symbol_table.values(): + g.add_node(module.module_name, file_path=module.file_path) + + for edge in jedi_edges: + src = sig_to_mod.get(edge.source) + dst = sig_to_mod.get(edge.target) + if src is None or dst is None or src == dst: + continue + if g.has_edge(src, dst): + g[src][dst]["weight"] += edge.weight + else: + g.add_edge(src, dst, weight=edge.weight) + return g + + +# ---------------------------------------------------------------------------- +# Partitioning +# ---------------------------------------------------------------------------- + +def _undirected_weighted(g: nx.DiGraph) -> nx.Graph: + """Collapse a directed weighted graph to undirected, summing both directions.""" + h = nx.Graph() + h.add_nodes_from(g.nodes()) + for u, v, w in g.edges(data="weight", default=1): + if h.has_edge(u, v): + h[u][v]["weight"] += w + else: + h.add_edge(u, v, weight=w) + return h + + +def _communities(h: nx.Graph, resolution: float, seed: int) -> List[Set[str]]: + """Weighted community detection, preferring Louvain with a graceful fallback.""" + if h.number_of_nodes() == 0: + return [] + if hasattr(nx.community, "louvain_communities"): + return [ + set(c) + for c in nx.community.louvain_communities( + h, weight="weight", resolution=resolution, seed=seed + ) + ] + # networkx < 3.0 (Python 3.9/3.10 floor): greedy modularity has no seed. + return [set(c) for c in nx.community.greedy_modularity_communities(h, weight="weight")] + + +def _greedy_bin_pack( + units: List[Tuple[frozenset, int]], budget: int +) -> List[Set[str]]: + """First-fit-decreasing pack (unit -> modules, size) into <= budget bins. + + Termination guarantee for the recursive splitter: a community that Louvain + refuses to cut still gets divided here. A single unit larger than the + budget (an atomic SCC — a real import cycle) is emitted on its own; that is + unavoidable without breaking edges. + """ + bins: List[Tuple[Set[str], int]] = [] + for members, size in sorted(units, key=lambda u: u[1], reverse=True): + placed = False + for b in bins: + if b[1] + size <= budget: + b[0].update(members) + bins[bins.index(b)] = (b[0], b[1] + size) + placed = True + break + if not placed: + bins.append((set(members), size)) + return [b[0] for b in bins] + + +def _split_oversized( + h: nx.Graph, + community: Set[str], + unit_size: Dict[str, int], + budget: int, + seed: int, + depth: int = 0, +) -> List[Set[str]]: + """Recursively partition a community whose file count exceeds the budget.""" + size = sum(unit_size[n] for n in community) + if size <= budget or len(community) == 1: + # Fits, or is a single atomic SCC unit that cannot be split further. + return [community] + + sub = h.subgraph(community) + # Escalate resolution so Louvain cuts more aggressively each level. + parts = _communities(sub, resolution=2.0 * (depth + 1), seed=seed) + parts = [p for p in parts if p] + + if len(parts) <= 1: + # Louvain could not divide it (one dense blob) — fall back to bin packing + # so the budget is still honoured. + units = [(frozenset([n]), unit_size[n]) for n in community] + return _greedy_bin_pack(units, budget) + + out: List[Set[str]] = [] + for p in parts: + out.extend(_split_oversized(h, p, unit_size, budget, seed, depth + 1)) + return out + + +def _merge_small( + shards: List[Set[str]], + h: nx.Graph, + unit_size: Dict[str, int], + budget: int, +) -> List[Set[str]]: + """Agglomeratively merge under-budget shards by inter-shard weight. + + Recovers cut edges: merging two shards turns the edges between them back + into intra-shard edges. Greedy — repeatedly merge the heaviest-coupled + pair that still fits the budget — which is enough to mop up the small + fragments Louvain leaves behind without reintroducing oversized shards. + """ + shards = [set(s) for s in shards] + sizes = [sum(unit_size[n] for n in s) for s in shards] + + def pair_weight(a: Set[str], b: Set[str]) -> float: + w = 0.0 + for u in a: + for v in h.adj[u]: + if v in b: + w += h[u][v]["weight"] + return w + + while True: + best: Optional[Tuple[int, int]] = None + best_w = 0.0 + for i in range(len(shards)): + for j in range(i + 1, len(shards)): + if sizes[i] + sizes[j] > budget: + continue + w = pair_weight(shards[i], shards[j]) + if w > best_w: + best_w, best = w, (i, j) + if best is None: + break + i, j = best + shards[i] |= shards[j] + sizes[i] += sizes[j] + del shards[j] + del sizes[j] + + # Final consolidation: first-fit pack any shards that still fit together. + # Merging zero-coupling shards (e.g. isolated leaf modules) costs no cut + # weight and avoids spawning a PyCG process per trivial file. Packing + # never severs an edge, so it is strictly safe; the budget still bounds + # per-shard PyCG divergence risk. + units = [(frozenset(s), sz) for s, sz in zip(shards, sizes)] + return _greedy_bin_pack(units, budget) + + +# ---------------------------------------------------------------------------- +# Public entry point +# ---------------------------------------------------------------------------- + +def plan_shards( + symbol_table: Dict[str, PyModule], + jedi_edges: List[PyCallEdge], + budget: int = 100, + seed: int = 42, + merge_small: bool = True, +) -> ShardPlan: + """Partition a project into coupling-aware PyCG shards. + + Args: + symbol_table: The level-1 symbol table (``file_path -> PyModule``). + jedi_edges: Jedi-provenance call edges from the level-1 call graph. + budget: Maximum number of files per shard. + seed: Determinism seed for Louvain. + merge_small: Agglomeratively merge under-budget shards to cut shard + count and recover edges. + + Returns: + A :class:`ShardPlan`. Shards never silently drop files: every project + module appears in exactly one shard (an atomic SCC larger than *budget* + yields a single oversized shard, flagged in ``metrics``). + """ + g = build_module_graph(symbol_table, jedi_edges) + total_weight = float(sum(w for _, _, w in g.edges(data="weight", default=1))) + + # SCC condensation: each strongly-connected component is one atomic unit. + condensation = nx.condensation(g) # DAG; node attr 'members' = set of modules + mapping: Dict[str, int] = condensation.graph["mapping"] # module -> scc id + unit_members: Dict[int, Set[str]] = { + scc: set(condensation.nodes[scc]["members"]) for scc in condensation.nodes + } + unit_size: Dict[int, int] = {scc: len(m) for scc, m in unit_members.items()} + + # Weighted undirected graph over SCC units. + hu = nx.Graph() + hu.add_nodes_from(condensation.nodes()) + for u, v, w in g.edges(data="weight", default=1): + su, sv = mapping[u], mapping[v] + if su == sv: + continue + if hu.has_edge(su, sv): + hu[su][sv]["weight"] += w + else: + hu.add_edge(su, sv, weight=w) + + # Community detection over units, then budget enforcement. + communities = _communities(hu, resolution=1.0, seed=seed) + unit_shards: List[Set[str]] = [] # sets of SCC ids + for community in communities: + unit_shards.extend( + _split_oversized(hu, community, unit_size, budget, seed) + ) + + if merge_small: + unit_shards = _merge_small(unit_shards, hu, unit_size, budget) + + # Expand SCC units back to modules, then to file paths. + module_shards: List[List[str]] = [] + for units in unit_shards: + mods: Set[str] = set() + for scc in units: + mods |= unit_members[scc] + if mods: + module_shards.append(sorted(mods)) + + mod_to_file = {n: g.nodes[n]["file_path"] for n in g.nodes} + file_shards = [[mod_to_file[m] for m in mods] for mods in module_shards] + + # Metrics: how much Jedi edge weight does this partition sever? + shard_of: Dict[str, int] = {} + for idx, mods in enumerate(module_shards): + for m in mods: + shard_of[m] = idx + cut_weight = 0.0 + for u, v, w in g.edges(data="weight", default=1): + if shard_of.get(u) != shard_of.get(v): + cut_weight += w + + sizes = [len(s) for s in module_shards] or [0] + metrics = { + "modules": float(g.number_of_nodes()), + "module_edges": float(g.number_of_edges()), + "total_edge_weight": total_weight, + "cut_weight": cut_weight, + "cut_ratio": (cut_weight / total_weight) if total_weight else 0.0, + "num_shards": float(len(module_shards)), + "max_shard_files": float(max(sizes)), + "oversized_shards": float(sum(1 for s in sizes if s > budget)), + } + + plan = ShardPlan(shards=file_shards, module_shards=module_shards, metrics=metrics) + logger.info("Shard planner: %s", plan) + return plan diff --git a/test/test_shard_planner.py b/test/test_shard_planner.py new file mode 100644 index 0000000..aff3e58 --- /dev/null +++ b/test/test_shard_planner.py @@ -0,0 +1,156 @@ +"""Unit tests for the Jedi-driven PyCG shard planner. + +These exercise the pure partitioning logic (no PyCG required): module-graph +construction, SCC atomicity, budget enforcement, and the cut-ratio metric. +""" +from typing import List + +import networkx as nx +import pytest + +from codeanalyzer.schema.py_schema import PyCallEdge, PyCallable, PyClass, PyModule +from codeanalyzer.semantic_analysis.pycg.shard_planner import ( + build_module_graph, + plan_shards, +) + + +# ---------------------------------------------------------------------------- +# Builders +# ---------------------------------------------------------------------------- + +def _module(name: str, file_path: str, func_names: List[str]) -> PyModule: + fns = { + fn: PyCallable(signature=f"{name}.{fn}", name=fn, path=file_path) + for fn in func_names + } + return PyModule(file_path=file_path, module_name=name, functions=fns) + + +def _edge(src: str, dst: str, w: int = 1) -> PyCallEdge: + return PyCallEdge(source=src, target=dst, weight=w, provenance=["jedi"]) + + +def _cut_ratio(g: nx.DiGraph, module_shards: List[List[str]]) -> float: + shard_of = {m: i for i, mods in enumerate(module_shards) for m in mods} + total = cut = 0.0 + for u, v, w in g.edges(data="weight", default=1): + total += w + if shard_of.get(u) != shard_of.get(v): + cut += w + return cut / total if total else 0.0 + + +# ---------------------------------------------------------------------------- +# build_module_graph +# ---------------------------------------------------------------------------- + +def test_module_graph_projects_callables_to_modules(): + st = { + "/a.py": _module("a", "/a.py", ["f", "g"]), + "/b.py": _module("b", "/b.py", ["h"]), + } + edges = [ + _edge("a.f", "b.h", 3), # cross-module -> kept + _edge("a.f", "a.g", 5), # intra-module -> dropped + _edge("a.g", "ext.lib.x"), # external target -> dropped + ] + g = build_module_graph(st, edges) + assert set(g.nodes) == {"a", "b"} + assert g.has_edge("a", "b") and g["a"]["b"]["weight"] == 3 + assert not g.has_edge("a", "a") + assert g.number_of_edges() == 1 + + +def test_isolated_modules_are_nodes(): + st = {"/a.py": _module("a", "/a.py", ["f"]), "/b.py": _module("b", "/b.py", ["g"])} + g = build_module_graph(st, []) # no edges + assert set(g.nodes) == {"a", "b"} + assert g.number_of_edges() == 0 + + +# ---------------------------------------------------------------------------- +# plan_shards +# ---------------------------------------------------------------------------- + +def _coupled_clusters_project(): + """pkg_a <-> pkg_b heavy cross-package coupling + isolated leaves.""" + st, edges = {}, [] + for i in range(4): + st[f"/pkg_a/m{i}.py"] = _module(f"pkg_a.m{i}", f"/pkg_a/m{i}.py", ["f"]) + st[f"/pkg_b/m{i}.py"] = _module(f"pkg_b.m{i}", f"/pkg_b/m{i}.py", ["g"]) + for i in range(4): + edges.append(_edge(f"pkg_a.m{i}.f", f"pkg_b.m{i}.g", 10)) + edges.append(_edge(f"pkg_b.m{i}.g", f"pkg_a.m{i}.f", 8)) + for i in range(3): + edges.append(_edge(f"pkg_a.m{i}.f", f"pkg_a.m{i+1}.f", 2)) + return st, edges + + +def test_budget_is_respected(): + st, edges = _coupled_clusters_project() + plan = plan_shards(st, edges, budget=4) + assert plan.metrics["max_shard_files"] <= 4 + assert plan.metrics["oversized_shards"] == 0 + + +def test_every_module_assigned_exactly_once(): + st, edges = _coupled_clusters_project() + plan = plan_shards(st, edges, budget=4) + assigned = [m for shard in plan.module_shards for m in shard] + assert sorted(assigned) == sorted(m.module_name for m in st.values()) + assert len(assigned) == len(set(assigned)) # no module duplicated + + +def test_beats_naive_per_package_cut_ratio(): + st, edges = _coupled_clusters_project() + g = build_module_graph(st, edges) + naive = {} + for m in st.values(): + naive.setdefault(m.module_name.split(".")[0], []).append(m.module_name) + naive_ratio = _cut_ratio(g, list(naive.values())) + + plan = plan_shards(st, edges, budget=4) + assert plan.metrics["cut_ratio"] < naive_ratio + + +def test_import_cycle_is_never_split(): + # util.x <-> helpers.y form a cross-package cycle; must co-locate even + # though they live in different top-level packages. + st = { + "/util/x.py": _module("util.x", "/util/x.py", ["a"]), + "/helpers/y.py": _module("helpers.y", "/helpers/y.py", ["b"]), + } + edges = [_edge("util.x.a", "helpers.y.b", 5), _edge("helpers.y.b", "util.x.a", 5)] + plan = plan_shards(st, edges, budget=10) + shard_with_util = next(s for s in plan.module_shards if "util.x" in s) + assert "helpers.y" in shard_with_util + + +def test_oversized_atomic_cycle_is_flagged_not_dropped(): + # A single import cycle of 6 modules with a budget of 3 cannot be split + # without breaking edges; it must survive as one oversized shard. + st, edges = {}, [] + names = [f"cyc.m{i}" for i in range(6)] + for i, n in enumerate(names): + st[f"/cyc/m{i}.py"] = _module(n, f"/cyc/m{i}.py", ["f"]) + for i in range(6): # ring -> one big SCC + edges.append(_edge(f"{names[i]}.f", f"{names[(i + 1) % 6]}.f", 1)) + plan = plan_shards(st, edges, budget=3) + assert plan.metrics["oversized_shards"] >= 1 + assigned = [m for shard in plan.module_shards for m in shard] + assert sorted(assigned) == sorted(names) # nothing dropped + + +def test_determinism(): + st, edges = _coupled_clusters_project() + a = plan_shards(st, edges, budget=4) + b = plan_shards(st, edges, budget=4) + norm = lambda p: sorted(sorted(s) for s in p.module_shards) + assert norm(a) == norm(b) + + +def test_empty_project(): + plan = plan_shards({}, [], budget=4) + assert plan.shards == [] + assert plan.metrics["cut_ratio"] == 0.0 From 5c02ba302933b839cc26427291f24dca13e4f4c7 Mon Sep 17 00:00:00 2001 From: Rahul Krishna Date: Fri, 26 Jun 2026 23:40:00 -0400 Subject: [PATCH 2/5] feat(pycg): Ray-parallel execution for Jedi-planned shards Materialise each planned file-set shard as a symlink mini-project up front (the trees must outlive their remote tasks), submit one Ray task per shard, and collect against a single wall-clock deadline (Ray workers can't use SIGALRM, so the timeout is enforced at the orchestrator, mirroring _build_sharded_ray). Symlink trees are cleaned up once the batch completes. Factor _materialize_shard_root out of the _shard_symlink_root context manager so both the sequential and Ray paths share tree construction. Under --ray the jedi strategy now parallelises instead of falling back to sequential. --- .../semantic_analysis/pycg/pycg_analysis.py | 172 ++++++++++++++---- 1 file changed, 136 insertions(+), 36 deletions(-) diff --git a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py index bb69a03..8beb03e 100644 --- a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py +++ b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py @@ -81,16 +81,15 @@ def _handler(signum: int, frame: object) -> None: from codeanalyzer.schema.py_schema import PyCallEdge, PyModule from codeanalyzer.semantic_analysis.call_graph import iter_callables_in_symbol_table from codeanalyzer.semantic_analysis.pycg.pycg_exceptions import PyCGExceptions -from codeanalyzer.semantic_analysis.pycg.shard_planner import plan_shards +from codeanalyzer.semantic_analysis.pycg.shard_planner import ShardPlan, plan_shards from codeanalyzer.utils import ProgressBar, logger -@contextlib.contextmanager -def _shard_symlink_root( +def _materialize_shard_root( files: List[str], project_dir: Path, -) -> Generator[Tuple[Path, List[str]], None, None]: - """Materialise a shard's files as a temporary mini-project. +) -> Tuple[Path, List[str]]: + """Build a temporary symlink mini-project for a shard; return ``(root, eps)``. PyCG bounds its import-following to the ``package`` directory — only modules whose resolved file lives under that root are followed; everything @@ -102,38 +101,50 @@ def _shard_symlink_root( this mirror as the package root confines analysis to the shard while emitting project-relative edge names (so ``prefix=""`` — no rename needed). - Yields ``(root, entry_points)`` where *entry_points* are the symlinked - paths inside *root*. The temp tree is removed on exit. + The caller owns the returned *root* and must ``shutil.rmtree`` it. """ root = Path(tempfile.mkdtemp(prefix="canpy_pycg_shard_")) entry_points: List[str] = [] linked_inits: Set[Path] = set() + for f in files: + src = Path(f).resolve() + try: + rel = src.relative_to(project_dir) + except ValueError: + continue # defensively skip files outside the project + dst = root / rel + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(src) + entry_points.append(str(dst)) + + # Symlink the __init__.py chain from project root down to this file's + # package so PyCG/importlib can resolve the dotted module name. These + # add ~0 analysis cost (usually empty) and keep out-of-shard siblings + # unresolved → ghost nodes. + for i in range(len(rel.parent.parts) + 1): + pkg_rel = Path(*rel.parent.parts[:i]) + real_init = project_dir / pkg_rel / "__init__.py" + link_init = root / pkg_rel / "__init__.py" + if real_init.exists() and link_init not in linked_inits: + link_init.parent.mkdir(parents=True, exist_ok=True) + if not link_init.exists(): + link_init.symlink_to(real_init.resolve()) + linked_inits.add(link_init) + return root, entry_points + + +@contextlib.contextmanager +def _shard_symlink_root( + files: List[str], + project_dir: Path, +) -> Generator[Tuple[Path, List[str]], None, None]: + """Context-manager wrapper around :func:`_materialize_shard_root`. + + Yields ``(root, entry_points)`` and removes the temp tree on exit. + """ + root, entry_points = _materialize_shard_root(files, project_dir) try: - for f in files: - src = Path(f).resolve() - try: - rel = src.relative_to(project_dir) - except ValueError: - continue # defensively skip files outside the project - dst = root / rel - dst.parent.mkdir(parents=True, exist_ok=True) - if not dst.exists(): - dst.symlink_to(src) - entry_points.append(str(dst)) - - # Symlink the __init__.py chain from project root down to this - # file's package so PyCG/importlib can resolve the dotted module - # name. These add ~0 analysis cost (usually empty) and keep - # out-of-shard siblings unresolved → ghost nodes. - for i in range(len(rel.parent.parts) + 1): - pkg_rel = Path(*rel.parent.parts[:i]) - real_init = project_dir / pkg_rel / "__init__.py" - link_init = root / pkg_rel / "__init__.py" - if real_init.exists() and link_init not in linked_inits: - link_init.parent.mkdir(parents=True, exist_ok=True) - if not link_init.exists(): - link_init.symlink_to(real_init.resolve()) - linked_inits.add(link_init) yield root, entry_points finally: shutil.rmtree(root, ignore_errors=True) @@ -573,10 +584,7 @@ def _build_sharded_planned( ) if self.using_ray: - logger.info( - "PyCG: Ray parallelism is not yet wired for the 'jedi' shard " - "strategy — running shards sequentially." - ) + return self._build_sharded_planned_ray(plan) all_edges: List[PyCallEdge] = [] skipped = 0 @@ -617,6 +625,98 @@ def _build_sharded_planned( ) return result + def _build_sharded_planned_ray(self, plan: "ShardPlan") -> List[PyCallEdge]: + """Ray-parallel execution of Jedi-planned file-set shards. + + Each shard is materialised as a symlink mini-project up front (the + trees must outlive their remote tasks), submitted as a Ray task, and + collected against a single wall-clock deadline — Ray workers cannot use + SIGALRM, so the timeout is enforced at the orchestrator level (mirroring + :meth:`_build_sharded_ray`). All symlink trees are removed once the + batch completes. + """ + import os + import ray + + os.environ.setdefault("RAY_IGNORE_UNHANDLED_ERRORS", "1") + remote_fn = ray.remote(_pycg_shard_worker) + + roots: List[Path] = [] + futures: List[Any] = [] + meta: Dict[Any, tuple] = {} # ObjectRef -> (shard_idx, n_files) + skipped = 0 + all_edges: List[PyCallEdge] = [] + try: + with ProgressBar(len(plan.shards), "Building call graph shards (parallel)", item_label="shards") as progress: + for idx, files in enumerate(plan.shards): + n = len(files) + if n > self.shard_ceiling: + skipped += 1 + progress.advance() + continue + root, eps = _materialize_shard_root(files, self.project_dir) + roots.append(root) + fut = remote_fn.remote(eps, str(root), "") + futures.append(fut) + meta[fut] = (idx, n) + + deadline = ( + time.perf_counter() + float(self.shard_timeout) + if self.shard_timeout > 0 else None + ) + pending = list(futures) + while pending: + if deadline is not None: + remaining = deadline - time.perf_counter() + if remaining <= 0: + break + else: + remaining = None + + ready, pending = ray.wait(pending, num_returns=1, timeout=remaining) + if not ready: + break + + fut = ready[0] + idx, n = meta[fut] + try: + triples = ray.get(fut) + all_edges.extend( + PyCallEdge(source=s, target=t, weight=w, provenance=["pycg"]) + for s, t, w in triples + ) + logger.debug("PyCG shard %d: %d edges from %d files (Ray)", idx, len(triples), n) + except Exception as exc: + logger.warning("PyCG shard %d failed — skipped: %s", idx, exc) + skipped += 1 + progress.advance() + + for fut in pending: + idx, _ = meta[fut] + logger.warning( + "PyCG shard %d timed out after %ds — skipped", + idx, self.shard_timeout, + ) + ray.cancel(fut, force=True) + skipped += 1 + progress.advance() + finally: + for root in roots: + shutil.rmtree(root, ignore_errors=True) + + if skipped: + logger.warning( + "PyCG: %d/%d shard(s) skipped (ceiling, %ds timeout, or failure)", + skipped, len(plan.shards), self.shard_timeout, + ) + + result = self._coalesce_edges(all_edges) + logger.info( + "PyCG: %d edges from %d/%d shard(s) (%d before dedup, Jedi-planned, Ray-parallel)", + len(result), len(plan.shards) - skipped, len(plan.shards), len(all_edges), + ) + return result + def _build_sharded( self, entry_points: List[str], From 3985d6960fe500e918ea363fafe9c070ed951831 Mon Sep 17 00:00:00 2001 From: Rahul Krishna Date: Sat, 27 Jun 2026 00:10:08 -0400 Subject: [PATCH 3/5] fix(pycg): key the shard module graph by file path, not module_name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PyModule.module_name is only the file stem (py_file.stem), which collides heavily across a real project — every __init__.py, models.py, views.py shares a name. Keying the partition graph by module_name collapsed all same-stem files into a single node and, via the last-wins module_name->file_path map, silently dropped every colliding file but one from the shards. Observed on odoo: a 1028-file symbol table produced a graph of only 399 nodes (4 fat shards), so ~600 files were never handed to PyCG. Key graph nodes by file_path (unique) instead; carry module_name as a node attribute for readable reporting. plan_shards now emits file-path shards directly (no name->file remap) with a parallel module_shards name view. Add a regression test asserting every file lands in exactly one shard under stem collisions, and update graph tests for file-keyed nodes. --- .../semantic_analysis/pycg/shard_planner.py | 69 ++++++++++--------- test/test_shard_planner.py | 54 ++++++++++++--- 2 files changed, 81 insertions(+), 42 deletions(-) diff --git a/codeanalyzer/semantic_analysis/pycg/shard_planner.py b/codeanalyzer/semantic_analysis/pycg/shard_planner.py index 43bf648..a1608dc 100644 --- a/codeanalyzer/semantic_analysis/pycg/shard_planner.py +++ b/codeanalyzer/semantic_analysis/pycg/shard_planner.py @@ -61,7 +61,7 @@ # ---------------------------------------------------------------------------- -# Symbol-table walking: callable / class signature -> owning module +# Symbol-table walking: callable / class signature -> defining file # ---------------------------------------------------------------------------- def _walk_callable_sigs(c: PyCallable) -> Iterator[str]: @@ -80,22 +80,25 @@ def _walk_class_sigs(cls: PyClass) -> Iterator[str]: yield from _walk_class_sigs(inner) -def _signature_to_module(symbol_table: Dict[str, PyModule]) -> Dict[str, str]: - """Map every callable/class signature in the project to its module name. +def _signature_to_file(symbol_table: Dict[str, PyModule]) -> Dict[str, str]: + """Map every callable/class signature in the project to its defining file. - Built by walking each module's full nesting tree, so the mapping is exact - rather than relying on longest-prefix matching against module names (which - is ambiguous when one module name is a prefix of another). + Built by walking each module's full nesting tree, recording the file that + actually defines each callable. Files are the partition unit because + ``PyModule.module_name`` is only the file *stem* (``py_file.stem``), which + collides heavily across a real project (every ``__init__.py``, ``models.py`` + …) — keying the module graph by name would collapse unrelated files into + one node and silently drop files from shards. ``file_path`` is unique. """ - sig_to_mod: Dict[str, str] = {} + sig_to_file: Dict[str, str] = {} for module in symbol_table.values(): for fn in module.functions.values(): for sig in _walk_callable_sigs(fn): - sig_to_mod[sig] = module.module_name + sig_to_file[sig] = module.file_path for cls in module.classes.values(): for sig in _walk_class_sigs(cls): - sig_to_mod[sig] = module.module_name - return sig_to_mod + sig_to_file[sig] = module.file_path + return sig_to_file # ---------------------------------------------------------------------------- @@ -133,22 +136,24 @@ def build_module_graph( symbol_table: Dict[str, PyModule], jedi_edges: List[PyCallEdge], ) -> nx.DiGraph: - """Project Jedi callable→callable edges onto a weighted module DiGraph. - - Every project module is a node (isolated modules included). Edge weight is - the summed Jedi weight of cross-module call sites; intra-module edges and - edges touching external/library symbols (no symbol-table entry) are - dropped — they cannot influence how the project is partitioned. + """Project Jedi callable→callable edges onto a weighted file DiGraph. + + Nodes are file paths (the unique, collision-free partition unit — see + :func:`_signature_to_file`); each carries a ``module_name`` attribute for + readable reporting. Every project file is a node (isolated files + included). Edge weight is the summed Jedi weight of cross-file call sites; + intra-file edges and edges touching external/library symbols (no + symbol-table entry) are dropped — they cannot influence the partition. """ - sig_to_mod = _signature_to_module(symbol_table) + sig_to_file = _signature_to_file(symbol_table) g = nx.DiGraph() for module in symbol_table.values(): - g.add_node(module.module_name, file_path=module.file_path) + g.add_node(module.file_path, module_name=module.module_name) for edge in jedi_edges: - src = sig_to_mod.get(edge.source) - dst = sig_to_mod.get(edge.target) + src = sig_to_file.get(edge.source) + dst = sig_to_file.get(edge.target) if src is None or dst is None or src == dst: continue if g.has_edge(src, dst): @@ -355,23 +360,25 @@ def plan_shards( if merge_small: unit_shards = _merge_small(unit_shards, hu, unit_size, budget) - # Expand SCC units back to modules, then to file paths. - module_shards: List[List[str]] = [] + # Expand SCC units back to file paths (graph nodes are files). + file_shards: List[List[str]] = [] for units in unit_shards: - mods: Set[str] = set() + files: Set[str] = set() for scc in units: - mods |= unit_members[scc] - if mods: - module_shards.append(sorted(mods)) + files |= unit_members[scc] + if files: + file_shards.append(sorted(files)) - mod_to_file = {n: g.nodes[n]["file_path"] for n in g.nodes} - file_shards = [[mod_to_file[m] for m in mods] for mods in module_shards] + # Parallel view in module names (file stems) for human-readable reporting. + module_shards = [ + [g.nodes[f].get("module_name", f) for f in files] for files in file_shards + ] # Metrics: how much Jedi edge weight does this partition sever? shard_of: Dict[str, int] = {} - for idx, mods in enumerate(module_shards): - for m in mods: - shard_of[m] = idx + for idx, files in enumerate(file_shards): + for f in files: + shard_of[f] = idx cut_weight = 0.0 for u, v, w in g.edges(data="weight", default=1): if shard_of.get(u) != shard_of.get(v): diff --git a/test/test_shard_planner.py b/test/test_shard_planner.py index aff3e58..fe6142c 100644 --- a/test/test_shard_planner.py +++ b/test/test_shard_planner.py @@ -31,8 +31,9 @@ def _edge(src: str, dst: str, w: int = 1) -> PyCallEdge: return PyCallEdge(source=src, target=dst, weight=w, provenance=["jedi"]) -def _cut_ratio(g: nx.DiGraph, module_shards: List[List[str]]) -> float: - shard_of = {m: i for i, mods in enumerate(module_shards) for m in mods} +def _cut_ratio(g: nx.DiGraph, file_shards: List[List[str]]) -> float: + # Graph nodes are file paths; partitions are lists of file paths. + shard_of = {f: i for i, files in enumerate(file_shards) for f in files} total = cut = 0.0 for u, v, w in g.edges(data="weight", default=1): total += w @@ -45,27 +46,41 @@ def _cut_ratio(g: nx.DiGraph, module_shards: List[List[str]]) -> float: # build_module_graph # ---------------------------------------------------------------------------- -def test_module_graph_projects_callables_to_modules(): +def test_module_graph_keys_nodes_by_file_not_name(): + # Two distinct files share the stem "models" (module_name collision); the + # graph must keep them as separate nodes keyed by path, not collapse them. + st = { + "/pkg_a/models.py": _module("models", "/pkg_a/models.py", ["f"]), + "/pkg_b/models.py": _module("models", "/pkg_b/models.py", ["h"]), + } + edges = [ + _edge("models.f", "models.h", 3), # but which file? mapped by definition + ] + g = build_module_graph(st, edges) + assert set(g.nodes) == {"/pkg_a/models.py", "/pkg_b/models.py"} + assert g.nodes["/pkg_a/models.py"]["module_name"] == "models" + + +def test_module_graph_projects_callables_to_files(): st = { "/a.py": _module("a", "/a.py", ["f", "g"]), "/b.py": _module("b", "/b.py", ["h"]), } edges = [ - _edge("a.f", "b.h", 3), # cross-module -> kept - _edge("a.f", "a.g", 5), # intra-module -> dropped + _edge("a.f", "b.h", 3), # cross-file -> kept + _edge("a.f", "a.g", 5), # intra-file -> dropped _edge("a.g", "ext.lib.x"), # external target -> dropped ] g = build_module_graph(st, edges) - assert set(g.nodes) == {"a", "b"} - assert g.has_edge("a", "b") and g["a"]["b"]["weight"] == 3 - assert not g.has_edge("a", "a") + assert set(g.nodes) == {"/a.py", "/b.py"} + assert g.has_edge("/a.py", "/b.py") and g["/a.py"]["/b.py"]["weight"] == 3 assert g.number_of_edges() == 1 -def test_isolated_modules_are_nodes(): +def test_isolated_files_are_nodes(): st = {"/a.py": _module("a", "/a.py", ["f"]), "/b.py": _module("b", "/b.py", ["g"])} g = build_module_graph(st, []) # no edges - assert set(g.nodes) == {"a", "b"} + assert set(g.nodes) == {"/a.py", "/b.py"} assert g.number_of_edges() == 0 @@ -105,9 +120,11 @@ def test_every_module_assigned_exactly_once(): def test_beats_naive_per_package_cut_ratio(): st, edges = _coupled_clusters_project() g = build_module_graph(st, edges) + # Naive baseline: one shard per top-level directory (e.g. /pkg_a/...). naive = {} for m in st.values(): - naive.setdefault(m.module_name.split(".")[0], []).append(m.module_name) + top = m.file_path.split("/")[1] + naive.setdefault(top, []).append(m.file_path) naive_ratio = _cut_ratio(g, list(naive.values())) plan = plan_shards(st, edges, budget=4) @@ -150,6 +167,21 @@ def test_determinism(): assert norm(a) == norm(b) +def test_no_file_dropped_on_stem_collision(): + # Regression: module_name is only the file stem, so many files collide on + # name (every __init__.py, models.py, ...). Keying by name would drop all + # but one per stem. Every FILE must land in exactly one shard. + st, edges = {}, [] + for pkg in ("a", "b", "c", "d"): + for stem in ("__init__", "models", "views"): + path = f"/{pkg}/{stem}.py" + st[path] = _module(stem, path, ["f"]) + plan = plan_shards(st, edges, budget=5) + assigned = sorted(f for shard in plan.shards for f in shard) + assert assigned == sorted(st.keys()) # all 12 files present + assert len(assigned) == len(set(assigned)) # none duplicated + + def test_empty_project(): plan = plan_shards({}, [], budget=4) assert plan.shards == [] From 81fd409d63dd34346ea78faa1015bd9c8586a72e Mon Sep 17 00:00:00 2001 From: Rahul Krishna Date: Sat, 27 Jun 2026 09:25:47 -0400 Subject: [PATCH 4/5] feat(pycg): bound the fixpoint with --pycg-max-iter; stop following in-tree deps Two robustness fixes for level-2 PyCG, motivated by odoo divergence analysis. 1. max_iter cap (--pycg-max-iter, default 50). PyCG runs its PostProcessor fixpoint with max_iter=-1 (until convergence). Its abstract domain is field-sensitive access paths with no k-limiting/widening, so on heavy metaclass/mixin code the def set balloons (measured: 23 odoo ORM files -> 7.3k defs pass 0, 8.4k pass 1) and convergence may need many O(defs^2) passes. Capping passes returns a sound-but-incomplete graph and guarantees termination even with --pycg-shard-timeout 0 (which previously hung forever on a single diverging shard). Threaded through _run_pycg_batch and the Ray worker. Note: the wall-clock timeout is still the guard for shards whose individual passes exceed it. 2. Dependency exclusion. PyCG bounds analysis to its package dir via "if mod_dir not in mod.__file__". The whole-project path used package=project_dir, but an in-tree .codeanalyzer venv / site-packages lives under project_dir, so PyCG followed imports into dependencies and exploded. Run the whole-project path inside a symlink mini-project (as the shards already do) whose root mirrors only the SKIP_DIRS-filtered source, so deps resolve outside mod_dir and stay ghost nodes. Add test/test_pycg_sharding.py (max_iter threading; in-tree dep stays a ghost and its internals are never analysed). --- codeanalyzer/__main__.py | 17 ++++++ codeanalyzer/core.py | 1 + codeanalyzer/options/options.py | 1 + .../semantic_analysis/pycg/pycg_analysis.py | 45 ++++++++++---- test/test_pycg_sharding.py | 61 +++++++++++++++++++ 5 files changed, 112 insertions(+), 13 deletions(-) create mode 100644 test/test_pycg_sharding.py diff --git a/codeanalyzer/__main__.py b/codeanalyzer/__main__.py index 8c6bf72..51f227f 100644 --- a/codeanalyzer/__main__.py +++ b/codeanalyzer/__main__.py @@ -200,6 +200,22 @@ def main( ), ), ] = ShardStrategy.JEDI, + pycg_max_iter: Annotated[ + int, + typer.Option( + "--pycg-max-iter", + help=( + "Cap on PyCG's fixpoint passes per shard/project (level 2; " + "default 50). PyCG iterates until its points-to state stops " + "changing, but its access-path domain has no convergence bound, " + "so heavy metaclass/mixin code (e.g. an ORM) can loop with each " + "pass costing seconds. The cap returns a sound-but-incomplete " + "call graph instead of looping until the timeout kills it. " + "Set to -1 for PyCG's unbounded run-to-convergence behaviour." + ), + min=-1, + ), + ] = 50, ): options = AnalysisOptions( input=input, @@ -224,6 +240,7 @@ def main( pycg_shard_ceiling=pycg_shard_ceiling, pycg_shard_timeout=pycg_shard_timeout, pycg_shard_strategy=pycg_shard_strategy, + pycg_max_iter=pycg_max_iter, ) _set_log_level(options.verbosity) diff --git a/codeanalyzer/core.py b/codeanalyzer/core.py index fc87f95..151446e 100644 --- a/codeanalyzer/core.py +++ b/codeanalyzer/core.py @@ -683,6 +683,7 @@ def _get_pycg_call_graph( shard_ceiling=self.options.pycg_shard_ceiling, shard_timeout=self.options.pycg_shard_timeout, shard_strategy=self.options.pycg_shard_strategy, + max_iter=self.options.pycg_max_iter, using_ray=self.using_ray, ) return pycg.build_call_graph_edges(symbol_table, jedi_edges=jedi_edges) diff --git a/codeanalyzer/options/options.py b/codeanalyzer/options/options.py index 5e3b56b..4e8662c 100644 --- a/codeanalyzer/options/options.py +++ b/codeanalyzer/options/options.py @@ -61,3 +61,4 @@ class AnalysisOptions: pycg_shard_ceiling: int = 100 pycg_shard_timeout: int = 120 pycg_shard_strategy: ShardStrategy = ShardStrategy.JEDI + pycg_max_iter: int = 50 diff --git a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py index 8beb03e..b109dc6 100644 --- a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py +++ b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py @@ -154,13 +154,14 @@ def _pycg_shard_worker( entry_points: List[str], package_dir: str, prefix: str, + max_iter: int = -1, ) -> List[tuple]: """Run PyCG on one shard; called in a Ray worker process. Returns a list of ``(source, target, weight)`` tuples that the caller converts to :class:`PyCallEdge` objects. This function is a plain module-level callable so it can be pickled by Ray without capturing any - class-level state. + class-level state. *max_iter* caps PyCG's fixpoint passes (-1 = unbounded). """ import importlib import sys @@ -191,7 +192,7 @@ def _pycg_shard_worker( cg = CallGraphGenerator( entry_points=entry_points, package=package_dir, - max_iter=-1, + max_iter=max_iter, operation="call-graph", ) cg.analyze() @@ -366,7 +367,22 @@ class PyCG: # --pycg-shard-timeout. Set to 0 to disable. _PYCG_SHARD_TIMEOUT: int = 120 - # Directory names that should never be fed to PyCG as entry points. + # Cap on PyCG's outer fixpoint passes. PyCG runs PostProcessor until the + # def/scope/MRO state stops changing; its abstract domain (field-sensitive + # access paths, no k-limiting or widening) has no ascending-chain bound, so + # on heavy metaclass/mixin code (e.g. an ORM) the def set can balloon into + # the thousands and each O(defs^2) pass costs seconds — convergence, if it + # comes, takes many passes. A finite cap turns "loop until killed" into a + # sound-but-incomplete result that still returns the edges found so far. + # 50 is generous — well-behaved code converges in well under 20 passes — + # while bounding the pathological case. Override via --pycg-max-iter; + # -1 restores PyCG's unbounded run-to-convergence behaviour. + _PYCG_MAX_ITER: int = 50 + + # Directory names that should never be fed to PyCG as entry points, nor + # followed into during import resolution (an in-tree .codeanalyzer venv / + # site-packages lives under project_dir and would otherwise be pulled into + # the package bound and analysed — see _shard_symlink_root). _SKIP_DIRS: frozenset = frozenset({ ".codeanalyzer", ".git", "__pycache__", "venv", ".venv", "virtualenv", "env", ".env", @@ -382,6 +398,7 @@ def __init__( shard_ceiling: Optional[int] = None, shard_timeout: Optional[int] = None, shard_strategy: str = "jedi", + max_iter: Optional[int] = None, using_ray: bool = False, ) -> None: self.project_dir = Path(project_dir).resolve() @@ -393,6 +410,7 @@ def __init__( self.shard_timeout = ( shard_timeout if shard_timeout is not None else self._PYCG_SHARD_TIMEOUT ) + self.max_iter = max_iter if max_iter is not None else self._PYCG_MAX_ITER # "jedi": partition the Jedi module graph (SCC + Louvain) so coupled # modules co-compute and few edges are severed (see shard_planner). # "package": legacy one-shard-per-package-directory grouping. @@ -519,7 +537,7 @@ def _run_pycg_batch( cg = self._CallGraphGenerator( entry_points=entry_points, package=str(package_dir), - max_iter=-1, + max_iter=self.max_iter, operation="call-graph", ) cg.analyze() @@ -656,7 +674,7 @@ def _build_sharded_planned_ray(self, plan: "ShardPlan") -> List[PyCallEdge]: continue root, eps = _materialize_shard_root(files, self.project_dir) roots.append(root) - fut = remote_fn.remote(eps, str(root), "") + fut = remote_fn.remote(eps, str(root), "", self.max_iter) futures.append(fut) meta[fut] = (idx, n) @@ -844,7 +862,7 @@ def _build_sharded_ray(self, shards: Dict[Path, List[str]]) -> List[PyCallEdge]: progress.advance() continue prefix = self._package_prefix(pkg_root, self.project_dir) - fut = remote_fn.remote(files, str(pkg_root), prefix) + fut = remote_fn.remote(files, str(pkg_root), prefix, self.max_iter) futures.append(fut) meta[fut] = (pkg_label, n) @@ -989,14 +1007,15 @@ def build_call_graph_edges( ) return [] else: - # Small project (≤ ceiling): whole-project analysis. + # Small project (≤ ceiling): whole-project analysis. Run inside a + # symlink mini-project mirroring only the (already SKIP_DIRS-filtered) + # entry points, so PyCG's package bound covers project source alone. + # Pointing PyCG at project_dir directly would put an in-tree + # .codeanalyzer venv / site-packages *under* mod_dir, and PyCG would + # follow imports into those dependencies and explode the analysis. logger.info("PyCG: starting whole-project call graph analysis (%d files)", n_files) - try: - edges = self._run_pycg_batch( - entry_points, self.project_dir, resolver, prefix="" - ) - except PyCGExceptions.PyCGAnalysisError as exc: - raise + with _shard_symlink_root(entry_points, self.project_dir) as (root, eps): + edges = self._run_pycg_batch(eps, root, resolver, prefix="") elapsed = time.perf_counter() - t0 logger.info("✅ PyCG: %d edges in %.1fs", len(edges), elapsed) diff --git a/test/test_pycg_sharding.py b/test/test_pycg_sharding.py new file mode 100644 index 0000000..ec43c68 --- /dev/null +++ b/test/test_pycg_sharding.py @@ -0,0 +1,61 @@ +"""Tests for PyCG executor scoping: dependency exclusion and the max_iter cap. + +These drive the real PyCG wrapper, so they require ``pycg`` (a level-2 install +dependency). They are deliberately tiny (a few files) so they run fast. +""" +from pathlib import Path + +import pytest + +pytest.importorskip("PyCG") + +from codeanalyzer.semantic_analysis.pycg.pycg_analysis import ( + PyCG, + _PyCGCallableResolver, + _shard_symlink_root, +) + + +def test_max_iter_default_and_override(tmp_path): + p = PyCG(tmp_path) + assert p.max_iter == PyCG._PYCG_MAX_ITER == 50 + assert PyCG(tmp_path, max_iter=7).max_iter == 7 + assert PyCG(tmp_path, max_iter=-1).max_iter == -1 + + +def test_pycg_does_not_follow_into_in_tree_dependency(tmp_path): + """An in-tree ``.codeanalyzer`` venv under project_dir must stay a ghost. + + PyCG bounds analysis to its ``package`` directory; running inside the + symlink mini-project keeps that bound on project source only, so imports + into a bundled dependency are recorded as ghost edges but never analysed. + Regression guard for the dep-reach blowup. + """ + proj = tmp_path + app = proj / "app" + app.mkdir() + (app / "__init__.py").write_text("") + (app / "main.py").write_text("import bigdep\ndef run():\n return bigdep.work()\n") + + # A bundled dependency with many internal functions: if PyCG followed into + # it, dozens of bigdep.fN definitions/edges would appear. + dep = proj / ".codeanalyzer" / "venv" / "site-packages" / "bigdep" + dep.mkdir(parents=True) + body = "".join(f"def f{i}(x):\n return f{(i + 1) % 50}(x)\n" for i in range(50)) + body += "def work():\n return f0(1)\n" + (dep / "__init__.py").write_text(body) + + pycg = PyCG(proj) + pycg._ensure_pycg_loaded() + resolver = _PyCGCallableResolver(set()) + entry_points = [str(app / "__init__.py"), str(app / "main.py")] + with _shard_symlink_root(entry_points, proj) as (root, eps): + edges = pycg._run_pycg_batch(eps, root, resolver, prefix="") + + nodes = {n for e in edges for n in (e.source, e.target)} + # bigdep is reachable as a ghost target ... + assert any(n.startswith("bigdep") for n in nodes) + # ... but none of its internals were analysed. + assert not [n for n in nodes if n.startswith("bigdep.f")] + # and the real app edge is present. + assert any(e.source == "app.main.run" and e.target == "bigdep.work" for e in edges) From f72ef0041228ec8b4e92aea9dcc151ddc9d24849 Mon Sep 17 00:00:00 2001 From: Rahul Krishna Date: Sat, 27 Jun 2026 11:26:16 -0400 Subject: [PATCH 5/5] feat(pycg): iterative decomposition of runaway shards A uniform shard ceiling forces a global choice: small shards everywhere (high cut, low recall) just to tame the few that diverge. Instead, start coarse and re-shard only the shards that time out. Algorithm: plan shards with SCC + Louvain at the ceiling, run each through PyCG, and treat any timed-out shard as a runaway. Re-partition that runaway's files alone at half the budget and re-run. Repeat down to a floor (10 files). Files that still diverge at the floor, or form an atomic cycle that will not split, fall back to Jedi-only coverage. Refactor the planned executor into a reusable primitive that returns (edges, runaways), used by both the sequential and Ray paths, and drive it from an adaptive loop. Odoo benchmark (1028 modules, level 2, Ray): 22210 PyCG edges, up from 17149 for the best uniform ceiling, with only 20 of 1028 files irreducible. Cost is wall time (about 12.7 min) since rounds run in sequence. Add a unit test driving the adaptive loop with a stubbed runner. --- .../semantic_analysis/pycg/pycg_analysis.py | 222 ++++++++++-------- test/test_pycg_sharding.py | 50 ++++ 2 files changed, 177 insertions(+), 95 deletions(-) diff --git a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py index b109dc6..e82c639 100644 --- a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py +++ b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py @@ -81,7 +81,7 @@ def _handler(signum: int, frame: object) -> None: from codeanalyzer.schema.py_schema import PyCallEdge, PyModule from codeanalyzer.semantic_analysis.call_graph import iter_callables_in_symbol_table from codeanalyzer.semantic_analysis.pycg.pycg_exceptions import PyCGExceptions -from codeanalyzer.semantic_analysis.pycg.shard_planner import ShardPlan, plan_shards +from codeanalyzer.semantic_analysis.pycg.shard_planner import plan_shards from codeanalyzer.utils import ProgressBar, logger @@ -379,6 +379,13 @@ class PyCG: # -1 restores PyCG's unbounded run-to-convergence behaviour. _PYCG_MAX_ITER: int = 50 + # Iterative decomposition of runaway (timed-out) shards: a shard that the + # wall-clock timeout kills is re-partitioned at half the budget and re-run, + # down to this file-count floor. Below the floor — or for an atomic import + # cycle that won't split — the residue falls back to Jedi-only coverage. + _PYCG_DECOMP_FLOOR: int = 10 + _PYCG_MAX_DECOMP_ROUNDS: int = 6 + # Directory names that should never be fed to PyCG as entry points, nor # followed into during import resolution (an in-tree .codeanalyzer venv / # site-packages lives under project_dir and would otherwise be pulled into @@ -417,6 +424,7 @@ def __init__( self.shard_strategy = shard_strategy self.using_ray = using_ray self._CallGraphGenerator: Optional[Any] = None + self._resolver: Optional["_PyCGCallableResolver"] = None @staticmethod def _coalesce_edges(edges: List[PyCallEdge]) -> List[PyCallEdge]: @@ -570,20 +578,26 @@ def _build_sharded_planned( symbol_table: Dict[str, PyModule], resolver: "_PyCGCallableResolver", ) -> List[PyCallEdge]: - """Coupling-aware sharding driven by the Jedi module graph. - - Unlike :meth:`_build_sharded` (one shard per package directory), the - shards here are chosen to *minimise the call edges severed between - shards*: :func:`shard_planner.plan_shards` condenses the Jedi call - graph by strongly-connected component (so import cycles never split) - and clusters it with Louvain so tightly-coupled modules land together. - Each shard — an arbitrary set of files — is run through PyCG via a - symlinked mini-project (:func:`_shard_symlink_root`) that bounds PyCG - to exactly those files. - - Reported ``cut_ratio`` is the fraction of Jedi edge weight crossing - shard boundaries — an upper bound on the PyCG edges lost to sharding. + """Coupling-aware sharding with iterative decomposition of runaways. + + Shards are chosen to *minimise the call edges severed between shards*: + :func:`shard_planner.plan_shards` condenses the Jedi call graph by + strongly-connected component (so import cycles never split) and clusters + it with Louvain so tightly-coupled modules land together. Each shard is + run through PyCG via a symlinked mini-project that bounds analysis to its + files. + + PyCG's fixpoint diverges on heavy metaclass/mixin clusters, and a uniform + ceiling would force *every* shard small (severing many edges) just to tame + the few that run away. Instead we start coarse (low cut, high recall on + healthy code) and **only re-decompose the shards that time out**: each + runaway's files are re-partitioned at half the budget and re-run, down to + a floor. A runaway shard contributes zero edges, so splitting it recovers + almost all of them while paying cut on its internal seams alone. The + residue that still diverges at the floor (or is an atomic cycle that won't + split) falls back to Jedi-only coverage. """ + self._resolver = resolver plan = plan_shards( symbol_table, jedi_edges, budget=self.shard_ceiling, merge_small=True ) @@ -594,64 +608,107 @@ def _build_sharded_planned( int(m["num_shards"]), m["cut_ratio"], int(m["max_shard_files"]), int(m["modules"]), ) - if m["oversized_shards"]: - logger.warning( - "PyCG: %d shard(s) exceed the %d-file ceiling — skipped " - "(atomic import cycles larger than the budget)", - int(m["oversized_shards"]), self.shard_ceiling, + + runner = ( + self._run_fileset_shards_ray if self.using_ray + else self._run_fileset_shards_seq + ) + all_edges: List[PyCallEdge] = [] + shards = plan.shards + budget = self.shard_ceiling + converged_total = 0 + irreducible_files = 0 + round_no = 0 + + while shards: + label = "decomposition round %d (budget %d, %d shard(s))" % ( + round_no, budget, len(shards), ) + logger.info("PyCG: %s", label) + edges, runaways = runner(shards) + all_edges.extend(edges) + converged_total += len(shards) - len(runaways) + if not runaways: + break - if self.using_ray: - return self._build_sharded_planned_ray(plan) + next_budget = max(self._PYCG_DECOMP_FLOOR, budget // 2) + stop_decomposing = ( + round_no >= self._PYCG_MAX_DECOMP_ROUNDS or next_budget >= budget + ) - all_edges: List[PyCallEdge] = [] - skipped = 0 - with ProgressBar(len(plan.shards), "Building call graph shards", item_label="shards") as progress: - for idx, files in enumerate(plan.shards): - n = len(files) - if n > self.shard_ceiling: - skipped += 1 - progress.advance() + next_shards: List[List[str]] = [] + for rf in runaways: + # Re-partition this runaway's files alone, at a tighter budget. + # An atomic cycle (or a lone file) that won't shrink is + # irreducible — accept Jedi-only rather than loop forever. + sub_st = {f: symbol_table[f] for f in rf if f in symbol_table} + if stop_decomposing or len(rf) <= 1: + irreducible_files += len(rf) continue - try: - with _shard_symlink_root(files, self.project_dir) as (root, eps): - with _shard_timeout(self.shard_timeout): - edges = self._run_pycg_batch(eps, root, resolver, prefix="") - all_edges.extend(edges) - logger.debug("PyCG shard %d: %d edges from %d files", idx, len(edges), n) - except TimeoutError: - logger.warning( - "PyCG shard %d timed out after %ds — skipped", - idx, self.shard_timeout, - ) - skipped += 1 - except PyCGExceptions.PyCGAnalysisError as exc: - logger.warning("PyCG shard %d failed — skipped: %s", idx, exc) - skipped += 1 - progress.advance() + sub_plan = plan_shards(sub_st, jedi_edges, budget=next_budget) + if len(sub_plan.shards) <= 1: + # did not actually split (one atomic SCC) — give up on it + irreducible_files += len(rf) + continue + next_shards.extend(sub_plan.shards) - if skipped: + if not next_shards: + break + logger.info( + "PyCG: %d shard(s) ran away — decomposing into %d sub-shard(s) " + "at budget %d", len(runaways), len(next_shards), next_budget, + ) + shards, budget = next_shards, next_budget + round_no += 1 + + if irreducible_files: logger.warning( - "PyCG: %d/%d shard(s) skipped (ceiling, %ds timeout, or failure)", - skipped, len(plan.shards), self.shard_timeout, + "PyCG: %d file(s) in irreducibly-divergent shards fall back to " + "Jedi-only coverage", irreducible_files, ) result = self._coalesce_edges(all_edges) logger.info( - "PyCG: %d edges from %d/%d shard(s) (%d before dedup, Jedi-planned)", - len(result), len(plan.shards) - skipped, len(plan.shards), len(all_edges), + "PyCG: %d edges from %d converged shard(s) over %d round(s) " + "(%d before dedup, Jedi-planned%s)", + len(result), converged_total, round_no + 1, len(all_edges), + ", Ray-parallel" if self.using_ray else "", ) return result - def _build_sharded_planned_ray(self, plan: "ShardPlan") -> List[PyCallEdge]: - """Ray-parallel execution of Jedi-planned file-set shards. + def _run_fileset_shards_seq( + self, shards: List[List[str]], + ) -> Tuple[List[PyCallEdge], List[List[str]]]: + """Run each file-set shard sequentially; return ``(edges, runaways)``. - Each shard is materialised as a symlink mini-project up front (the - trees must outlive their remote tasks), submitted as a Ray task, and - collected against a single wall-clock deadline — Ray workers cannot use - SIGALRM, so the timeout is enforced at the orchestrator level (mirroring - :meth:`_build_sharded_ray`). All symlink trees are removed once the - batch completes. + A shard that times out or raises is returned in *runaways* (its file + list) for the caller to re-decompose; it contributes no edges. + """ + resolver = self._resolver + edges_all: List[PyCallEdge] = [] + runaways: List[List[str]] = [] + with ProgressBar(len(shards), "Building call graph shards", item_label="shards") as progress: + for files in shards: + try: + with _shard_symlink_root(files, self.project_dir) as (root, eps): + with _shard_timeout(self.shard_timeout): + edges = self._run_pycg_batch(eps, root, resolver, prefix="") + edges_all.extend(edges) + except (TimeoutError, PyCGExceptions.PyCGAnalysisError): + runaways.append(files) + progress.advance() + return edges_all, runaways + + def _run_fileset_shards_ray( + self, shards: List[List[str]], + ) -> Tuple[List[PyCallEdge], List[List[str]]]: + """Ray-parallel variant of :meth:`_run_fileset_shards_seq`. + + Each shard is materialised as a symlink mini-project up front (the trees + must outlive their remote tasks), submitted as a Ray task, and collected + against one wall-clock deadline — Ray workers cannot use SIGALRM, so the + timeout is enforced orchestrator-side. Timed-out/failed shards become + runaways; symlink trees are removed once the batch completes. """ import os import ray @@ -661,22 +718,17 @@ def _build_sharded_planned_ray(self, plan: "ShardPlan") -> List[PyCallEdge]: roots: List[Path] = [] futures: List[Any] = [] - meta: Dict[Any, tuple] = {} # ObjectRef -> (shard_idx, n_files) - skipped = 0 - all_edges: List[PyCallEdge] = [] + meta: Dict[Any, List[str]] = {} # ObjectRef -> shard file list + edges_all: List[PyCallEdge] = [] + runaways: List[List[str]] = [] try: - with ProgressBar(len(plan.shards), "Building call graph shards (parallel)", item_label="shards") as progress: - for idx, files in enumerate(plan.shards): - n = len(files) - if n > self.shard_ceiling: - skipped += 1 - progress.advance() - continue + with ProgressBar(len(shards), "Building call graph shards (parallel)", item_label="shards") as progress: + for files in shards: root, eps = _materialize_shard_root(files, self.project_dir) roots.append(root) fut = remote_fn.remote(eps, str(root), "", self.max_iter) futures.append(fut) - meta[fut] = (idx, n) + meta[fut] = files deadline = ( time.perf_counter() + float(self.shard_timeout) @@ -696,44 +748,24 @@ def _build_sharded_planned_ray(self, plan: "ShardPlan") -> List[PyCallEdge]: break fut = ready[0] - idx, n = meta[fut] try: triples = ray.get(fut) - all_edges.extend( + edges_all.extend( PyCallEdge(source=s, target=t, weight=w, provenance=["pycg"]) for s, t, w in triples ) - logger.debug("PyCG shard %d: %d edges from %d files (Ray)", idx, len(triples), n) - except Exception as exc: - logger.warning("PyCG shard %d failed — skipped: %s", idx, exc) - skipped += 1 + except Exception: + runaways.append(meta[fut]) progress.advance() - for fut in pending: - idx, _ = meta[fut] - logger.warning( - "PyCG shard %d timed out after %ds — skipped", - idx, self.shard_timeout, - ) + for fut in pending: # exceeded the deadline ray.cancel(fut, force=True) - skipped += 1 + runaways.append(meta[fut]) progress.advance() finally: for root in roots: shutil.rmtree(root, ignore_errors=True) - - if skipped: - logger.warning( - "PyCG: %d/%d shard(s) skipped (ceiling, %ds timeout, or failure)", - skipped, len(plan.shards), self.shard_timeout, - ) - - result = self._coalesce_edges(all_edges) - logger.info( - "PyCG: %d edges from %d/%d shard(s) (%d before dedup, Jedi-planned, Ray-parallel)", - len(result), len(plan.shards) - skipped, len(plan.shards), len(all_edges), - ) - return result + return edges_all, runaways def _build_sharded( self, diff --git a/test/test_pycg_sharding.py b/test/test_pycg_sharding.py index ec43c68..c0629e6 100644 --- a/test/test_pycg_sharding.py +++ b/test/test_pycg_sharding.py @@ -23,6 +23,56 @@ def test_max_iter_default_and_override(tmp_path): assert PyCG(tmp_path, max_iter=-1).max_iter == -1 +def test_adaptive_decomposition_splits_runaways(tmp_path, monkeypatch): + """A shard that 'runs away' is re-decomposed until its pieces converge. + + Drives the real adaptive loop + planner, but stubs the PyCG runner with a + size threshold: shards larger than the threshold time out (runaway); smaller + ones converge and yield one synthetic edge per file. A 16-file coupled + cluster must therefore be split across rounds until every piece is small + enough, with no files lost. + """ + from codeanalyzer.schema.py_schema import PyCallEdge, PyCallable, PyModule + from codeanalyzer.semantic_analysis.pycg.pycg_analysis import _PyCGCallableResolver + + # A loosely-coupled chain of 40 modules: splittable (not one atomic cycle). + st, jedi = {}, [] + for i in range(40): + path = f"/proj/m{i}.py" + st[path] = PyModule( + file_path=path, module_name=f"m{i}", + functions={"f": PyCallable(signature=f"m{i}.f", name="f", path=path)}, + ) + if i: + jedi.append(PyCallEdge(source=f"m{i-1}.f", target=f"m{i}.f", weight=1, + provenance=["jedi"])) + + # threshold >= the decomposition floor (10) so pieces can shrink enough to converge. + pycg = PyCG(tmp_path, shard_ceiling=40) + threshold = 12 # shards with > 12 files "diverge" + rounds_seen = [] + + def fake_runner(shards): + rounds_seen.append([len(s) for s in shards]) + edges, runaways = [], [] + for files in shards: + if len(files) > threshold: + runaways.append(files) + else: + edges += [PyCallEdge(source=f, target="x", weight=1, provenance=["pycg"]) + for f in files] + return edges, runaways + + monkeypatch.setattr(pycg, "_run_fileset_shards_seq", fake_runner) + edges = pycg._build_sharded_planned(jedi, st, _PyCGCallableResolver(set())) + + assert len(rounds_seen) >= 2, "runaway shard was never decomposed" + # Every shard that was finally accepted is within the convergence threshold. + assert all(sz <= threshold for sz in rounds_seen[-1]) + # No files lost: one pycg edge per file across all 16 modules. + assert len({e.source for e in edges}) == 40 + + def test_pycg_does_not_follow_into_in_tree_dependency(tmp_path): """An in-tree ``.codeanalyzer`` venv under project_dir must stay a ghost.