Coverage for agentos/concurrent/parallel.py: 0%
180 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 12:29 +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 async def _run_one(coro: Coroutine, task_id: str) -> None:
182 tr = TaskResult(task_id=task_id, status=TaskStatus.RUNNING, started_at=time.monotonic())
183 results.append(tr)
185 async with self.throttler:
186 try:
187 tr.result = await coro
188 tr.status = TaskStatus.COMPLETED
189 except asyncio.CancelledError:
190 tr.status = TaskStatus.CANCELLED
191 if self.fail_fast:
192 raise
193 except TimeoutError:
194 tr.status = TaskStatus.TIMEOUT
195 tr.error = TimeoutError(f"Task {task_id} timed out")
196 if self.fail_fast:
197 raise
198 except Exception as e:
199 tr.status = TaskStatus.FAILED
200 tr.error = e
201 if self.fail_fast:
202 raise
203 finally:
204 tr.finished_at = time.monotonic()
205 tr.duration_ms = (tr.finished_at - tr.started_at) * 1000
207 # Launch all tasks
208 for i, coro in enumerate(coros):
209 task_id = uuid.uuid4().hex[:10]
210 t = asyncio.create_task(_run_one(coro, task_id))
211 tasks[task_id] = t
213 # Wait with global timeout
214 try:
215 done, pending = await asyncio.wait(
216 tasks.values(),
217 timeout=effective_timeout,
218 return_when=(
219 asyncio.ALL_COMPLETED if not self.fail_fast else asyncio.FIRST_EXCEPTION
220 ),
221 )
223 # Cancel remaining on fail-fast
224 if pending and self.fail_fast:
225 for t in pending:
226 t.cancel()
228 # Handle timeout: cancel remaining if not return_partial
229 if pending and not return_partial:
230 for t in pending:
231 t.cancel()
232 raise TimeoutError(f"Gather timed out after {effective_timeout}s")
234 # Mark timed-out tasks
235 for t in pending:
236 t.cancel()
237 for tr in results:
238 if tr.status == TaskStatus.RUNNING:
239 tr.status = TaskStatus.TIMEOUT
240 tr.finished_at = time.monotonic()
241 tr.duration_ms = (tr.finished_at - tr.started_at) * 1000
243 except Exception:
244 for t in tasks.values():
245 if not t.done():
246 t.cancel()
247 raise
249 # Build aggregate result
250 total_duration = (time.monotonic() - total_start) * 1000
251 completed = sum(1 for r in results if r.status == TaskStatus.COMPLETED)
252 failed = sum(1 for r in results if r.status == TaskStatus.FAILED)
253 timed_out = sum(1 for r in results if r.status == TaskStatus.TIMEOUT)
254 cancelled = sum(1 for r in results if r.status == TaskStatus.CANCELLED)
256 return GatherResult(
257 results=results,
258 total=len(results),
259 completed=completed,
260 failed=failed,
261 timed_out=timed_out,
262 cancelled=cancelled,
263 total_duration_ms=total_duration,
264 all_succeeded=(completed == len(results)),
265 )
267 async def map(
268 self,
269 func: Callable[[T], Coroutine],
270 items: list[T],
271 timeout: float | None = None,
272 ) -> GatherResult:
273 """
274 Map an async function over a list of items with bounded concurrency.
276 Args:
277 func: Async function taking one item and returning a value
278 items: List of input items
279 timeout: Global timeout
281 Returns:
282 GatherResult with results
283 """
284 coros = [func(item) for item in items]
285 return await self.gather(coros, timeout=timeout)
288# ── Convenience Functions ────────────────────────────────────────
291async def parallel_gather(
292 *coros: Coroutine,
293 max_concurrent: int = 8,
294 timeout: float = 60.0,
295 return_partial: bool = True,
296) -> GatherResult:
297 """
298 Convenience function: await multiple coroutines in parallel.
300 Usage:
301 result = await parallel_gather(
302 fetch_url(url1),
303 fetch_url(url2),
304 fetch_url(url3),
305 max_concurrent=5, timeout=30.0,
306 )
307 for r in result.get_results():
308 print(r)
309 """
310 executor = ParallelExecutor(max_concurrent=max_concurrent, timeout=timeout)
311 return await executor.gather(list(coros), return_partial=return_partial)
314async def parallel_map(
315 func: Callable[[T], Coroutine],
316 items: list[T],
317 max_concurrent: int = 8,
318 timeout: float = 60.0,
319) -> GatherResult:
320 """
321 Convenience function: map async function over items concurrently.
323 Usage:
324 result = await parallel_map(process_document, documents, max_concurrent=4)
325 print(f"Processed {result.completed}/{result.total} docs")
326 """
327 executor = ParallelExecutor(max_concurrent=max_concurrent, timeout=timeout)
328 return await executor.map(func, items)
331# ── Fan-Out / Fan-In with Aggregation ────────────────────────────
334@dataclass
335class FanOutConfig:
336 """Configuration for fan-out pattern."""
338 max_concurrent: int = 8
339 timeout: float = 60.0
340 aggregation: str = "all" # "all" | "first" | "merge"
341 retry_failed: bool = False
342 max_retries: int = 2
345class FanOutExecutor:
346 """
347 Fan-out pattern: dispatch tasks to N workers, collect results.
349 Supports aggregation modes:
350 - "all": Wait for all, return list of results
351 - "first": Return first successful result (race)
352 - "merge": Run all, merge results with a merge function
353 """
355 def __init__(self, config: FanOutConfig | None = None, max_concurrent: int = 0):
356 if config is None and max_concurrent > 0:
357 config = FanOutConfig(max_concurrent=max_concurrent)
358 self.config = config or FanOutConfig()
359 self.executor = ParallelExecutor(
360 max_concurrent=self.config.max_concurrent,
361 timeout=self.config.timeout,
362 )
364 async def fan_out(
365 self,
366 worker_coros: list[Coroutine],
367 merge_fn: Callable[[list[Any]], Any] | None = None,
368 ) -> list[Any] | Any | GatherResult:
369 """
370 Fan out tasks to workers and collect results.
372 Args:
373 worker_coros: List of worker coroutines
374 merge_fn: Merge function for "merge" mode (list of results -> merged value)
376 Returns:
377 Depends on aggregation mode:
378 - "all": list of results
379 - "first": first successful result
380 - "merge": merged result via merge_fn
381 """
382 mode = self.config.aggregation
384 if mode == "first":
385 # Race: return first successful
386 gather_result = await self.executor.gather(
387 worker_coros,
388 return_partial=True,
389 )
390 successes = gather_result.get_results()
391 if successes:
392 return successes[0]
393 # All failed — raise first error
394 errors = gather_result.get_errors()
395 if errors:
396 raise errors[0][1]
397 raise RuntimeError("All workers failed with no result")
399 elif mode == "merge":
400 gather_result = await self.executor.gather(worker_coros)
401 if not merge_fn:
402 raise ValueError("merge_fn required for 'merge' mode")
403 return merge_fn(gather_result.get_results())
405 else: # "all"
406 gather_result = await self.executor.gather(worker_coros)
407 if self.config.retry_failed and gather_result.failed > 0:
408 # Retry failed tasks
409 for r in gather_result.results:
410 if r.status == TaskStatus.FAILED:
411 # Note: retry requires caller to provide a way to rebuild the coro
412 pass
413 return gather_result
416# ── Agent Loop Integration ────────────────────────────────────────
419def create_parallel_agent_gather(
420 max_concurrent: int = 8,
421 timeout: float = 60.0,
422) -> Callable:
423 """
424 Create a gather function for use in Agent tool definitions.
426 Usage:
427 agent_tools["parallel_gather"] = create_parallel_agent_gather(max_concurrent=5)
428 """
430 async def agent_gather(*coros: Coroutine) -> GatherResult:
431 return await parallel_gather(*coros, max_concurrent=max_concurrent, timeout=timeout)
433 return agent_gather