Coverage for /home/runner/work/coxbuild/coxbuild/src/coxbuild/pipelines.py: 78%

1import inspect 

2import logging 

3import sys 

4import traceback 

5from dataclasses import asdict, dataclass, field 

6from datetime import datetime, timedelta 

7from enum import Enum 

8from graphlib import TopologicalSorter 

9from queue import Queue 

10from typing import TYPE_CHECKING, Any, Awaitable, Callable 

11 

12from coxbuild.configurations import Configuration 

13from coxbuild.hooks import Hook 

14from coxbuild.runtime import ExecutionState 

15 

16from .exceptions import CoxbuildRuntimeException, CoxbuildSchemaException 

17from .runners import Runner 

18from .tasks import Task, TaskContext, TaskHook, TaskResult 

19 

20if TYPE_CHECKING: 

21 from .extensions import Extension 

22 

23logger = logging.getLogger("pipelines") 

24 

25 

26class PipelineRuntimeException(CoxbuildRuntimeException): 

27 """Exception for pipeline running.""" 

28 pass 

29 

30 

31@dataclass 

32class PipelineContext: 

33 """Execution context for pipeline.""" 

34 pipeline: "Pipeline" 

35 """pipeline""" 

36 tasks: list[TaskResult] 

37 """tasks in the pipeline""" 

38 unmatchedNames: list[str] = field(default_factory=list) 

39 """not found task names""" 

40 config: Configuration | None = None 

41 """configuration""" 

42 

43 

44@dataclass 

45class PipelineResult: 

46 """Execution result for pipeline.""" 

47 duration: timedelta 

48 """execution duration""" 

49 tasks: list[TaskResult] 

50 """tasks in the pipeline""" 

51 exception: PipelineRuntimeException | None = None 

52 """exception when running""" 

53 

54 def __bool__(self): 

55 return self.exception is None and all((t for t in self.tasks if not t.task.continueOnError)) 

56 

57 @property 

58 def description(self) -> str: 

59 """Return result's description string.""" 

60 return "🟢 SUCCESS" if self else "🔴 FAILING" 

61 

62 

63@dataclass 

64class PipelineHook(Hook): 

65 """Hooks for a pipeline.""" 

66 extension: "Extension | None" = None 

67 

68 

69@dataclass 

70class PipelineBeforeTaskHook(PipelineHook): 

71 """Before task starting""" 

72 hook: Callable[[TaskContext], Awaitable[bool] | bool | None] | None = None 

73 

74 

75@dataclass 

76class PipelineBeforeHook(PipelineHook): 

77 """Before pipeline starting""" 

78 hook: Callable[[PipelineContext], 

79 Awaitable[bool] | bool | None] | None = None 

80 

81 

82@dataclass 

83class PipelineAfterHook(PipelineHook): 

84 """After pipeline finished""" 

85 hook: Callable[[PipelineContext, PipelineResult], 

86 Awaitable | None] | None = None 

87 

88 

89@dataclass 

90class PipelineAfterTaskHook(PipelineHook): 

91 """After task finished""" 

92 hook: Callable[[TaskContext, TaskResult], Awaitable | None] | None = None 

93 

94 

95def beforePipeline(hook: Callable[[PipelineContext], Awaitable[bool] | bool | None] | PipelineBeforeHook) -> PipelineBeforeHook: 

96 """ 

97 Decorator to configure before hook of a pipeline. 

98 Run before pipeline. 

99 

100 hook: hook function 

101 """ 

102 return hook if isinstance(hook, PipelineBeforeHook) else PipelineBeforeHook(hook) 

103 

104 

105def afterPipeline(hook: Callable[[PipelineContext, PipelineResult], Awaitable | None] | PipelineAfterHook) -> PipelineAfterHook: 

106 """ 

107 Decorator to configure after hook of a pipeline. 

108 Run after pipeline. 

109 

110 hook: hook function 

111 """ 

112 return hook if isinstance(hook, PipelineAfterHook) else PipelineAfterHook(hook) 

113 

114 

115def beforeTask(hook: Callable[[TaskContext], Awaitable[bool] | bool | None] | PipelineBeforeTaskHook) -> PipelineBeforeTaskHook: 

116 """ 

117 Decorator to configure before task hook of a task. 

118 Run before task. 

119 

120 hook: hook function 

121 """ 

122 return hook if isinstance(hook, PipelineBeforeTaskHook) else PipelineBeforeTaskHook(hook) 

123 

124 

125def afterTask(hook: Callable[[TaskContext, TaskResult], Awaitable | None] | PipelineAfterTaskHook) -> PipelineAfterTaskHook: 

126 """ 

127 Decorator to configure after task hook of a task. 

128 Run after task. 

129 

130 hook: hook function 

131 """ 

132 return hook if isinstance(hook, PipelineAfterTaskHook) else PipelineAfterTaskHook(hook) 

133 

134 

