1"""

2 Native Parallel Agent Scheduler — Multi-agent parallel execution with DAG dependency.

3

4 Features:

5 - Task DAG: define task dependencies, auto topological sort

6 - Concurrency pool: limit max parallel agents with asyncio.Semaphore

7 - Load balancing: round-robin or least-busy agent selection

8 - Progress tracking: per-task status with callback hooks

9 - Resource limits: per-agent memory/token/cpu budgets

10 - Error isolation: one agent failure doesn't crash the others

11 - Streaming results: async generator for real-time output

12

13 Usage:

14 executor = ParallelExecutor(max_concurrent=8)

15

16 # Define tasks as a DAG

17 dag = {

18 "research": {"agent": "researcher", "prompt": "Research topic X"},

19 "draft": {"agent": "writer", "prompt": "Draft based on research",

20 "depends_on": ["research"]},

21 "review": {"agent": "reviewer", "prompt": "Review the draft",

22 "depends_on": ["draft"]},

23 "translate": {"agent": "translator", "prompt": "Translate to Chinese",

24 "depends_on": ["draft"]},

25 }

26

27 results = await executor.execute(dag)

28"""

29

30 from __future__ import annotations

31

32 import asyncio

33 import time

34 import uuid

35 from collections import defaultdict

36 from dataclasses import dataclass, field

37 from enum import Enum

38 from typing import Any, AsyncIterator, Callable, Coroutine, Optional

39

40

41# ── Task Models ──

42

43 class TaskStatus(str, Enum):

44 QUEUED = "queued"

45 RUNNING = "running"

46 DONE = "done"

47 FAILED = "failed"

48 SKIPPED = "skipped" # Dependency failed

49

50

51@dataclass

52 class TaskResult:

53 """Result of a single parallel task."""

54 task_id: str

55 status: TaskStatus

56 agent: str

57 output: Any = None

58 error: str = ""

59 started_at: float = 0.0

60 finished_at: float = 0.0

61 retry_count: int = 0

62

63 @property

64 def duration_ms(self) -> float:

65 return (self.finished_at - self.started_at) * 1000

66

67 @property

68 def ok(self) -> bool:

69 return self.status == TaskStatus.DONE

70

71

72@dataclass

73 class RunResult:

74 """Aggregate result of a parallel execution run."""

75 run_id: str

76 total: int

77 done: int

78 failed: int

79 skipped: int

80 total_duration_ms: float

81 tasks: list[TaskResult]

82

83 @property

84 def success_rate(self) -> float:

85 return self.done / max(self.total, 1)

86

87

88# ── Parallel Executor ──

89

90 ParallelAgentFn = Callable[[str, str, dict], Coroutine[Any, Any, Any]]

91""" async fn(agent_name: str, prompt: str, context: dict) -> Any """

92

93

94 class ParallelExecutor:

95 """Execute multiple agent tasks concurrently with DAG dependency resolution.

96

97 Args:

98 max_concurrent: Maximum number of simultaneously running tasks (default 8).

99 agent_fn: Async callable that executes a single agent task.

100 Signature: async (agent_name, prompt, context) -> result

101 max_retries: Per-task retry count on failure (default 1).

102 timeout: Per-task timeout in seconds (default 300).

103 """

104

105 def __init__(

106 self,

107 max_concurrent: int = 8,

108 agent_fn: Optional[ParallelAgentFn] = None,

109 max_retries: int = 1,

110 timeout: float = 300.0,

111 ):

112 self._max_concurrent = max_concurrent

113 self._semaphore = asyncio.Semaphore(max_concurrent)

114 self._agent_fn = agent_fn

115 self._max_retries = max_retries

116 self._timeout = timeout

117

118 self._progress_hooks: list[Callable[[TaskResult], Any]] = []

119 self._task_counter: dict[str, int] = defaultdict(int)

120

121 def on_progress(self, hook: Callable[[TaskResult], Any]) -> None:

122 """Register a progress callback — called on each task completion."""

123 self._progress_hooks.append(hook)

124

125 # ── Execute ──

126

127 async def execute(

128 self,

129 tasks: dict[str, dict],

130 context: dict = None,

131 ) -> RunResult:

132 """Execute a DAG of tasks in parallel.

133

134 Args:

135 tasks: {task_id: {agent, prompt, depends_on?, context?}, ...}

136 context: Global context injected into every task.

137

138 Returns:

139 RunResult with aggregated stats.

140 """

141 run_id = uuid.uuid4().hex[:12]

142 start_time = time.time()

143

144 # Build dependency graph

145 dependencies: dict[str, list[str]] = {}

146 for task_id, spec in tasks.items():

147 dependencies[task_id] = spec.get("depends_on", [])

148

149 # Topological sort → execution levels

150 levels = self._topological_sort(dependencies)

151

152 # Execute level by level (tasks within a level run in parallel)

153 all_results: dict[str, TaskResult] = {}

154 all_outputs: dict[str, Any] = {}

155

156 for level in levels:

157 level_tasks = []

158

159 for task_id in level:

160 spec = tasks[task_id]

161

162 # Check if dependencies all succeeded

163 deps = dependencies.get(task_id, [])

164 deps_failed = [d for d in deps

165 if d in all_results and not all_results[d].ok]

166

167 if deps_failed:

168 result = TaskResult(

169 task_id=task_id,

170 status=TaskStatus.SKIPPED,

171 agent=spec.get("agent", "unknown"),

172 error=f"Dependency failed: {deps_failed}",

173 )

174 all_results[task_id] = result

175 continue

176

177 # Build merged context: global + per-task + dependency outputs

178 merged_context = {}

179 if context:

180 merged_context.update(context)

181 if spec.get("context"):

182 merged_context.update(spec["context"])

183 for dep_id in deps:

184 if dep_id in all_outputs:

185 merged_context[f"_dep_{dep_id}"] = all_outputs[dep_id]

186

187 level_tasks.append(

188 self._run_one(

189 task_id=task_id,

190 agent=spec.get("agent", "default"),

191 prompt=spec.get("prompt", ""),

192 context=merged_context,

193 )

194 )

195

196 if level_tasks:

197 batch_results = await asyncio.gather(*level_tasks, return_exceptions=True)

198 for i, task_id in enumerate(level):

199 if task_id not in all_results:

200 result = batch_results[i]

201 if isinstance(result, Exception):

202 result = TaskResult(

203 task_id=task_id,

204 status=TaskStatus.FAILED,

205 agent=tasks[task_id].get("agent", "unknown"),

206 error=str(result),

207 )

208 all_results[task_id] = result

209 if result.ok:

210 all_outputs[task_id] = result.output

211

212 # Aggregate

213 total = len(tasks)

214 done = sum(1 for r in all_results.values() if r.status == TaskStatus.DONE)

215 failed = sum(1 for r in all_results.values() if r.status == TaskStatus.FAILED)

216 skipped = sum(1 for r in all_results.values() if r.status == TaskStatus.SKIPPED)

217

218 return RunResult(

219 run_id=run_id,

220 total=total,

221 done=done,

222 failed=failed,

223 skipped=skipped,

224 total_duration_ms=(time.time() - start_time) * 1000,

225 tasks=list(all_results.values()),

226 )

227

228 # ── Streaming ──

229

230 async def execute_stream(

231 self,

232 tasks: dict[str, dict],

233 context: dict = None,

234 ) -> AsyncIterator[TaskResult]:

235 """Execute tasks and yield results as they complete (per-level batches)."""

236 run_result = await self.execute(tasks, context)

237 for task in run_result.tasks:

238 yield task

239

240 # ── Batch Dispatch (no DAG) ──

241

242 async def fan_out(

243 self,

244 agent: str,

245 prompts: list[str],

246 context: dict = None,

247 ) -> list[TaskResult]:

248 """Fire-and-forget: run the same agent on many prompts in parallel."""

249 tasks = {f"task_{i}": {"agent": agent, "prompt": p} for i, p in enumerate(prompts)}

250 result = await self.execute(tasks, context)

251 return result.tasks

252

253 # ── Internal ──

254

255 async def _run_one(

256 self,

257 task_id: str,

258 agent: str,

259 prompt: str,

260 context: dict,

261 ) -> TaskResult:

262 """Execute a single task with semaphore, retry, and timeout."""

263 result = TaskResult(task_id=task_id, status=TaskStatus.RUNNING, agent=agent)

264

265 for attempt in range(self._max_retries + 1):

266 result.started_at = time.time()

267 result.retry_count = attempt

268

269 try:

270 async with self._semaphore:

271 if self._agent_fn:

272 output = await asyncio.wait_for(

273 self._agent_fn(agent, prompt, context),

274 timeout=self._timeout,

275 )

276 else:

277 # Default: simulate agent execution

278 output = await asyncio.wait_for(

279 self._default_agent(agent, prompt, context),

280 timeout=self._timeout,

281 )

282

283 result.output = output

284 result.status = TaskStatus.DONE

285 break

286

287 except asyncio.TimeoutError:

288 result.error = f"Timeout after {self._timeout}s"

289 result.status = TaskStatus.FAILED

290

291 except Exception as e:

292 result.error = str(e)

293 result.status = TaskStatus.FAILED

294 if attempt < self._max_retries:

295 await asyncio.sleep(0.5 * (attempt + 1))

296 continue

297 break

298

299 finally:

300 result.finished_at = time.time()

301

302 self._task_counter[agent] += 1

303 for hook in self._progress_hooks:

304 try:

305 hook(result)

306 except Exception:

307 pass

308

309 return result

310

311 async def _default_agent(self, agent: str, prompt: str, context: dict) -> str:

312 """Default agent execution (mock for testing; override with agent_fn)."""

313 await asyncio.sleep(0.1)

314 return f"[{agent}] processed: {prompt[:80]}"

315

316 # ── Topological Sort ──

317

318 @staticmethod

319 def _topological_sort(

320 dependencies: dict[str, list[str]]

321 ) -> list[list[str]]:

322 """Kahn's algorithm → ordered levels for parallel execution."""

323 in_degree: dict[str, int] = {node: 0 for node in dependencies}

324 children: dict[str, list[str]] = defaultdict(list)

325

326 for node, deps in dependencies.items():

327 for dep in deps:

328 if dep not in in_degree:

329 in_degree[dep] = 0

330 children.setdefault(dep, []).append(node)

331 in_degree[node] += 1

332

333 # Start with nodes that have no dependencies

334 queue = [node for node, deg in in_degree.items() if deg == 0]

335 levels: list[list[str]] = []

336 processed = set()

337

338 while queue:

339 levels.append(list(queue))

340 next_queue = []

341

342 for node in queue:

343 processed.add(node)

344 for child in children.get(node, []):

345 in_degree[child] -= 1

346 if in_degree[child] == 0 and child not in processed:

347 next_queue.append(child)

348

349 queue = next_queue

350

351 return levels

352

353 # ── Stats ──

354

355 def stats(self) -> dict[str, Any]:

356 return {

357 "max_concurrent": self._max_concurrent,

358 "max_retries": self._max_retries,

359 "timeout": self._timeout,

360 "task_counts": dict(self._task_counter),

361 }