Coverage for agentos/concurrency/batch.py: 55%

1""" 

2AsyncBatchExecutor — Concurrent agent task dispatch with configurable 

3parallelism, timeout, retry, and result aggregation. 

4 

5Designed for running multiple AgentOS tasks in parallel (e.g., batch 

6evaluation, multi-model comparison, bulk processing). 

7""" 

8 

9from __future__ import annotations 

10 

11import asyncio 

12import logging 

13import time 

14from collections.abc import Awaitable, Callable 

15from dataclasses import dataclass, field 

16from enum import Enum 

17from typing import Any 

18 

19logger = logging.getLogger(__name__) 

20 

21 

22class TaskStatus(Enum): 

23 """任务状态枚举。""" 

24 

25 PENDING = "pending" 

26 RUNNING = "running" 

27 SUCCESS = "success" 

28 FAILED = "failed" 

29 TIMEOUT = "timeout" 

30 RETRYING = "retrying" 

31 CANCELLED = "cancelled" 

32 

33 

34class BatchStrategy(Enum): 

35 """Execution strategy for batch tasks.""" 

36 

37 PARALLEL = "parallel" # All tasks run concurrently (limited by max_concurrency) 

38 SEQUENTIAL = "sequential" # One after another 

39 SMART = "smart" # Dynamically adjust based on system load 

40 

41 

42@dataclass 

43class TaskSpec: 

44 """Specification for a single task in a batch.""" 

45 

46 task_id: str 

47 coro_or_func: Callable[..., Awaitable[Any]] 

48 args: tuple = () 

49 kwargs: dict[str, Any] = field(default_factory=dict) 

50 timeout: float = 60.0 

51 max_retries: int = 0 

52 metadata: dict[str, Any] = field(default_factory=dict) 

53 

54 

55@dataclass 

56class TaskResult: 

57 """Result of a single task execution.""" 

58 

59 task_id: str 

60 status: TaskStatus 

61 result: Any = None 

62 error: str | None = None 

63 duration_ms: float = 0.0 

64 retries: int = 0 

65 started_at: float = 0.0 

66 finished_at: float = 0.0 

67 

68 @property 

69 def success(self) -> bool: 

70 return self.status == TaskStatus.SUCCESS 

71 

72 

73@dataclass 

74class BatchConfig: 

75 """Configuration for AsyncBatchExecutor.""" 

76 

77 max_concurrency: int = 5 

78 default_timeout: float = 60.0 

79 max_retries: int = 1 

80 retry_delay: float = 1.0 

81 strategy: BatchStrategy = BatchStrategy.PARALLEL 

82 fail_fast: bool = False 

83 collect_errors: bool = True 

84 

85 

86@dataclass 

87class BatchResult: 

88 """Aggregated result of a batch execution.""" 

89 

90 results: list[TaskResult] = field(default_factory=list) 

91 total: int = 0 

92 succeeded: int = 0 

93 failed: int = 0 

94 timed_out: int = 0 

95 total_duration_ms: float = 0.0 

96 started_at: float = 0.0 

97 finished_at: float = 0.0 

98 

99 @property 

100 def success_rate(self) -> float: 

101 if self.total == 0: 

102 return 0.0 

103 return self.succeeded / self.total 

104 

105 @property 

106 def all_success(self) -> bool: 

107 return self.succeeded == self.total 

108 

109 def get_failed_ids(self) -> list[str]: 

110 return [ 

111 r.task_id for r in self.results if r.status in (TaskStatus.FAILED, TaskStatus.TIMEOUT) 

112 ] 

113 

114 

115class AsyncBatchExecutor: 

116 """Concurrently dispatches multiple AgentOS tasks and aggregates results.""" 

117 

118 def __init__(self, config: BatchConfig | None = None): 

119 self.config = config or BatchConfig() 

120 self._semaphore: asyncio.Semaphore | None = None 

121 self._cancel_event: asyncio.Event | None = None 

122 

123 async def execute(self, tasks: list[TaskSpec]) -> BatchResult: 