135class PipelineRunner(Runner): 

136 """Runner for pipeline.""" 

137 

138 def __init__(self, pipeline: "Pipeline", tasks: list[Task], hooks: list[PipelineHook], unmatchedNames: list[str]) -> None: 

139 """ 

140 Create runner. 

141 

142 pipeline: pipeline 

143 tasks: tasks in the pipeline 

144 hooks: hooks for pipeline 

145 unmatchedNames: not found task names 

146 """ 

147 self.pipeline = pipeline 

148 self.tasks = tasks 

149 self.beforeTask = [beforeTask for beforeTask in hooks if isinstance( 

150 beforeTask, PipelineBeforeTaskHook)] 

151 self.afterTask = [afterTask for afterTask in hooks if isinstance( 

152 afterTask, PipelineAfterTaskHook)] 

153 self.before = [before for before in hooks if isinstance( 

154 before, PipelineBeforeHook)] 

155 self.after = [after for after in hooks if isinstance( 

156 after, PipelineAfterHook)] 

157 self.context = PipelineContext( 

158 self.pipeline, self.tasks, unmatchedNames) 

159 self.result = None 

160 

161 super().__init__(self._run) 

162 

163 async def _before(self): 

164 logger.debug(f"Run pipeline before hook") 

165 try: 

166 for hook in self.before: 

167 pre = hook.hook(self.context) 

168 if inspect.isawaitable(pre): 

169 pre: bool = await pre 

170 

171 if pre == False: 

172 return False 

173 except Exception as ex: 

174 logger.error(f"Run pipeline before hook failed.", exc_info=ex) 

175 print(f"Run pipeline before hook failed: {ex}") 

176 return False 

177 

178 async def _beforeTask(self, context: TaskContext): 

179 logger.debug(f"Run pipeline before task hook for {context.task.name}") 

180 for hook in self.beforeTask: 

181 pre = hook.hook(context) 

182 if inspect.isawaitable(pre): 

183 pre: bool = await pre 

184 if pre == False: 

185 return False 

186 

187 async def _afterTask(self, context: TaskContext, result: TaskResult): 

188 logger.debug(f"Run pipeline after task hook for {context.task.name}") 

189 for hook in self.afterTask: 

190 res = hook.hook(context, result) 

191 if inspect.isawaitable(res): 

192 await res 

193 

194 async def _after(self): 

195 logger.debug(f"Pipeline after hook.") 

196 try: 

197 for hook in self.after: 

198 res = hook.hook(self.context, self.result) 

199 if inspect.isawaitable(res): 

200 await res 

201 except Exception as ex: 

202 logger.error(f"Run pipeline after hook failed.", exc_info=ex) 

203 print(f"Run pipeline after hook failed: {ex}") 

204 

205 async def _run(self): 

206 n = len(self.tasks) 

207 

208 for i, task in enumerate(self.tasks): 

209 logger.debug(f"Run task {i+1}({task.name}) of {n} tasks") 

210 print(f"{'-'*15} ({i+1}/{n}) 📜 Task {task.name} {'-'*15}") 

211 print("") 

212 

213 self._executionState.task = task 

214 

215 runner = task() 

216 runner.context.config = self.context.config 

217 

218 pre = await self._beforeTask(runner.context) 

219 if pre == False: 

220 message = f"Stop task {task.name} running by pipeline before hook" 

221 logger.info(message) 

222 print(message) 

223 continue 

224 

225 res = await runner 

226 

227 self._results.append(res) 

228 

229 await self._afterTask(runner.context, res) 

230 

231 self._executionState.task = None 

232 

233 if not res: 

234 if task.continueOnError: 

235 message = f"Task {task.name} failed, but continue on error." 

236 logger.error(message) 

237 print(message) 

238 print("") 

239 else: 

240 print("") 

241 break 

242 

243 async def __aenter__(self) -> Callable[[], Awaitable | None]: 

244 logger.debug(f"Running pipeline: {self.tasks}") 

245 print(f"{'-'*20} ⌛ Running 🕰️ {datetime.now()} {'-'*20}") 

246 

247 self.result = None 

248 self._results: list[TaskResult] = [] 

249 

250 self.context.config = self.context.config or Configuration() 

251 self._executionState = ExecutionState(self.context.config) 

252 self._executionState.matchedTasks = self.tasks 

253 self._executionState.unmatchedTasks = self.context.unmatchedNames 

254 

255 res = await super().__aenter__() 

256 

257 pre = await self._before() 

258 if pre == False: 

259 message = "Stop pipeline running by pipeline setup hook" 

260 logger.info(message) 

261 print(message) 

262 res = None 

263 

264 print("") 

265 

266 return res or (lambda: None) 

267 

268 async def __aexit__(self, exc_type, exc_value, exc_tb) -> bool: 

269 await super().__aexit__(exc_type, exc_value, exc_tb) 

270 

