aaronrene / gabriel-muse public
shard.py python
480 lines 16.8 KB
Raw
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 32 days ago
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)")
File History 2 commits
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 32 days ago
sha256:d8316ffae901be06347e16ab55be11868eb519dd16ade3e8aa16a99e662f7e62 baseline: rc14 re-baseline after rc3 store corruption recovery Human patch 32 days ago