shard.py
file-level
1
files
1
commits
0
hotspots
0
🧊 dead
0
💥 blast risk
| 1 | """``muse coord shard`` — partition a codebase into low-coupling work zones. |
| 2 | |
| 3 | Divides files into N groups (shards) such that files within each shard are |
| 4 | tightly coupled to each other (many imports between them) and loosely coupled |
| 5 | to files in other shards. Each shard is a safe parallel-work zone for agents. |
| 6 | |
| 7 | Agents assigned to different shards are unlikely to create merge conflicts |
| 8 | because they work on different parts of the dependency graph. |
| 9 | |
| 10 | Algorithm |
| 11 | --------- |
| 12 | 1. Build the import graph from the committed snapshot. |
| 13 | 2. Compute connected components of the import graph. |
| 14 | 3. Distribute components greedily into N shards, balancing by symbol count. |
| 15 | 4. Report each shard with its files, symbol count, and coupling score. |
| 16 | |
| 17 | The coupling score is the count of cross-shard import edges (lower = better). |
| 18 | |
| 19 | Usage:: |
| 20 | |
| 21 | muse coord shard --agents 4 |
| 22 | muse coord shard --agents 8 --commit HEAD~10 |
| 23 | muse coord shard --agents 4 --language Python |
| 24 | muse coord shard --agents 4 --json |
| 25 | |
| 26 | Output (text):: |
| 27 | |
| 28 | Shard plan — 4 agents, commit a1b2c3d4 |
| 29 | ────────────────────────────────────────────────────────────── |
| 30 | |
| 31 | Shard 1 (12 symbols, 3 files, coupling 1): |
| 32 | src/billing.py |
| 33 | src/billing_utils.py |
| 34 | src/tax.py |
| 35 | |
| 36 | Shard 2 (8 symbols, 2 files, coupling 0): |
| 37 | src/auth.py |
| 38 | src/session.py |
| 39 | ... |
| 40 | |
| 41 | Cross-shard edges: 1 |
| 42 | |
| 43 | JSON output schema:: |
| 44 | |
| 45 | { |
| 46 | "schema_version": str, |
| 47 | "commit": str, // 8-char short ID (display only) |
| 48 | "full_commit_id": str, // full commit ID for cross-referencing |
| 49 | "agents": int, |
| 50 | "shards_created": int, |
| 51 | "total_files": int, |
| 52 | "total_symbols": int, |
| 53 | "cross_shard_edges": int, |
| 54 | "shards": [ |
| 55 | { |
| 56 | "shard": int, |
| 57 | "files": [str, ...], |
| 58 | "symbol_count": int, |
| 59 | "coupling_score": int |
| 60 | }, |
| 61 | ... |
| 62 | ], |
| 63 | "elapsed_seconds": float |
| 64 | } |
| 65 | |
| 66 | Flags: |
| 67 | |
| 68 | ``--agents N`` |
| 69 | Number of parallel work zones to create (1–256, default: 4). |
| 70 | |
| 71 | ``--commit, -c REF`` |
| 72 | Use a historical snapshot instead of HEAD. |
| 73 | |
| 74 | ``--language LANG`` |
| 75 | Restrict to files of this language. |
| 76 | |
| 77 | ``--format`` / ``--json`` |
| 78 | Emit the shard plan as compact JSON. |
| 79 | |
| 80 | Exit codes:: |
| 81 | |
| 82 | 0 — success (includes no-files-found case) |
| 83 | 1 — bad arguments (--agents out of range) or commit not found |
| 84 | """ |
| 85 | |
| 86 | from __future__ import annotations |
| 87 | |
| 88 | import argparse |
| 89 | import json |
| 90 | import logging |
| 91 | import pathlib |
| 92 | import sys |
| 93 | import time |
| 94 | |
| 95 | from muse._version import __version__ |
| 96 | from muse.core.errors import ExitCode |
| 97 | from muse.core.object_store import read_object |
| 98 | from muse.core.repo import read_repo_id, require_repo |
| 99 | from muse.core.store import Manifest, get_commit_snapshot_manifest, read_current_branch, resolve_commit_ref |
| 100 | from muse.plugins.code._query import language_of, symbols_for_snapshot |
| 101 | from muse.plugins.code.ast_parser import parse_symbols |
| 102 | from muse.core.validation import sanitize_display |
| 103 | |
| 104 | type _AdjMap = dict[str, set[str]] |
| 105 | type _CounterMap = dict[str, int] |
| 106 | type _ShardMap = dict[str, int] |
| 107 | |
| 108 | logger = logging.getLogger(__name__) |
| 109 | |
| 110 | # ── Input constraints ───────────────────────────────────────────────────────── |
| 111 | |
| 112 | #: Allowed range for ``--agents``. |
| 113 | _MIN_AGENTS: int = 1 |
| 114 | _MAX_AGENTS: int = 256 |
| 115 | |
| 116 | |
| 117 | # ── Error helper ────────────────────────────────────────────────────────────── |
| 118 | |
| 119 | |
| 120 | def _err(msg: str, as_json: bool, status: str = "error") -> None: |
| 121 | """Print an error and return. Caller raises SystemExit.""" |
| 122 | if as_json: |
| 123 | print(json.dumps({"error": msg, "status": status})) |
| 124 | else: |
| 125 | print(f"❌ {msg}", file=sys.stderr) |
| 126 | |
| 127 | |
| 128 | # ── Internal helpers ────────────────────────────────────────────────────────── |
| 129 | |
| 130 | |
| 131 | def _file_stem(fp: str) -> str: |
| 132 | return pathlib.PurePosixPath(fp).stem |
| 133 | |
| 134 | |
| 135 | def _build_import_edges( |
| 136 | root: pathlib.Path, |
| 137 | manifest: Manifest, |
| 138 | language_filter: str | None, |
| 139 | ) -> list[tuple[str, str]]: |
| 140 | """Return ``(importer, importee)`` file-path pairs from the snapshot. |
| 141 | |
| 142 | Reads each file's object from the object store, parses its import symbols, |
| 143 | then resolves each imported name to a file path in the manifest via a |
| 144 | stem-to-path lookup table. |
| 145 | |
| 146 | Only files matching *language_filter* (when provided) are included. |
| 147 | Files whose object is missing from the store are silently skipped. |
| 148 | |
| 149 | Args: |
| 150 | root: Repository root. |
| 151 | manifest: Snapshot file-path → object-ID mapping. |
| 152 | language_filter: Language name to restrict to, or ``None`` for all. |
| 153 | |
| 154 | Returns: |
| 155 | List of ``(importer_path, importee_path)`` tuples. |
| 156 | """ |
| 157 | stem_to_file: Manifest = {} |
| 158 | for fp in manifest: |
| 159 | if language_filter and language_of(fp) != language_filter: |
| 160 | continue |
| 161 | stem_to_file[_file_stem(fp)] = fp |
| 162 | |
| 163 | edges: list[tuple[str, str]] = [] |
| 164 | for file_path, obj_id in sorted(manifest.items()): |
| 165 | if language_filter and language_of(file_path) != language_filter: |
| 166 | continue |
| 167 | raw = read_object(root, obj_id) |
| 168 | if raw is None: |
| 169 | continue |
| 170 | tree = parse_symbols(raw, file_path) |
| 171 | for rec in tree.values(): |
| 172 | if rec["kind"] != "import": |
| 173 | continue |
| 174 | imported = rec["qualified_name"].split(".")[-1].replace("import::", "") |
| 175 | target = stem_to_file.get(imported) |
| 176 | if target and target != file_path: |
| 177 | edges.append((file_path, target)) |
| 178 | return edges |
| 179 | |
| 180 | |
| 181 | def _connected_components( |
| 182 | files: list[str], |
| 183 | edges: list[tuple[str, str]], |
| 184 | ) -> list[frozenset[str]]: |
| 185 | """Return weakly-connected components of the import graph. |
| 186 | |
| 187 | Uses an iterative DFS to avoid stack overflow on large graphs. |
| 188 | |
| 189 | Args: |
| 190 | files: All file paths (graph nodes). |
| 191 | edges: ``(importer, importee)`` pairs (directed edges treated as |
| 192 | undirected for connectivity purposes). |
| 193 | |
| 194 | Returns: |
| 195 | List of :class:`frozenset` — one per connected component. |
| 196 | """ |
| 197 | adj: _AdjMap = {f: set() for f in files} |
| 198 | for a, b in edges: |
| 199 | adj.setdefault(a, set()).add(b) |
| 200 | adj.setdefault(b, set()).add(a) |
| 201 | |
| 202 | visited: set[str] = set() |
| 203 | components: list[frozenset[str]] = [] |
| 204 | |
| 205 | for start in files: |
| 206 | if start in visited: |
| 207 | continue |
| 208 | component: set[str] = set() |
| 209 | stack = [start] |
| 210 | while stack: |
| 211 | node = stack.pop() |
| 212 | if node in visited: |
| 213 | continue |
| 214 | visited.add(node) |
| 215 | component.add(node) |
| 216 | for neighbour in adj.get(node, set()): |
| 217 | if neighbour not in visited: |
| 218 | stack.append(neighbour) |
| 219 | components.append(frozenset(component)) |
| 220 | |
| 221 | return components |
| 222 | |
| 223 | |
| 224 | def _greedy_partition( |
| 225 | components: list[frozenset[str]], |
| 226 | sym_counts: _CounterMap, |
| 227 | n_shards: int, |
| 228 | ) -> list[frozenset[str]]: |
| 229 | """Distribute connected components into *n_shards* shards. |
| 230 | |
| 231 | Uses a greedy largest-first strategy: components are sorted by total |
| 232 | symbol count (descending) and each is assigned to the shard with the |
| 233 | currently smallest symbol total. This approximates the Longest Processing |
| 234 | Time (LPT) algorithm for makespan minimisation. |
| 235 | |
| 236 | Args: |
| 237 | components: Connected components to distribute. |
| 238 | sym_counts: File-path → symbol count mapping. |
| 239 | n_shards: Number of shards to create. |
| 240 | |
| 241 | Returns: |
| 242 | List of *n_shards* :class:`frozenset` objects (some may be empty when |
| 243 | there are fewer components than shards). |
| 244 | """ |
| 245 | shards: list[set[str]] = [set() for _ in range(n_shards)] |
| 246 | shard_sizes: list[int] = [0] * n_shards |
| 247 | |
| 248 | # Sort components largest-first for better balance (LPT heuristic). |
| 249 | sorted_comps = sorted( |
| 250 | components, |
| 251 | key=lambda c: sum(sym_counts.get(f, 0) for f in c), |
| 252 | reverse=True, |
| 253 | ) |
| 254 | for comp in sorted_comps: |
| 255 | smallest = min(range(n_shards), key=lambda i: shard_sizes[i]) |
| 256 | shards[smallest].update(comp) |
| 257 | shard_sizes[smallest] += sum(sym_counts.get(f, 0) for f in comp) |
| 258 | |
| 259 | return [frozenset(s) for s in shards] |
| 260 | |
| 261 | |
| 262 | # ── CLI registration ────────────────────────────────────────────────────────── |
| 263 | |
| 264 | |
| 265 | def register(subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]") -> None: |
| 266 | """Register the ``shard`` subcommand on *subparsers* (under ``muse coord``). |
| 267 | |
| 268 | Wires all flags with their defaults, choices, and help text so that |
| 269 | ``--help`` output is accurate. Sets ``func`` to :func:`run`. |
| 270 | """ |
| 271 | parser = subparsers.add_parser( |
| 272 | "shard", |
| 273 | help="Partition the codebase into N low-coupling work zones for parallel agents.", |
| 274 | description=__doc__, |
| 275 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 276 | ) |
| 277 | parser.add_argument( |
| 278 | "--agents", "-n", |
| 279 | type=int, |
| 280 | default=4, |
| 281 | metavar="N", |
| 282 | help=f"Number of work zones ({_MIN_AGENTS}–{_MAX_AGENTS}, default: 4).", |
| 283 | ) |
| 284 | parser.add_argument( |
| 285 | "--commit", "-c", |
| 286 | dest="ref", |
| 287 | default=None, |
| 288 | metavar="REF", |
| 289 | help="Use this commit instead of HEAD.", |
| 290 | ) |
| 291 | parser.add_argument( |
| 292 | "--language", "-l", |
| 293 | default=None, |
| 294 | metavar="LANG", |
| 295 | help="Restrict to files of this language.", |
| 296 | ) |
| 297 | parser.add_argument( |
| 298 | "--format", "-f", |
| 299 | default="text", |
| 300 | dest="fmt", |
| 301 | choices=("text", "json"), |
| 302 | help="Output format: text (default) or json.", |
| 303 | ) |
| 304 | parser.add_argument( |
| 305 | "--json", |
| 306 | action="store_const", |
| 307 | const="json", |
| 308 | dest="fmt", |
| 309 | help="Shorthand for --format json.", |
| 310 | ) |
| 311 | parser.set_defaults(func=run) |
| 312 | |
| 313 | |
| 314 | # ── Command implementation ──────────────────────────────────────────────────── |
| 315 | |
| 316 | |
| 317 | def run(args: argparse.Namespace) -> None: |
| 318 | """Partition the codebase into N low-coupling work zones for parallel agents. |
| 319 | |
| 320 | Uses the import graph connectivity to group files that are tightly coupled |
| 321 | together and loosely coupled to other shards. Agents assigned to different |
| 322 | shards minimise merge conflicts. |
| 323 | |
| 324 | Execution order |
| 325 | --------------- |
| 326 | 1. **Validate inputs** — ``--agents`` range checked before any file I/O. |
| 327 | Failure exits :attr:`~muse.core.errors.ExitCode.USER_ERROR` (1) with a |
| 328 | message on *stderr* or compact JSON on *stdout* when ``--format json``. |
| 329 | 2. **Resolve repo** — :func:`~muse.core.repo.require_repo`. |
| 330 | 3. **Resolve commit** — :func:`~muse.core.store.resolve_commit_ref`. |
| 331 | 4. **Load symbol map** — :func:`~muse.plugins.code._query.symbols_for_snapshot`. |
| 332 | 5. **Build import graph** — :func:`_build_import_edges` reads each file's |
| 333 | object from the store and parses import symbols. |
| 334 | 6. **Partition** — :func:`_connected_components` then :func:`_greedy_partition`. |
| 335 | 7. **Emit output** — compact JSON or human-readable text. |
| 336 | |
| 337 | Security |
| 338 | -------- |
| 339 | * ``--agents`` is validated against :data:`_MIN_AGENTS` / :data:`_MAX_AGENTS` |
| 340 | before any I/O. |
| 341 | * ``--commit`` refs are sanitised by :func:`~muse.core.store.resolve_commit_ref`. |
| 342 | * ``--language`` values are passed through |
| 343 | :func:`~muse.core.validation.sanitize_display` before text output to |
| 344 | prevent ANSI injection. |
| 345 | * All file paths in text output pass through |
| 346 | :func:`~muse.core.validation.sanitize_display`. |
| 347 | * Error messages go to *stdout* as compact JSON when ``--format json``, or |
| 348 | to *stderr* with an ``❌`` prefix otherwise. |
| 349 | |
| 350 | Performance |
| 351 | ----------- |
| 352 | * Import-graph build is O(F × S) where F is the number of files in the |
| 353 | manifest and S is the average number of symbols per file. |
| 354 | * Connected-components DFS is O(F + E) where E is the number of import |
| 355 | edges. |
| 356 | * Greedy partition is O(C log C + C × N) where C is the number of |
| 357 | components and N the number of shards. |
| 358 | * Shards are balanced by symbol count using the LPT heuristic |
| 359 | (largest-processing-time first), which gives a ≤ 4/3 OPT approximation. |
| 360 | |
| 361 | Args: |
| 362 | args: Parsed ``argparse.Namespace`` with attributes ``agents``, |
| 363 | ``ref``, ``language``, and ``fmt``. |
| 364 | |
| 365 | Exit codes: |
| 366 | 0 — success (no files found is also success). |
| 367 | 1 — bad arguments or commit not found. |
| 368 | """ |
| 369 | as_json: bool = args.fmt == "json" |
| 370 | t0 = time.monotonic() |
| 371 | |
| 372 | # ── Input validation (before any file I/O) ──────────────────────────────── |
| 373 | |
| 374 | agents_raw: int = args.agents |
| 375 | if not (_MIN_AGENTS <= agents_raw <= _MAX_AGENTS): |
| 376 | msg = f"--agents must be {_MIN_AGENTS}–{_MAX_AGENTS}, got {agents_raw}" |
| 377 | _err(msg, as_json, "bad_args") |
| 378 | raise SystemExit(ExitCode.USER_ERROR) |
| 379 | |
| 380 | agents: int = agents_raw |
| 381 | ref: str | None = args.ref |
| 382 | language: str | None = args.language |
| 383 | |
| 384 | root = require_repo() |
| 385 | repo_id = read_repo_id(root) |
| 386 | branch = read_current_branch(root) |
| 387 | |
| 388 | commit = resolve_commit_ref(root, repo_id, branch, ref) |
| 389 | if commit is None: |
| 390 | msg = f"commit '{sanitize_display(ref or 'HEAD')}' not found" |
| 391 | _err(msg, as_json, "commit_not_found") |
| 392 | raise SystemExit(ExitCode.USER_ERROR) |
| 393 | |
| 394 | manifest = get_commit_snapshot_manifest(root, commit.commit_id) or {} |
| 395 | sym_map = symbols_for_snapshot(root, manifest, language_filter=language) |
| 396 | sym_counts: _CounterMap = {fp: len(tree) for fp, tree in sym_map.items()} |
| 397 | |
| 398 | files = sorted(sym_counts.keys()) |
| 399 | if not files: |
| 400 | elapsed = round(time.monotonic() - t0, 4) |
| 401 | if as_json: |
| 402 | print(json.dumps({ |
| 403 | "schema_version": __version__, |
| 404 | "commit": commit.commit_id[:8], |
| 405 | "full_commit_id": commit.commit_id, |
| 406 | "agents": agents, |
| 407 | "shards_created": 0, |
| 408 | "total_files": 0, |
| 409 | "total_symbols": 0, |
| 410 | "cross_shard_edges": 0, |
| 411 | "shards": [], |
| 412 | "elapsed_seconds": elapsed, |
| 413 | })) |
| 414 | else: |
| 415 | print(" (no semantic files found in snapshot)") |
| 416 | return |
| 417 | |
| 418 | edges = _build_import_edges(root, manifest, language) |
| 419 | components = _connected_components(files, edges) |
| 420 | n = min(agents, len(components)) # Can't have more shards than components. |
| 421 | shard_sets = _greedy_partition(components, sym_counts, n) |
| 422 | |
| 423 | # Compute cross-shard edges. |
| 424 | file_to_shard: _ShardMap = {} |
| 425 | for i, s in enumerate(shard_sets): |
| 426 | for f in s: |
| 427 | file_to_shard[f] = i |
| 428 | cross_edges = sum( |
| 429 | 1 for a, b in edges |
| 430 | if file_to_shard.get(a, -1) != file_to_shard.get(b, -2) |
| 431 | ) |
| 432 | |
| 433 | total_symbols = sum(sym_counts.values()) |
| 434 | elapsed = round(time.monotonic() - t0, 4) |
| 435 | |
| 436 | if as_json: |
| 437 | print(json.dumps({ |
| 438 | "schema_version": __version__, |
| 439 | "commit": commit.commit_id[:8], |
| 440 | "full_commit_id": commit.commit_id, |
| 441 | "agents": agents, |
| 442 | "shards_created": n, |
| 443 | "total_files": len(files), |
| 444 | "total_symbols": total_symbols, |
| 445 | "cross_shard_edges": cross_edges, |
| 446 | "shards": [ |
| 447 | { |
| 448 | "shard": i + 1, |
| 449 | "files": sorted(s), |
| 450 | "symbol_count": sum(sym_counts.get(f, 0) for f in s), |
| 451 | "coupling_score": sum( |
| 452 | 1 for a, b in edges |
| 453 | if (a in s) != (b in s) |
| 454 | ), |
| 455 | } |
| 456 | for i, s in enumerate(shard_sets) if s |
| 457 | ], |
| 458 | "elapsed_seconds": elapsed, |
| 459 | })) |
| 460 | return |
| 461 | |
| 462 | # ── Text output ─────────────────────────────────────────────────────────── |
| 463 | print(f"\nShard plan — {n} agent(s), commit {commit.commit_id[:8]}") |
| 464 | if language: |
| 465 | print(f" (language: {sanitize_display(language)})") |
| 466 | print("─" * 62) |
| 467 | |
| 468 | for i, s in enumerate(shard_sets): |
| 469 | if not s: |
| 470 | continue |
| 471 | sym_total = sum(sym_counts.get(f, 0) for f in s) |
| 472 | coupling = sum(1 for a, b in edges if (a in s) != (b in s)) |
| 473 | print(f"\n Shard {i + 1} ({sym_total} symbols, {len(s)} files, coupling {coupling}):") |
| 474 | for fp in sorted(s): |
| 475 | print(f" {sanitize_display(fp)}") |
| 476 | |
| 477 | print(f"\n Cross-shard edges: {cross_edges}") |
| 478 | if cross_edges == 0: |
| 479 | print(" ✅ Perfect isolation — no cross-shard dependencies") |
| 480 | print(f"\n ({elapsed:.3f}s)") |