Coverage for agentos/concurrent/parallel.py: 0%
182 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 01:26 +0800
1"""
2Async Parallel Execution Primitives — fan-out/fan-in for multi-agent tasks.
4Provides structured concurrency patterns for AgentOS:
5- TaskGroup: structured async task grouping with collective result/error handling
6- ParallelExecutor: fan-out/fan-in with timeout, cancellation, throttling
7- parallel_gather: await multiple coroutines with timeout & partial results
8- parallel_map: map a function over items concurrently with bounded parallelism
10Key features:
11- Structured concurrency (all-or-nothing or partial results)
12- Semaphore-based throttling (bounded parallelism)
13- Per-task timeout + global timeout
14- Result aggregation with success/failure tracking
15- Graceful cancellation propagation
16"""
18from __future__ import annotations
20import asyncio
21import time
22import uuid
23from collections.abc import Callable, Coroutine
24from dataclasses import dataclass, field
25from enum import StrEnum
26from typing import Any, TypeVar
28T = TypeVar("T")
29R = TypeVar("R")
32# ── Data Structures ──────────────────────────────────────────────
35class TaskStatus(StrEnum):
36 """Individual task execution status."""
38 PENDING = "pending"
39 RUNNING = "running"
40 COMPLETED = "completed"
41 FAILED = "failed"
42 TIMEOUT = "timeout"
43 CANCELLED = "cancelled"
46@dataclass
47class TaskResult:
48 """Result of a single parallel task."""
50 task_id: str
51 status: TaskStatus = TaskStatus.PENDING
52 result: Any = None
53 error: Exception | None = None
54 started_at: float = 0.0
55 finished_at: float = 0.0
56 duration_ms: float = 0.0
57 retries: int = 0
60@dataclass
61class GatherResult:
62 """Aggregated result from parallel_gather."""
64 results: list[TaskResult] = field(default_factory=list)
65 total: int = 0
66 completed: int = 0
67 failed: int = 0
68 timed_out: int = 0
69 cancelled: int = 0
70 total_duration_ms: float = 0.0
71 all_succeeded: bool = False
73 @property
74 def success_rate(self) -> float:
75 return self.completed / self.total if self.total > 0 else 0.0
77 def get_results(self) -> list[Any]:
78 """Extract successful result values."""
79 return [r.result for r in self.results if r.status == TaskStatus.COMPLETED]
81 def get_errors(self) -> list[tuple[str, Exception]]:
82 """Extract (task_id, error) pairs for failed tasks."""
83 return [(r.task_id, r.error) for r in self.results if r.error]
86# ── Semaphore-based Task Throttler ───────────────────────────────
89class TaskThrottler:
90 """
91 Bounded concurrency controller using asyncio.Semaphore.
93 Usage:
94 throttler = TaskThrottler(max_concurrent=5)
95 async with throttler:
96 await do_work()
97 """
99 def __init__(self, max_concurrent: int = 10):
100 if max_concurrent < 1:
101 raise ValueError("max_concurrent must be >= 1")
102 self.max_concurrent = max_concurrent
103 self._semaphore = asyncio.Semaphore(max_concurrent)
104 self._active = 0
105 self._peak = 0
107 @property
108 def active(self) -> int:
109 return self._active
111 @property
112 def peak(self) -> int:
113 return self._peak
115 async def __aenter__(self):
116 await self._semaphore.acquire()
117 self._active += 1
118 if self._active > self._peak:
119 self._peak = self._active
120 return self
122 async def __aexit__(self, *args):
123 self._active -= 1
124 self._semaphore.release()
127# ── Parallel Executor ────────────────────────────────────────────
130class ParallelExecutor:
131 """
132 Fan-out/fan-in executor for running multiple coroutines concurrently.
134 Supports structured concurrency: either wait for all (fail-fast or tolerant),
135 or collect partial results on timeout.
137 Usage:
138 executor = ParallelExecutor(max_concurrent=8, timeout=30.0)
139 result = await executor.gather([
140 agent1.run(task_a),
141 agent2.run(task_b),
142 agent3.run(task_c),
143 ])
144 print(f"{result.completed}/{result.total} succeeded")
145 """
147 def __init__(
148 self,
149 max_concurrent: int = 8,
150 timeout: float = 60.0,
151 fail_fast: bool = False,
152 ):
153 self.max_concurrent = max_concurrent
154 self.timeout = timeout
155 self.fail_fast = fail_fast
156 self.throttler = TaskThrottler(max_concurrent)
158 async def gather(
159 self,
160 coros: list[Coroutine],
161 timeout: float | None = None,
162 return_partial: bool = True,
163 ) -> GatherResult:
164 """
165 Execute multiple coroutines in parallel with bounded concurrency.
167 Args:
168 coros: List of coroutines to execute
169 timeout: Global timeout (overrides executor default)
170 return_partial: If True, return partial results on timeout; if False, raise
172 Returns:
173 GatherResult with aggregated status
174 """
175 effective_timeout = timeout if timeout is not None else self.timeout
176 total_start = time.monotonic()
178 results: list[TaskResult] = []
179 tasks: dict[str, asyncio.Task] = {}
181 if not coros:
182 return GatherResult(total=0)
184 async def _run_one(coro: Coroutine, task_id: str) -> None:
185 tr = TaskResult(task_id=task_id, status=TaskStatus.RUNNING, started_at=time.monotonic())
186 results.append(tr)
188 async with self.throttler:
189 try:
190 tr.result = await coro
191 tr.status = TaskStatus.COMPLETED
192 except asyncio.CancelledError:
193 tr.status = TaskStatus.CANCELLED
194 if self.fail_fast:
195 raise
196 except TimeoutError:
197 tr.status = TaskStatus.TIMEOUT
198 tr.error = TimeoutError(f"Task {task_id} timed out")
199 if self.fail_fast:
200 raise
201 except Exception as e:
202 tr.status = TaskStatus.FAILED
203 tr.error = e
204 if self.fail_fast:
205 raise
206 finally:
207 tr.finished_at = time.monotonic()
208 tr.duration_ms = (tr.finished_at - tr.started_at) * 1000
210 # Launch all tasks
211 for i, coro in enumerate(coros):
212 task_id = uuid.uuid4().hex[:10]
213 t = asyncio.create_task(_run_one(coro, task_id))
214 tasks[task_id] = t
216 # Wait with global timeout
217 try:
218 done, pending = await asyncio.wait(
219 tasks.values(),
220 timeout=effective_timeout,
221 return_when=(
222 asyncio.ALL_COMPLETED if not self.fail_fast else asyncio.FIRST_EXCEPTION
223 ),
224 )
226 # Cancel remaining on fail-fast
227 if pending and self.fail_fast:
228 for t in pending:
229 t.cancel()
231 # Handle timeout: cancel remaining if not return_partial
232 if pending and not return_partial:
233 for t in pending:
234 t.cancel()
235 raise TimeoutError(f"Gather timed out after {effective_timeout}s")
237 # Mark timed-out tasks
238 for t in pending:
239 t.cancel()
240 for tr in results:
241 if tr.status == TaskStatus.RUNNING:
242 tr.status = TaskStatus.TIMEOUT
243 tr.finished_at = time.monotonic()
244 tr.duration_ms = (tr.finished_at - tr.started_at) * 1000
246 except Exception:
247 for t in tasks.values():
248 if not t.done():
249 t.cancel()
250 raise
252 # Build aggregate result
253 total_duration = (time.monotonic() - total_start) * 1000
254 completed = sum(1 for r in results if r.status == TaskStatus.COMPLETED)
255 failed = sum(1 for r in results if r.status == TaskStatus.FAILED)
256 timed_out = sum(1 for r in results if r.status == TaskStatus.TIMEOUT)
257 cancelled = sum(1 for r in results if r.status == TaskStatus.CANCELLED)
259 return GatherResult(
260 results=results,
261 total=len(results),
262 completed=completed,
263 failed=failed,
264 timed_out=timed_out,
265 cancelled=cancelled,
266 total_duration_ms=total_duration,
267 all_succeeded=(completed == len(results)),
268 )
270 async def map(
271 self,
272 func: Callable[[T], Coroutine],
273 items: list[T],
274 timeout: float | None = None,
275 ) -> GatherResult:
276 """
277 Map an async function over a list of items with bounded concurrency.
279 Args:
280 func: Async function taking one item and returning a value
281 items: List of input items
282 timeout: Global timeout
284 Returns:
285 GatherResult with results
286 """
287 coros = [func(item) for item in items]
288 return await self.gather(coros, timeout=timeout)
291# ── Convenience Functions ────────────────────────────────────────
294async def parallel_gather(
295 *coros: Coroutine,
296 max_concurrent: int = 8,
297 timeout: float = 60.0,
298 return_partial: bool = True,
299) -> GatherResult:
300 """
301 Convenience function: await multiple coroutines in parallel.
303 Usage:
304 result = await parallel_gather(
305 fetch_url(url1),
306 fetch_url(url2),
307 fetch_url(url3),
308 max_concurrent=5, timeout=30.0,
309 )
310 for r in result.get_results():
311 print(r)
312 """
313 executor = ParallelExecutor(max_concurrent=max_concurrent, timeout=timeout)
314 return await executor.gather(list(coros), return_partial=return_partial)
317async def parallel_map(
318 func: Callable[[T], Coroutine],
319 items: list[T],
320 max_concurrent: int = 8,
321 timeout: float = 60.0,
322) -> GatherResult:
323 """
324 Convenience function: map async function over items concurrently.
326 Usage:
327 result = await parallel_map(process_document, documents, max_concurrent=4)
328 print(f"Processed {result.completed}/{result.total} docs")
329 """
330 executor = ParallelExecutor(max_concurrent=max_concurrent, timeout=timeout)
331 return await executor.map(func, items)
334# ── Fan-Out / Fan-In with Aggregation ────────────────────────────
337@dataclass
338class FanOutConfig:
339 """Configuration for fan-out pattern."""
341 max_concurrent: int = 8
342 timeout: float = 60.0
343 aggregation: str = "all" # "all" | "first" | "merge"
344 retry_failed: bool = False
345 max_retries: int = 2
348class FanOutExecutor:
349 """
350 Fan-out pattern: dispatch tasks to N workers, collect results.
352 Supports aggregation modes:
353 - "all": Wait for all, return list of results
354 - "first": Return first successful result (race)
355 - "merge": Run all, merge results with a merge function
356 """
358 def __init__(self, config: FanOutConfig | None = None, max_concurrent: int = 0):
359 if config is None and max_concurrent > 0:
360 config = FanOutConfig(max_concurrent=max_concurrent)
361 self.config = config or FanOutConfig()
362 self.executor = ParallelExecutor(
363 max_concurrent=self.config.max_concurrent,
364 timeout=self.config.timeout,
365 )
367 async def fan_out(
368 self,
369 worker_coros: list[Coroutine],
370 merge_fn: Callable[[list[Any]], Any] | None = None,
371 ) -> list[Any] | Any | GatherResult:
372 """
373 Fan out tasks to workers and collect results.
375 Args:
376 worker_coros: List of worker coroutines
377 merge_fn: Merge function for "merge" mode (list of results -> merged value)
379 Returns:
380 Depends on aggregation mode:
381 - "all": list of results
382 - "first": first successful result
383 - "merge": merged result via merge_fn
384 """
385 mode = self.config.aggregation
387 if mode == "first":
388 # Race: return first successful
389 gather_result = await self.executor.gather(
390 worker_coros,
391 return_partial=True,
392 )
393 successes = gather_result.get_results()
394 if successes:
395 return successes[0]
396 # All failed — raise first error
397 errors = gather_result.get_errors()
398 if errors:
399 raise errors[0][1]
400 raise RuntimeError("All workers failed with no result")
402 elif mode == "merge":
403 gather_result = await self.executor.gather(worker_coros)
404 if not merge_fn:
405 raise ValueError("merge_fn required for 'merge' mode")
406 return merge_fn(gather_result.get_results())
408 else: # "all"
409 gather_result = await self.executor.gather(worker_coros)
410 if self.config.retry_failed and gather_result.failed > 0:
411 # Retry failed tasks
412 for r in gather_result.results:
413 if r.status == TaskStatus.FAILED:
414 # Note: retry requires caller to provide a way to rebuild the coro
415 pass
416 return gather_result
419# ── Agent Loop Integration ────────────────────────────────────────
422def create_parallel_agent_gather(
423 max_concurrent: int = 8,
424 timeout: float = 60.0,
425) -> Callable:
426 """
427 Create a gather function for use in Agent tool definitions.
429 Usage:
430 agent_tools["parallel_gather"] = create_parallel_agent_gather(max_concurrent=5)
431 """
433 async def agent_gather(*coros: Coroutine) -> GatherResult:
434 return await parallel_gather(*coros, max_concurrent=max_concurrent, timeout=timeout)
436 return agent_gather