Coverage for agentos/orchestration/parallel.py: 32%
167 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 09:19 +0800
1"""
2Native Parallel Agent Scheduler — Multi-agent parallel execution with DAG dependency.
4Features:
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
13Usage:
14 executor = ParallelExecutor(max_concurrent=8)
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 }
27 results = await executor.execute(dag)
28"""
30from __future__ import annotations
32import asyncio
33import time
34import uuid
35from collections import defaultdict
36from collections.abc import AsyncIterator, Callable, Coroutine
37from dataclasses import dataclass
38from enum import StrEnum
39from typing import Any
41# ── Task Models ──
44class TaskStatus(StrEnum):
45 QUEUED = "queued"
46 RUNNING = "running"
47 DONE = "done"
48 FAILED = "failed"
49 SKIPPED = "skipped" # Dependency failed
52@dataclass
53class TaskResult:
54 """Result of a single parallel task."""
56 task_id: str
57 status: TaskStatus
58 agent: str
59 output: Any = None
60 error: str = ""
61 started_at: float = 0.0
62 finished_at: float = 0.0
63 retry_count: int = 0
65 @property
66 def duration_ms(self) -> float:
67 return (self.finished_at - self.started_at) * 1000
69 @property
70 def ok(self) -> bool:
71 return self.status == TaskStatus.DONE
74@dataclass
75class RunResult:
76 """Aggregate result of a parallel execution run."""
78 run_id: str
79 total: int
80 done: int
81 failed: int
82 skipped: int
83 total_duration_ms: float
84 tasks: list[TaskResult]
86 @property
87 def success_rate(self) -> float:
88 return self.done / max(self.total, 1)
91# ── Parallel Executor ──
93ParallelAgentFn = Callable[[str, str, dict], Coroutine[Any, Any, Any]]
94""" async fn(agent_name: str, prompt: str, context: dict) -> Any """
97class ParallelExecutor:
98 """Execute multiple agent tasks concurrently with DAG dependency resolution.
100 Args:
101 max_concurrent: Maximum number of simultaneously running tasks (default 8).
102 agent_fn: Async callable that executes a single agent task.
103 Signature: async (agent_name, prompt, context) -> result
104 max_retries: Per-task retry count on failure (default 1).
105 timeout: Per-task timeout in seconds (default 300).
106 """
108 def __init__(
109 self,
110 max_concurrent: int = 8,
111 agent_fn: ParallelAgentFn | None = None,
112 max_retries: int = 1,
113 timeout: float = 300.0,
114 ):
115 self._max_concurrent = max_concurrent
116 self._semaphore = asyncio.Semaphore(max_concurrent)
117 self._agent_fn = agent_fn
118 self._max_retries = max_retries
119 self._timeout = timeout
121 self._progress_hooks: list[Callable[[TaskResult], Any]] = []
122 self._task_counter: dict[str, int] = defaultdict(int)
124 def on_progress(self, hook: Callable[[TaskResult], Any]) -> None:
125 """Register a progress callback — called on each task completion."""
126 self._progress_hooks.append(hook)
128 # ── Execute ──
130 async def execute(
131 self,
132 tasks: dict[str, dict],
133 context: dict = None,
134 ) -> RunResult:
135 """Execute a DAG of tasks in parallel.
137 Args:
138 tasks: {task_id: {agent, prompt, depends_on?, context?}, ...}
139 context: Global context injected into every task.
141 Returns:
142 RunResult with aggregated stats.
143 """
144 run_id = uuid.uuid4().hex[:12]
145 start_time = time.time()
147 # Build dependency graph
148 dependencies: dict[str, list[str]] = {}
149 for task_id, spec in tasks.items():
150 dependencies[task_id] = spec.get("depends_on", [])
152 # Topological sort → execution levels
153 levels = self._topological_sort(dependencies)
155 # Execute level by level (tasks within a level run in parallel)
156 all_results: dict[str, TaskResult] = {}
157 all_outputs: dict[str, Any] = {}
159 for level in levels:
160 level_tasks = []
162 for task_id in level:
163 spec = tasks[task_id]
165 # Check if dependencies all succeeded
166 deps = dependencies.get(task_id, [])
167 deps_failed = [d for d in deps if d in all_results and not all_results[d].ok]
169 if deps_failed:
170 result = TaskResult(
171 task_id=task_id,
172 status=TaskStatus.SKIPPED,
173 agent=spec.get("agent", "unknown"),
174 error=f"Dependency failed: {deps_failed}",
175 )
176 all_results[task_id] = result
177 continue
179 # Build merged context: global + per-task + dependency outputs
180 merged_context = {}
181 if context:
182 merged_context.update(context)
183 if spec.get("context"):
184 merged_context.update(spec["context"])
185 for dep_id in deps:
186 if dep_id in all_outputs:
187 merged_context[f"_dep_{dep_id}"] = all_outputs[dep_id]
189 level_tasks.append(
190 self._run_one(
191 task_id=task_id,
192 agent=spec.get("agent", "default"),
193 prompt=spec.get("prompt", ""),
194 context=merged_context,
195 )
196 )
198 if level_tasks:
199 batch_results = await asyncio.gather(*level_tasks, return_exceptions=True)
200 for i, task_id in enumerate(level):
201 if task_id not in all_results:
202 result = batch_results[i]
203 if isinstance(result, Exception):
204 result = TaskResult(
205 task_id=task_id,
206 status=TaskStatus.FAILED,
207 agent=tasks[task_id].get("agent", "unknown"),
208 error=str(result),
209 )
210 all_results[task_id] = result
211 if result.ok:
212 all_outputs[task_id] = result.output
214 # Aggregate
215 total = len(tasks)
216 done = sum(1 for r in all_results.values() if r.status == TaskStatus.DONE)
217 failed = sum(1 for r in all_results.values() if r.status == TaskStatus.FAILED)
218 skipped = sum(1 for r in all_results.values() if r.status == TaskStatus.SKIPPED)
220 return RunResult(
221 run_id=run_id,
222 total=total,
223 done=done,
224 failed=failed,
225 skipped=skipped,
226 total_duration_ms=(time.time() - start_time) * 1000,
227 tasks=list(all_results.values()),
228 )
230 # ── Streaming ──
232 async def execute_stream(
233 self,
234 tasks: dict[str, dict],
235 context: dict = None,
236 ) -> AsyncIterator[TaskResult]:
237 """Execute tasks and yield results as they complete (per-level batches)."""
238 run_result = await self.execute(tasks, context)
239 for task in run_result.tasks:
240 yield task
242 # ── Batch Dispatch (no DAG) ──
244 async def fan_out(
245 self,
246 agent: str,
247 prompts: list[str],
248 context: dict = None,
249 ) -> list[TaskResult]:
250 """Fire-and-forget: run the same agent on many prompts in parallel."""
251 tasks = {f"task_{i}": {"agent": agent, "prompt": p} for i, p in enumerate(prompts)}
252 result = await self.execute(tasks, context)
253 return result.tasks
255 # ── Internal ──
257 async def _run_one(
258 self,
259 task_id: str,
260 agent: str,
261 prompt: str,
262 context: dict,
263 ) -> TaskResult:
264 """Execute a single task with semaphore, retry, and timeout."""
265 result = TaskResult(task_id=task_id, status=TaskStatus.RUNNING, agent=agent)
267 for attempt in range(self._max_retries + 1):
268 result.started_at = time.time()
269 result.retry_count = attempt
271 try:
272 async with self._semaphore:
273 if self._agent_fn:
274 output = await asyncio.wait_for(
275 self._agent_fn(agent, prompt, context),
276 timeout=self._timeout,
277 )
278 else:
279 # Default: simulate agent execution
280 output = await asyncio.wait_for(
281 self._default_agent(agent, prompt, context),
282 timeout=self._timeout,
283 )
285 result.output = output
286 result.status = TaskStatus.DONE
287 break
289 except TimeoutError:
290 result.error = f"Timeout after {self._timeout}s"
291 result.status = TaskStatus.FAILED
293 except Exception as e:
294 result.error = str(e)
295 result.status = TaskStatus.FAILED
296 if attempt < self._max_retries:
297 await asyncio.sleep(0.5 * (attempt + 1))
298 continue
299 break
301 finally:
302 result.finished_at = time.time()
304 self._task_counter[agent] += 1
305 for hook in self._progress_hooks:
306 try:
307 hook(result)
308 except Exception:
309 pass
311 return result
313 async def _default_agent(self, agent: str, prompt: str, context: dict) -> str:
314 """Default agent execution (mock for testing; override with agent_fn)."""
315 await asyncio.sleep(0.1)
316 return f"[{agent}] processed: {prompt[:80]}"
318 # ── Topological Sort ──
320 @staticmethod
321 def _topological_sort(dependencies: dict[str, list[str]]) -> 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)
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
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()
338 while queue:
339 levels.append(list(queue))
340 next_queue = []
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)
349 queue = next_queue
351 return levels
353 # ── Stats ──
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 }