Coverage for src/keel/swarm_runtime.py: 100%
100 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 12:05 +0000
1"""Keel Swarm Runtime — Isolated multi-worktree execution & cluster orchestration.
3Thin I/O execution layer for running parallel swarm workers in isolated Git worktrees,
4dispatching keel ship jobs, handling worker state persistence, and managing fail-soft
5rebalancing across waves.
6"""
8from __future__ import annotations
10import concurrent.futures
11import datetime
12import shutil
13import subprocess # nosec B404
14import sys
15from collections.abc import Callable
16from pathlib import Path
17from typing import Any
19from .runner import CommandResult
20from .swarm import (
21 SwarmPlan,
22 SwarmRunResult,
23 SwarmRunState,
24 SwarmWorkerStatus,
25 rebalance_swarm_plan,
26 save_swarm_state,
27 update_worker_state,
28)
30SubprocessRunner = Callable[[list[str], Path], CommandResult]
33def default_runner(cmd: list[str], cwd: Path) -> CommandResult:
34 """Run a subprocess command in cwd and return a CommandResult."""
35 try:
36 proc = subprocess.run( # nosec B603
37 cmd,
38 cwd=cwd,
39 stdout=subprocess.PIPE,
40 stderr=subprocess.STDOUT,
41 text=True,
42 timeout=300,
43 check=False,
44 )
45 return CommandResult(
46 ok=(proc.returncode == 0),
47 code=proc.returncode,
48 output=proc.stdout or "",
49 timed_out=False,
50 stdout=proc.stdout or "",
51 stderr="",
52 )
53 except subprocess.TimeoutExpired as exc:
54 out = (
55 exc.output
56 if isinstance(exc.output, str)
57 else (exc.stdout if isinstance(exc.stdout, str) else "")
58 )
59 return CommandResult(
60 ok=False,
61 code=124,
62 output=out or "",
63 timed_out=True,
64 stdout=out or "",
65 stderr="",
66 )
67 except Exception as exc: # noqa: BLE001
68 return CommandResult(
69 ok=False,
70 code=1,
71 output=str(exc),
72 timed_out=False,
73 stdout="",
74 stderr=str(exc),
75 )
78def build_worktree_path(swarm_id: str, cluster_id: str, root: str | Path = ".") -> Path:
79 """Return the isolated worktree filesystem path for a specific swarm cluster worker."""
80 return Path(root) / ".keel" / "worktrees" / swarm_id / cluster_id
83def create_swarm_worktree(
84 repo_root: Path,
85 worktree_path: Path,
86 branch_name: str,
87 base_branch: str = "main",
88 runner: SubprocessRunner | None = None,
89) -> bool:
90 """Create an isolated git worktree for a swarm cluster worker."""
91 worktree_path.parent.mkdir(parents=True, exist_ok=True)
92 run = runner or default_runner
93 cmd = [
94 "git",
95 "worktree",
96 "add",
97 "-B",
98 branch_name,
99 str(worktree_path),
100 base_branch,
101 ]
102 res = run(cmd, repo_root)
103 return res.ok
106def remove_swarm_worktree(
107 repo_root: Path,
108 worktree_path: Path,
109 runner: SubprocessRunner | None = None,
110) -> bool:
111 """Remove a previously created isolated git worktree."""
112 run = runner or default_runner
113 cmd = ["git", "worktree", "remove", "--force", str(worktree_path)]
114 res = run(cmd, repo_root)
115 if not res.ok and worktree_path.exists():
116 shutil.rmtree(worktree_path, ignore_errors=True)
117 return True
120def execute_cluster_worker(
121 project_yaml: str,
122 issue: int,
123 root: Path,
124 worktree_dir: Path,
125 *,
126 dry_run: bool = True,
127 role: str = "core",
128 extra_args: list[str] | None = None,
129 runner: SubprocessRunner | None = None,
130) -> dict[str, Any]:
131 """Execute a single cluster worker pipeline (keel ship) within its worktree."""
132 run = runner or default_runner
133 cmd = [
134 sys.executable,
135 "-m",
136 "keel",
137 "ship",
138 project_yaml,
139 "--root",
140 str(worktree_dir if worktree_dir.exists() else root),
141 "--issue",
142 str(issue),
143 "--json",
144 ]
145 if dry_run:
146 cmd.append("--dry-run")
147 if extra_args:
148 cmd.extend(extra_args)
150 target_cwd = worktree_dir if worktree_dir.exists() else root
151 res = run(cmd, target_cwd)
153 return {
154 "issue": issue,
155 "role": role,
156 "ok": res.ok,
157 "code": res.code,
158 "output": res.stdout or res.output,
159 }
162def run_swarm_orchestration(
163 plan: SwarmPlan,
164 project_yaml: str,
165 *,
166 root: str | Path = ".",
167 dry_run: bool = True,
168 max_workers: int = 4,
169 runner: SubprocessRunner | None = None,
170 create_worktrees: bool = True,
171) -> SwarmRunResult:
172 """Execute the waves and clusters of a SwarmPlan with fail-soft isolation."""
173 root_path = Path(root).resolve()
174 workers_list: list[SwarmWorkerStatus] = []
176 # Initialize workers
177 for w in plan.waves:
178 for c in w.clusters:
179 issue_num = c.issues[0] if c.issues else 0
180 workers_list.append(
181 SwarmWorkerStatus(
182 cluster_id=c.cluster_id,
183 issue=issue_num,
184 role=c.role,
185 step="s0",
186 status="queued",
187 updated_at=datetime.datetime.now(datetime.UTC).isoformat(),
188 )
189 )
191 state = SwarmRunState(
192 swarm_id=plan.swarm_id,
193 total_workers=len(workers_list),
194 active_wave=1,
195 workers=tuple(workers_list),
196 started_at=datetime.datetime.now(datetime.UTC).isoformat(),
197 )
198 save_swarm_state(state, root=root_path)
200 passed_count = 0
201 failed_count = 0
202 wave_results: list[dict[str, Any]] = []
203 current_plan = plan
205 for wave in current_plan.waves:
206 state = SwarmRunState(
207 swarm_id=state.swarm_id,
208 total_workers=state.total_workers,
209 active_wave=wave.wave_index,
210 workers=state.workers,
211 started_at=state.started_at,
212 )
214 cluster_tasks = list(wave.clusters)
215 if not cluster_tasks:
216 continue
218 wave_record: dict[str, Any] = {
219 "wave_index": wave.wave_index,
220 "mode": wave.mode,
221 "eligible_direct_landing": wave.eligible_direct_landing,
222 "cluster_results": {},
223 }
225 # Mark clusters in this wave as running
226 for c in cluster_tasks:
227 state = update_worker_state(
228 state, c.cluster_id, step="s4", status="running", details="executing ship pipeline"
229 )
230 save_swarm_state(state, root=root_path)
232 def _worker_fn(cluster: Any) -> tuple[str, dict[str, Any]]:
233 c_id = cluster.cluster_id
234 issue_n = cluster.issues[0] if cluster.issues else 0
235 wt_path = build_worktree_path(plan.swarm_id, c_id, root=root_path)
237 if create_worktrees and not dry_run:
238 branch_name = f"swarm/{plan.swarm_id}/{c_id}"
239 create_swarm_worktree(root_path, wt_path, branch_name, runner=runner)
241 res = execute_cluster_worker(
242 project_yaml=project_yaml,
243 issue=issue_n,
244 root=root_path,
245 worktree_dir=wt_path,
246 dry_run=dry_run,
247 role=cluster.role,
248 runner=runner,
249 )
251 if create_worktrees and not dry_run:
252 remove_swarm_worktree(root_path, wt_path, runner=runner)
254 return c_id, res
256 # Run wave clusters in parallel thread pool
257 pool_workers = min(max_workers, len(cluster_tasks)) if len(cluster_tasks) > 0 else 1
258 with concurrent.futures.ThreadPoolExecutor(max_workers=pool_workers) as executor:
259 future_to_cluster = {
260 executor.submit(_worker_fn, cluster): cluster for cluster in cluster_tasks
261 }
262 for future in concurrent.futures.as_completed(future_to_cluster):
263 c_id, worker_res = future.result()
264 wave_record["cluster_results"][c_id] = worker_res
265 issue_val = worker_res.get("issue", 0)
267 if worker_res.get("ok", False):
268 passed_count += 1
269 state = update_worker_state(
270 state, c_id, step="s10", status="passed", details="pipeline completed"
271 )
272 else:
273 failed_count += 1
274 state = update_worker_state(
275 state,
276 c_id,
277 step="s4",
278 status="failed",
279 details=worker_res.get("output", ""),
280 )
281 # Dynamically rebalance subsequent waves if needed
282 current_plan = rebalance_swarm_plan(current_plan, issue_val)
284 save_swarm_state(state, root=root_path)
286 wave_results.append(wave_record)
288 # Finalize state
289 overall_status = (
290 "success"
291 if failed_count == 0 and passed_count > 0
292 else ("partial_failure" if passed_count > 0 else "failed")
293 )
294 if passed_count == 0 and failed_count == 0:
295 overall_status = "success"
297 state = SwarmRunState(
298 swarm_id=state.swarm_id,
299 total_workers=state.total_workers,
300 active_wave=state.active_wave,
301 workers=state.workers,
302 started_at=state.started_at,
303 completed_at=datetime.datetime.now(datetime.UTC).isoformat(),
304 )
305 save_swarm_state(state, root=root_path)
307 return SwarmRunResult(
308 swarm_id=plan.swarm_id,
309 status=overall_status,
310 total_workers=len(workers_list),
311 passed_count=passed_count,
312 failed_count=failed_count,
313 dry_run=dry_run,
314 wave_results=tuple(wave_results),
315 )