124 """Execute a list of tasks and return aggregated results.""" 

125 if not tasks: 

126 return BatchResult(total=0) 

127 

128 start = time.perf_counter() 

129 self._semaphore = asyncio.Semaphore(self.config.max_concurrency) 

130 self._cancel_event = asyncio.Event() 

131 results: list[TaskResult] = [] 

132 

133 if self.config.strategy == BatchStrategy.SEQUENTIAL: 

134 for task in tasks: 

135 result = await self._execute_one(task) 

136 results.append(result) 

137 if self.config.fail_fast and not result.success: 

138 break 

139 else: 

140 # PARALLEL or SMART 

141 tasks_coros = [self._execute_one(task) for task in tasks] 

142 results = list(await asyncio.gather(*tasks_coros)) 

143 

144 elapsed = (time.perf_counter() - start) * 1000 

145 succeeded = sum(1 for r in results if r.status == TaskStatus.SUCCESS) 

146 failed = sum(1 for r in results if r.status == TaskStatus.FAILED) 

147 timed_out = sum(1 for r in results if r.status == TaskStatus.TIMEOUT) 

148 

149 return BatchResult( 

150 results=results, 

151 total=len(tasks), 

152 succeeded=succeeded, 

153 failed=failed, 

154 timed_out=timed_out, 

155 total_duration_ms=elapsed, 

156 started_at=start, 

157 finished_at=time.perf_counter(), 

158 ) 

159 

160 async def _execute_one(self, task: TaskSpec) -> TaskResult: 

161 """Execute a single task with retry support.""" 

162 timeout = task.timeout if task.timeout is not None else self.config.default_timeout 

163 max_retries = task.max_retries if task.max_retries is not None else self.config.max_retries 

164 retries = 0 

165 

166 assert self._semaphore is not None, "semaphore must be set before calling execute()" 

167 async with self._semaphore: 

168 while True: 

169 started = time.perf_counter() 

170 try: 

171 coro = task.coro_or_func(*task.args, **task.kwargs) 

172 result_value = await asyncio.wait_for(coro, timeout=timeout) 

173 elapsed = (time.perf_counter() - started) * 1000 

174 return TaskResult( 

175 task_id=task.task_id, 

176 status=TaskStatus.SUCCESS, 

177 result=result_value, 

178 duration_ms=elapsed, 

179 retries=retries, 

180 started_at=started, 

181 finished_at=time.perf_counter(), 

182 ) 

183 except TimeoutError: 

184 if retries < max_retries: 

185 retries += 1 

186 logger.warning( 

187 f"Task '{task.task_id}' timed out (attempt {retries}/{max_retries}), retrying..." 

188 ) 

189 await asyncio.sleep(self.config.retry_delay) 

190 continue 

191 elapsed = (time.perf_counter() - started) * 1000 

192 return TaskResult( 

193 task_id=task.task_id, 

194 status=TaskStatus.TIMEOUT, 

195 error=f"Timed out after {timeout}s (retries: {retries})", 

196 duration_ms=elapsed, 

197 retries=retries, 

198 started_at=started, 

199 finished_at=time.perf_counter(), 

200 ) 

201 except Exception as e: 

202 if retries < max_retries: 

203 retries += 1 

204 logger.warning( 

205 f"Task '{task.task_id}' failed (attempt {retries}/{max_retries}): {e}" 

206 ) 

207 await asyncio.sleep(self.config.retry_delay) 

208 continue 

209 elapsed = (time.perf_counter() - started) * 1000 

210 return TaskResult( 

211 task_id=task.task_id, 

212 status=TaskStatus.FAILED, 

213 error=str(e), 

214 duration_ms=elapsed, 

215 retries=retries, 

216 started_at=started, 

217 finished_at=time.perf_counter(), 

218 ) 

219 

220 def cancel_all(self) -> None: 

221 """Cancel all pending tasks.""" 

222 if self._cancel_event: 

223 self._cancel_event.set()