Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/orchestration/global_policy.py: 88%

1""" 

2Bookkeeping logic that fires on every state transition. 

3 

4For clarity, `GlobalFlowpolicy` and `GlobalTaskPolicy` contain all transition logic 

5implemented using `BaseUniversalTransform`. None of these operations modify state, and regardless of what orchestration Prefect REST API might 

6enforce on a transition, the global policies contain Prefect's necessary bookkeeping. 

7Because these transforms record information about the validated state committed to the 

8state database, they should be the most deeply nested contexts in orchestration loop. 

9""" 

10 

11from __future__ import annotations 

12 

13from typing import Any, Union, cast 

14 

15from packaging.version import Version 

16 

17import prefect.server.models as models 

18from prefect.server.database import orm_models 

19from prefect.server.orchestration.policies import BaseOrchestrationPolicy 

20from prefect.server.orchestration.rules import ( 

21 BaseOrchestrationRule, 

22 BaseUniversalTransform, 

23 FlowOrchestrationContext, 

24 FlowRunUniversalTransform, 

25 GenericOrchestrationContext, 

26 OrchestrationContext, 

27 TaskOrchestrationContext, 

28 TaskRunUniversalTransform, 

29) 

30from prefect.server.schemas import core 

31from prefect.server.schemas.core import FlowRunPolicy 

32 

33 

34def COMMON_GLOBAL_TRANSFORMS() -> list[ 

35 type[ 

36 BaseUniversalTransform[ 

37 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

38 ] 

39 ] 

40]: 

41 return [ 

42 SetRunStateType, 

43 SetRunStateName, 

44 SetRunStateTimestamp, 

45 SetStartTime, 

46 SetEndTime, 

47 IncrementRunTime, 

48 SetExpectedStartTime, 

49 SetNextScheduledStartTime, 

50 UpdateStateDetails, 

51 ] 

52 

53 

54class GlobalFlowPolicy(BaseOrchestrationPolicy[orm_models.FlowRun, core.FlowRunPolicy]): 

55 """ 

56 Global transforms that run against flow-run-state transitions in priority order. 

57 

58 These transforms are intended to run immediately before and after a state transition 

59 is validated. 

60 """ 

61 

62 @staticmethod 

63 def priority() -> list[ 

64 Union[ 

65 type[BaseUniversalTransform[orm_models.FlowRun, core.FlowRunPolicy]], 

66 type[BaseOrchestrationRule[orm_models.FlowRun, core.FlowRunPolicy]], 

67 ] 

68 ]: 

69 return cast( 

70 list[ 

71 Union[ 

72 type[ 

73 BaseUniversalTransform[orm_models.FlowRun, core.FlowRunPolicy] 

74 ], 

75 type[BaseOrchestrationRule[orm_models.FlowRun, core.FlowRunPolicy]], 

76 ] 

77 ], 

78 COMMON_GLOBAL_TRANSFORMS(), 

79 ) + [ 

80 UpdateSubflowParentTask, 

81 UpdateSubflowStateDetails, 

82 IncrementFlowRunCount, 

83 RemoveResumingIndicator, 

84 ] 

85 

86 

87class GlobalTaskPolicy(BaseOrchestrationPolicy[orm_models.TaskRun, core.TaskRunPolicy]): 

88 """ 

89 Global transforms that run against task-run-state transitions in priority order. 

90 

91 These transforms are intended to run immediately before and after a state transition 

92 is validated. 

93 """ 

94 

95 @staticmethod 

96 def priority() -> list[ 

97 Union[ 

98 type[BaseUniversalTransform[orm_models.TaskRun, core.TaskRunPolicy]], 

99 type[BaseOrchestrationRule[orm_models.TaskRun, core.TaskRunPolicy]], 

100 ] 

101 ]: 

102 return cast( 

103 list[ 

104 Union[ 

105 type[ 

106 BaseUniversalTransform[orm_models.TaskRun, core.TaskRunPolicy] 

107 ], 

108 type[BaseOrchestrationRule[orm_models.TaskRun, core.TaskRunPolicy]], 

109 ] 

110 ], 

111 COMMON_GLOBAL_TRANSFORMS(), 

112 ) + [IncrementTaskRunCount] 