271 exception = None if self.exc_value is None else PipelineRuntimeException( 

272 f"Failed to run pipeline", cause=self.exc_value) 

273 self.result = PipelineResult( 

274 duration=self.duration, tasks=self._results, exception=exception) 

275 

276 if self.exc_value is not None: 

277 logger.error("Task execute exception.", exc_info=self.exc_value) 

278 traceback.print_exception( 

279 self.exc_type, self.exc_value, self.exc_tb, file=sys.stdout) 

280 

281 await self._after() 

282 

283 logger.info(f"Finish pipeline: {self.result}") 

284 

285 print( 

286 f"{'-'*20} 📋 Done {self.result.description} (⏱️ {self.result.duration}) {'-'*20}") 

287 

288 cnt = len(self.result.tasks) 

289 

290 for i, tr in enumerate(self.result.tasks): 

291 print( 

292 f"({i+1}/{cnt})\t{tr.description}\t⏱️ {tr.duration}\t📜 {tr.task.name}") 

293 

294 del self._results 

295 del self._executionState 

296 

297 return True 

298 

299 def __await__(self): 

300 yield from super().__await__() 

301 return self.result 

302 

303 

304@dataclass 

305class Pipeline: 

306 """Pipeline to run tasks.""" 

307 

308 tasks: dict[str, Task] = field(default_factory=dict) 

309 """tasks in the pipeline""" 

310 hooks: list[PipelineHook] = field(default_factory=list) 

311 """hook for pipeline""" 

312 

313 def copy(self) -> "Pipeline": 

314 """Copy pipeline.""" 

315 return Pipeline(tasks=self.tasks.copy(), hooks=self.hooks.copy()) 

316 

317 def register(self, task: Task) -> None: 

318 """ 

319 Register a task into pipeline. 

320 

321 task: task to register 

322 """ 

323 if task.name in self.tasks: 

324 logger.warning( 

325 f"Replace registered task: {self.tasks[task.name]}.") 

326 self.tasks[task.name] = task 

327 logger.debug(f"Register task {task}") 

328 

329 def hook(self, hook: PipelineHook) -> None: 

330 """ 

331 Hook pipeline. 

332 

333 thook: the hook 

334 """ 

335 

336 self.hooks.append(hook) 

337 

338 def beforeTask(self, body: Callable[[TaskContext], Awaitable[bool] | bool | None] | PipelineBeforeTaskHook) -> PipelineBeforeTaskHook: 

339 """Add before task hook.""" 

340 hook = beforeTask(body) 

341 self.hook(hook) 

342 return hook 

343 

344 def afterTask(self, body: Callable[[TaskContext, TaskResult], Awaitable[bool] | bool | None] | PipelineAfterTaskHook) -> PipelineAfterTaskHook: 

345 """Add after task hook.""" 

346 hook = afterTask(body) 

347 self.hook(hook) 

348 return hook 

349 

350 def before(self, body: Callable[[PipelineContext], Awaitable[bool] | bool | None] | PipelineBeforeHook) -> PipelineBeforeHook: 

351 """Add before hook.""" 

352 hook = beforePipeline(body) 

353 self.hook(hook) 

354 return hook 

355 

356 def after(self, body: Callable[[PipelineContext, PipelineResult], Awaitable | None] | PipelineAfterHook) -> PipelineAfterHook: 

357 """Add after hook.""" 

358 hook = afterPipeline(body) 

359 self.hook(hook) 

360 return hook 

361 

362 def __call__(self, *args: str | Task, **kwds: Any): 

363 """ 

364 Build runner of pipeline. 

365 

366 args: list of tasks or task names 

367 """ 

368 

369 logger.debug(f"Build pipeline by {args}.") 

370 

371 unmatchedNames = [] 

372 

373 tks: set[Task] = set() 

374 q: Queue[Task] = Queue() 

375 for task in args: 

376 if isinstance(task, str): 

377 if task in self.tasks: 

378 task = self.tasks[task] 

379 else: 

380 unmatchedNames.append(task) 

381 continue 

382 if task not in tks: 

383 tks.add(task) 

384 q.put(task) 

385 

386 if len(unmatchedNames) > 0: 

387 message = f"Not found task names: {unmatchedNames}" 

388 logger.warning(message) 

389 print(message) 

390 

391 while not q.empty(): 

392 t = q.get() 

393 for d in t.deps: 

394 if d not in tks: 

395 tks.add(d) 

396 q.put(d) 

397 

398 logger.debug(f"Build pipeline for tasks: {tks}.") 

399 

400 graph: dict[Task, set[Task]] = {} 

401 

402 for key in tks: 

403 graph[key] = {d for d in key.deps} 

404 

405 tasks = list((task 

406 for task in TopologicalSorter(graph).static_order())) 

407 

408 logger.debug(f"Tasks to run: {', '.join((t.name for t in tasks))}") 

409 

410 return PipelineRunner(self, tasks, self.hooks, unmatchedNames)