113 

114 

115class SetRunStateType( 

116 BaseUniversalTransform[ 

117 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

118 ] 

119): 

120 """ 

121 Updates the state type of a run on a state transition. 

122 """ 

123 

124 async def before_transition( 

125 self, context: GenericOrchestrationContext[orm_models.Run, Any] 

126 ) -> None: 

127 if self.nullified_transition(): 

128 return 

129 

130 # record the new state's type 

131 if context.proposed_state is not None: 131 ↛ exitline 131 didn't return from function 'before_transition' because the condition on line 131 was always true

132 context.run.state_type = context.proposed_state.type 

133 

134 

135class SetRunStateName( 

136 BaseUniversalTransform[ 

137 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

138 ] 

139): 

140 """ 

141 Updates the state name of a run on a state transition. 

142 """ 

143 

144 async def before_transition( 

145 self, context: GenericOrchestrationContext[orm_models.Run, Any] 

146 ) -> None: 

147 if self.nullified_transition(): 

148 return 

149 

150 if context.proposed_state is not None: 150 ↛ exitline 150 didn't return from function 'before_transition' because the condition on line 150 was always true

151 # record the new state's name 

152 context.run.state_name = context.proposed_state.name 

153 

154 

155class SetStartTime( 

156 BaseUniversalTransform[ 

157 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

158 ] 

159): 

160 """ 

161 Records the time a run enters a running state for the first time. 

162 """ 

163 

164 async def before_transition( 

165 self, context: GenericOrchestrationContext[orm_models.Run, Any] 

166 ) -> None: 

167 if self.nullified_transition(): 

168 return 

169 

170 if context.proposed_state is not None: 170 ↛ exitline 170 didn't return from function 'before_transition' because the condition on line 170 was always true

171 # if entering a running state and no start time is set... 

172 if context.proposed_state.is_running() and context.run.start_time is None: 

173 # set the start time 

174 context.run.start_time = context.proposed_state.timestamp 

175 

176 

177class SetRunStateTimestamp( 

178 BaseUniversalTransform[ 

179 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

180 ] 

181): 

182 """ 

183 Records the time a run changes states. 

184 """ 

185 

186 async def before_transition( 

187 self, context: GenericOrchestrationContext[orm_models.Run, Any] 

188 ) -> None: 

189 if self.nullified_transition(): 

190 return 

191 

192 if context.proposed_state is not None: 192 ↛ exitline 192 didn't return from function 'before_transition' because the condition on line 192 was always true

193 # record the new state's timestamp 

194 context.run.state_timestamp = context.proposed_state.timestamp 

195 

196 

197class SetEndTime( 

198 BaseUniversalTransform[ 

199 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

200 ] 

201): 

202 """ 

203 Records the time a run enters a terminal state. 

204 

205 With normal client usage, a run will not transition out of a terminal state. 

206 However, it's possible to force these transitions manually via the API. While 

207 leaving a terminal state, the end time will be unset. 

208 """ 

209 

210 async def before_transition( 

211 self, context: GenericOrchestrationContext[orm_models.Run, Any] 

212 ) -> None: 

213 if self.nullified_transition(): 

214 return 

215 

216 if context.proposed_state is not None: 216 ↛ exitline 216 didn't return from function 'before_transition' because the condition on line 216 was always true

217 # if exiting a final state for a non-final state... 

218 if ( 

219 context.initial_state 

220 and context.initial_state.is_final() 

221 and not context.proposed_state.is_final() 

222 ): 

223 # clear the end time 

224 context.run.end_time = None 

225 

226 # if entering a final state... 

227 if context.proposed_state.is_final(): 

228 if not context.run.end_time: 

229 # Set start_time if not set (for flows that fail before execution) 

230 if not context.run.start_time: 

231 context.run.start_time = context.proposed_state.timestamp 

232 context.run.end_time = context.proposed_state.timestamp 

233 

234 

235class IncrementRunTime( 

236 BaseUniversalTransform[ 

237 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

238 ] 

239): 

240 """ 

241 Records the amount of time a run spends in the running state. 

242 """ 

243 

244 async def before_transition( 

245 self, context: GenericOrchestrationContext[orm_models.Run, Any] 

246 ) -> None: 

247 if self.nullified_transition(): 

248 return 

249 

250 if context.proposed_state is not None: 250 ↛ exitline 250 didn't return from function 'before_transition' because the condition on line 250 was always true

251 # if exiting a running state... 

252 if context.initial_state and context.initial_state.is_running(): 

253 # increment the run time by the time spent in the previous state 

254 context.run.total_run_time += ( 

255 context.proposed_state.timestamp - context.initial_state.timestamp 

256 ) 

257 

258 

259class IncrementFlowRunCount( 

260 FlowRunUniversalTransform[orm_models.FlowRun, core.FlowRunPolicy] 

261): 

262 """ 

263 Records the number of times a run enters a running state. For use with retries. 

264 """ 

265 

266 async def before_transition( 

267 self, context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy] 

268 ) -> None: 

269 if self.nullified_transition(): 

270 return 

271 

272 if context.proposed_state is not None: 272 ↛ exitline 272 didn't return from function 'before_transition' because the condition on line 272 was always true

273 # if entering a running state... 

274 if context.proposed_state.is_running(): 

275 # do not increment the run count if resuming a paused flow 

276 api_version = context.parameters.get("api-version", None) 

277 if api_version is None or api_version >= Version("0.8.4"): 277 ↛ 282line 277 didn't jump to line 282 because the condition on line 277 was always true

278 if context.run.empirical_policy.resuming: 

279 return 

280 

281 # increment the run count 

282 context.run.run_count += 1 

283 

284 

285class RemoveResumingIndicator( 

286 BaseUniversalTransform[orm_models.FlowRun, core.FlowRunPolicy] 

287): 

288 """ 

289 Removes the indicator on a flow run that marks it as resuming. 

290 """ 

291 

292 async def before_transition( 

293 self, context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy] 

294 ) -> None: 

295 if self.nullified_transition(): 

296 return 

297 

298 proposed_state = context.proposed_state 

299 

300 api_version = context.parameters.get("api-version", None) 

301 if api_version is None or api_version >= Version("0.8.4"): 301 ↛ exitline 301 didn't return from function 'before_transition' because the condition on line 301 was always true

302 if proposed_state is not None and ( 

303 proposed_state.is_running() or proposed_state.is_final() 

304 ): 

305 if context.run.empirical_policy.resuming: 

306 updated_policy = context.run.empirical_policy.model_dump() 

307 updated_policy["resuming"] = False 

308 context.run.empirical_policy = FlowRunPolicy(**updated_policy) 

309 

310 

311class IncrementTaskRunCount(TaskRunUniversalTransform): 

312 """ 

313 Records the number of times a run enters a running state. For use with retries. 

314 """ 

315 

316 async def before_transition( 

317 self, context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy] 

318 ) -> None: 

319 if self.nullified_transition(): 319 ↛ 320line 319 didn't jump to line 320 because the condition on line 319 was never true

320 return 

321 

322 proposed_state = context.proposed_state 

323 

324 # if entering a running state... 

325 if proposed_state is not None and proposed_state.is_running(): 325 ↛ 327line 325 didn't jump to line 327 because the condition on line 325 was never true

326 # increment the run count 

327 context.run.run_count += 1 

328 

329 

330class SetExpectedStartTime( 

331 BaseUniversalTransform[ 

332 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

333 ] 

334): 

335 """ 

336 Estimates the time a state is expected to start running if not set. 

337 

338 For scheduled states, this estimate is simply the scheduled time. For other states, 

339 this is set to the time the proposed state was created by Prefect. 

340 """ 

341 

342 async def before_transition( 

343 self, context: GenericOrchestrationContext[orm_models.Run, Any] 

344 ) -> None: 

345 if self.nullified_transition(): 

346 return 

347 

348 # set expected start time if this is the first state 

349 if not context.run.expected_start_time and context.proposed_state is not None: 

350 if context.proposed_state.is_scheduled(): 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true

351 context.run.expected_start_time = ( 

352 context.proposed_state.state_details.scheduled_time 

353 ) 

354 else: 

355 context.run.expected_start_time = context.proposed_state.timestamp 

356 

357 

358class SetNextScheduledStartTime( 

359 BaseUniversalTransform[ 

360 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

361 ] 

362): 

363 """ 

364 Records the scheduled time on a run. 

365 

366 When a run enters a scheduled state, `run.next_scheduled_start_time` is set to 

367 the state's scheduled time. When leaving a scheduled state, 

368 `run.next_scheduled_start_time` is unset. 

369 """ 

370 

371 async def before_transition( 

372 self, context: GenericOrchestrationContext[orm_models.Run, Any] 

373 ) -> None: 

374 if self.nullified_transition(): 

375 return 

376 

377 # remove the next scheduled start time if exiting a scheduled state 

378 if context.initial_state and context.initial_state.is_scheduled(): 378 ↛ 379line 378 didn't jump to line 379 because the condition on line 378 was never true

379 context.run.next_scheduled_start_time = None 

380 

381 # set next scheduled start time if entering a scheduled state 

382 if context.proposed_state is not None and context.proposed_state.is_scheduled(): 382 ↛ 383line 382 didn't jump to line 383 because the condition on line 382 was never true

383 context.run.next_scheduled_start_time = ( 

384 context.proposed_state.state_details.scheduled_time 

385 ) 

386 

387 

388class UpdateSubflowParentTask( 

389 BaseUniversalTransform[orm_models.FlowRun, core.FlowRunPolicy] 

390): 

391 """ 

392 Whenever a subflow changes state, it must update its parent task run's state. 

393 """ 

394 

395 async def after_transition( 

396 self, context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy] 

397 ) -> None: 

398 if self.nullified_transition(): 

399 return 

400 

401 # only applies to flow runs with a parent task run id 

402 if ( 402 ↛ 407line 402 didn't jump to line 407 because the condition on line 402 was never true

403 context.run.parent_task_run_id is not None 

404 and context.validated_state is not None 

405 ): 

406 # avoid mutation of the flow run state 

407 subflow_parent_task_state = context.validated_state.fresh_copy() 

408 

409 # set the task's "child flow run id" to be the subflow run id 

410 subflow_parent_task_state.state_details.child_flow_run_id = context.run.id 

411 

412 await models.task_runs.set_task_run_state( 

413 session=context.session, 

414 task_run_id=context.run.parent_task_run_id, 

415 state=subflow_parent_task_state, 

416 force=True, 

417 ) 

418 

419 

420class UpdateSubflowStateDetails( 

421 BaseUniversalTransform[orm_models.FlowRun, core.FlowRunPolicy] 

422): 

423 """ 

424 Update a child subflow state's references to a corresponding tracking task run id 

425 in the parent flow run 

426 """ 

427 

428 async def before_transition( 

429 self, context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy] 

430 ) -> None: 

431 if self.nullified_transition(): 

432 return 

433 

434 # only applies to flow runs with a parent task run id 

435 if ( 435 ↛ 439line 435 didn't jump to line 439 because the condition on line 435 was never true

436 context.run.parent_task_run_id is not None 

437 and context.proposed_state is not None 

438 ): 

439 context.proposed_state.state_details.task_run_id = ( 

440 context.run.parent_task_run_id 

441 ) 

442 

443 

444class UpdateStateDetails( 

445 BaseUniversalTransform[ 

446 orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy] 

447 ] 

448): 

449 """ 

450 Update a state's references to a corresponding flow- or task- run. 

451 """ 

452 

453 async def before_transition( 

454 self, 

455 context: GenericOrchestrationContext, 

456 ) -> None: 

457 if self.nullified_transition(): 

458 return 

459 

460 if isinstance(context, FlowOrchestrationContext): 

461 flow_run = await context.flow_run() 

462 context.proposed_state.state_details.flow_run_id = flow_run.id 

463 

464 elif isinstance(context, TaskOrchestrationContext): 464 ↛ exitline 464 didn't return from function 'before_transition' because the condition on line 464 was always true

465 task_run = await context.task_run() 

466 context.proposed_state.state_details.flow_run_id = task_run.flow_run_id 

467 context.proposed_state.state_details.task_run_id = task_run.id