Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/services/cancellation_cleanup.py: 48%
157 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1"""
2The CancellationCleanup service. Responsible for cancelling tasks and subflows that haven't finished.
3"""
5import datetime
6import logging
7from typing import Annotated
8from uuid import NAMESPACE_URL, UUID, uuid5
10import sqlalchemy as sa
11from docket import CurrentDocket, Depends, Docket, Logged, Perpetual, Retry
13import prefect.server.models as models
14from prefect._internal.uuid7 import uuid7
15from prefect.client.schemas.worker_channel import CANCELLING_TIMEOUT_TEARDOWN
16from prefect.logging import get_logger
17from prefect.server.database import PrefectDBInterface, provide_database_interface
18from prefect.server.exceptions import ObjectNotFoundError
19from prefect.server.schemas import filters, responses, states
20from prefect.server.services.perpetual_services import perpetual_service
21from prefect.server.worker_communication.cleanup_queue import (
22 CleanupQueueMessage,
23 WorkerCleanupQueue,
24 get_worker_cleanup_queue,
25)
26from prefect.settings.context import get_current_settings
27from prefect.types._datetime import now
29NON_TERMINAL_STATES = list(set(states.StateType) - states.TERMINAL_STATES)
30CANCELLING_TIMEOUT_CANCELLED_MESSAGE = (
31 "Flow run cancellation timed out; marked this flow run as Cancelled."
32)
33CANCELLING_TIMEOUT_CHECK_KEY_PREFIX = "cancelling-timeout"
34PUSH_WORK_POOL_TYPE_SUFFIX = ":push"
35MANAGED_WORK_POOL_TYPE_SUFFIX = ":managed"
36logger: logging.Logger = get_logger(__name__)
38_service_cleanup_queue: WorkerCleanupQueue | None = None
39_service_cleanup_queue_storage: str | None = None
42def _get_service_worker_cleanup_queue() -> WorkerCleanupQueue:
43 global _service_cleanup_queue, _service_cleanup_queue_storage
45 storage = get_current_settings().server.worker_channel.cleanup_queue_storage
46 if _service_cleanup_queue is None or _service_cleanup_queue_storage != storage:
47 _service_cleanup_queue = get_worker_cleanup_queue()
48 _service_cleanup_queue_storage = storage
50 return _service_cleanup_queue
53def cancelling_timeout_check_key(flow_run_id: UUID) -> str:
54 return f"{CANCELLING_TIMEOUT_CHECK_KEY_PREFIX}:{flow_run_id}"
57async def schedule_cancelling_timeout_check(
58 *,
59 docket: Docket,
60 flow_run_id: UUID,
61 flow_run_state_id: UUID,
62 when: datetime.datetime,
63) -> None:
64 await docket.replace(
65 handle_cancelling_timeout,
66 key=cancelling_timeout_check_key(flow_run_id),
67 when=when,
68 )(
69 flow_run_id=flow_run_id,
70 flow_run_state_id=flow_run_state_id,
71 timeout_cancelled_state_id=uuid7(),
72 )
75async def maybe_schedule_cancelling_timeout_check_for_state(
76 *,
77 docket: Docket,
78 flow_run_id: UUID,
79 state: states.State | None,
80) -> None:
81 settings = get_current_settings().server.services.cancellation_cleanup
82 if ( 82 ↛ 91line 82 didn't jump to line 91 because the condition on line 82 was always true
83 not settings.enabled
84 or state is None
85 or state.type != states.StateType.CANCELLING
86 or state.id is None
87 or state.timestamp is None
88 ):
89 return
91 await schedule_cancelling_timeout_check(
92 docket=docket,
93 flow_run_id=flow_run_id,
94 flow_run_state_id=state.id,
95 when=state.timestamp
96 + datetime.timedelta(seconds=settings.cancelling_timeout_seconds),
97 )
100# Docket task function for cancelling child task runs of a cancelled flow run
101async def cancel_child_task_runs(
102 flow_run_id: Annotated[UUID, Logged],
103 *,
104 db: PrefectDBInterface = Depends(provide_database_interface),
105) -> None:
106 """Cancel child task runs of a cancelled flow run (docket task)."""
107 async with db.session_context() as session:
108 child_task_runs = await models.task_runs.read_task_runs(
109 session,
110 flow_run_filter=filters.FlowRunFilter(
111 id=filters.FlowRunFilterId(any_=[flow_run_id])
112 ),
113 task_run_filter=filters.TaskRunFilter(
114 state=filters.TaskRunFilterState(
115 type=filters.TaskRunFilterStateType(any_=NON_TERMINAL_STATES)
116 )
117 ),
118 limit=100,
119 )
121 for task_run in child_task_runs:
122 async with db.session_context(begin_transaction=True) as session:
123 await models.task_runs.set_task_run_state(
124 session=session,
125 task_run_id=task_run.id,
126 state=states.Cancelled(message="The parent flow run was cancelled."),
127 force=True,
128 )
131# Docket task function for cancelling a subflow run whose parent was cancelled
132async def cancel_subflow_run(
133 subflow_run_id: Annotated[UUID, Logged],
134 *,
135 docket: Docket = CurrentDocket(),
136 db: PrefectDBInterface = Depends(provide_database_interface),
137) -> None:
138 """Cancel a subflow run whose parent flow run was cancelled (docket task)."""
139 async with db.session_context() as session:
140 flow_run = await models.flow_runs.read_flow_run(
141 session, flow_run_id=subflow_run_id
142 )
144 if not flow_run or not flow_run.parent_task_run_id or not flow_run.state:
145 return
147 if flow_run.state.type in states.TERMINAL_STATES:
148 return
150 parent_task_run = await models.task_runs.read_task_run(
151 session, task_run_id=flow_run.parent_task_run_id
152 )
154 if not parent_task_run or not parent_task_run.flow_run_id: 154 ↛ 157line 154 didn't jump to line 157 because the condition on line 154 was always true
155 return
157 containing_flow_run = await models.flow_runs.read_flow_run(
158 session, flow_run_id=parent_task_run.flow_run_id
159 )
161 if ( 161 ↛ 168line 161 didn't jump to line 168 because the condition on line 161 was always true
162 containing_flow_run
163 and containing_flow_run.state
164 and containing_flow_run.state.type != states.StateType.CANCELLED
165 ):
166 return
168 if ( 168 ↛ 174line 168 didn't jump to line 174 because the condition on line 168 was always true
169 flow_run.deployment_id
170 and flow_run.state.type == states.StateType.CANCELLING
171 ):
172 return
174 if flow_run.deployment_id:
175 state = states.Cancelling(message="The parent flow run was cancelled.")
176 else:
177 state = states.Cancelled(message="The parent flow run was cancelled.")
179 async with db.session_context(begin_transaction=True) as session:
180 state_result = await models.flow_runs.set_flow_run_state(
181 session=session,
182 flow_run_id=subflow_run_id,
183 state=state,
184 )
186 if state_result.status == responses.SetStateStatus.ACCEPT: 186 ↛ exitline 186 didn't return from function 'cancel_subflow_run' because the condition on line 186 was always true
187 await maybe_schedule_cancelling_timeout_check_for_state(
188 docket=docket,
189 flow_run_id=subflow_run_id,
190 state=state_result.state,
191 )
194async def handle_cancelling_timeout(
195 flow_run_id: Annotated[UUID, Logged],
196 flow_run_state_id: Annotated[UUID, Logged],
197 timeout_cancelled_state_id: Annotated[UUID, Logged],
198 *,
199 db: PrefectDBInterface = Depends(provide_database_interface),
200 cleanup_queue: WorkerCleanupQueue = Depends(_get_service_worker_cleanup_queue),
201 retry: Retry = Retry.forever(delay=datetime.timedelta(seconds=0.5)),
202) -> CleanupQueueMessage | None:
203 """Handle a scheduled CANCELLING timeout check for a single flow run."""
204 settings = get_current_settings().server.services.cancellation_cleanup
206 async with db.session_context(begin_transaction=True) as session:
207 flow_run_result = await session.execute(
208 sa.select(
209 db.FlowRun,
210 db.WorkQueue.work_pool_id,
211 db.WorkPool.type,
212 )
213 .outerjoin(db.WorkQueue, db.FlowRun.work_queue_id == db.WorkQueue.id)
214 .outerjoin(db.WorkPool, db.WorkQueue.work_pool_id == db.WorkPool.id)
215 .where(db.FlowRun.id == flow_run_id)
216 .with_for_update(of=db.FlowRun)
217 )
218 row = flow_run_result.first()
219 if row is None:
220 logger.info(
221 "Flow run %s no longer exists, skipping CANCELLING timeout",
222 flow_run_id,
223 )
224 return None
226 flow_run, work_pool_id, work_pool_type = row
227 timeout_cancelled_state_already_committed = (
228 flow_run.state_type == states.StateType.CANCELLED
229 and flow_run.state_id == timeout_cancelled_state_id
230 )
231 if not timeout_cancelled_state_already_committed:
232 if flow_run.state_id != flow_run_state_id:
233 logger.info(
234 "Flow run %s is in a new state, skipping CANCELLING timeout",
235 flow_run.id,
236 )
237 return None
239 if flow_run.state_type != states.StateType.CANCELLING:
240 logger.info(
241 "Flow run %s is no longer CANCELLING, skipping timeout",
242 flow_run.id,
243 )
244 return None
246 if flow_run.state_timestamp is None:
247 logger.info(
248 "Flow run %s has no CANCELLING state timestamp, skipping timeout",
249 flow_run.id,
250 )
251 return None
253 deadline = flow_run.state_timestamp + datetime.timedelta(
254 seconds=settings.cancelling_timeout_seconds
255 )
256 if deadline > now("UTC"):
257 logger.info(
258 "Flow run %s has not reached the CANCELLING timeout, rescheduling",
259 flow_run.id,
260 )
261 retry.at(deadline)
262 return None
264 try:
265 from prefect.server.orchestration.core_policy import CoreFlowPolicy
267 state_result = await models.flow_runs.set_flow_run_state(
268 session=session,
269 flow_run_id=flow_run.id,
270 state=states.Cancelled(
271 id=timeout_cancelled_state_id,
272 message=CANCELLING_TIMEOUT_CANCELLED_MESSAGE,
273 ),
274 flow_policy=CoreFlowPolicy,
275 )
276 except ObjectNotFoundError:
277 logger.info(
278 "Flow run %s was removed during CANCELLING timeout, skipping",
279 flow_run.id,
280 )
281 return None
283 if state_result.status != responses.SetStateStatus.ACCEPT:
284 logger.info(
285 "CANCELLING timeout state transition for flow run %s was not "
286 "accepted",
287 flow_run.id,
288 extra={
289 "status": state_result.status.value,
290 "reason": getattr(state_result.details, "reason", None),
291 },
292 )
293 return None
295 if (
296 state_result.state is None
297 or state_result.state.id != timeout_cancelled_state_id
298 or state_result.state.type != states.StateType.CANCELLED
299 ):
300 return None
302 if (
303 flow_run.state_type != states.StateType.CANCELLED
304 or flow_run.state_id != timeout_cancelled_state_id
305 ):
306 logger.info(
307 "Flow run %s is no longer in the timeout Cancelled state, skipping "
308 "CANCELLING timeout cleanup",
309 flow_run.id,
310 )
311 return None
313 if flow_run.work_queue_id is None or work_pool_id is None: 313 ↛ 322line 313 didn't jump to line 322 because the condition on line 313 was always true
314 logger.warning(
315 "Skipping CANCELLING timeout cleanup for unroutable flow run: "
316 "flow_run_id=%s work_queue_id=%s",
317 flow_run.id,
318 flow_run.work_queue_id,
319 )
320 return None
322 work_pool_type = str(work_pool_type) if work_pool_type else None
323 if work_pool_type and work_pool_type.endswith(
324 (PUSH_WORK_POOL_TYPE_SUFFIX, MANAGED_WORK_POOL_TYPE_SUFFIX)
325 ):
326 logger.info(
327 "Skipping CANCELLING timeout cleanup for workerless work pool flow "
328 "run: flow_run_id=%s work_queue_id=%s work_pool_id=%s "
329 "work_pool_type=%s",
330 flow_run.id,
331 flow_run.work_queue_id,
332 work_pool_id,
333 work_pool_type,
334 )
335 return None
337 cleanup_enqueue_parameters = {
338 "message_id": uuid5(
339 NAMESPACE_URL,
340 "prefect:"
341 f"{CANCELLING_TIMEOUT_TEARDOWN}:"
342 f"{flow_run.id}:{timeout_cancelled_state_id}",
343 ),
344 "idempotency_key": (
345 f"{CANCELLING_TIMEOUT_TEARDOWN}:"
346 f"{flow_run.id}:{timeout_cancelled_state_id}"
347 ),
348 "work_pool_id": work_pool_id,
349 "work_queue_id": flow_run.work_queue_id,
350 "kind": CANCELLING_TIMEOUT_TEARDOWN,
351 "target": {
352 "flow_run_id": str(flow_run.id),
353 "infrastructure_pid": flow_run.infrastructure_pid,
354 },
355 }
357 message = await cleanup_queue.enqueue(**cleanup_enqueue_parameters)
359 logger.info(
360 "Enqueued CANCELLING timeout cleanup message: flow_run_id=%s message_id=%s",
361 flow_run_id,
362 message.message_id,
363 )
364 return message
367@perpetual_service(
368 enabled_getter=lambda: (
369 get_current_settings().server.services.cancellation_cleanup.enabled
370 ),
371)
372async def ensure_cancelling_timeout_checks(
373 docket: Docket = CurrentDocket(),
374 db: PrefectDBInterface = Depends(provide_database_interface),
375 perpetual: Perpetual = Perpetual(
376 automatic=True,
377 every=datetime.timedelta(
378 seconds=get_current_settings().server.services.cancellation_cleanup.loop_seconds
379 ),
380 ),
381) -> None:
382 """Seed Docket timeout checks for flow runs already in CANCELLING."""
384 settings = get_current_settings().server.services.cancellation_cleanup
385 batch_size = 200
386 last_state_timestamp: datetime.datetime | None = None
387 last_flow_run_id: UUID | None = None
389 while True:
390 query_conditions = [
391 db.FlowRun.state_type == states.StateType.CANCELLING,
392 db.FlowRun.state_id.is_not(None),
393 db.FlowRun.state_timestamp.is_not(None),
394 ]
395 if last_state_timestamp is not None and last_flow_run_id is not None: 395 ↛ 396line 395 didn't jump to line 396 because the condition on line 395 was never true
396 query_conditions.append(
397 sa.or_(
398 db.FlowRun.state_timestamp > last_state_timestamp,
399 sa.and_(
400 db.FlowRun.state_timestamp == last_state_timestamp,
401 db.FlowRun.id > last_flow_run_id,
402 ),
403 )
404 )
406 cancelling_flow_query = (
407 sa.select(
408 db.FlowRun.id,
409 db.FlowRun.state_id,
410 db.FlowRun.state_timestamp,
411 )
412 .where(*query_conditions)
413 .order_by(db.FlowRun.state_timestamp, db.FlowRun.id)
414 .limit(batch_size)
415 )
417 async with db.session_context() as session:
418 result = await session.execute(cancelling_flow_query)
419 cancelling_flow_runs = result.all()
421 if not cancelling_flow_runs: 421 ↛ 424line 421 didn't jump to line 424 because the condition on line 421 was always true
422 break
424 for flow_run_id, flow_run_state_id, state_timestamp in cancelling_flow_runs:
425 await docket.add(
426 handle_cancelling_timeout,
427 key=cancelling_timeout_check_key(flow_run_id),
428 when=state_timestamp
429 + datetime.timedelta(seconds=settings.cancelling_timeout_seconds),
430 )(
431 flow_run_id=flow_run_id,
432 flow_run_state_id=flow_run_state_id,
433 timeout_cancelled_state_id=uuid7(),
434 )
436 last_flow_run_id, _, last_state_timestamp = cancelling_flow_runs[-1]
437 if len(cancelling_flow_runs) < batch_size:
438 break
441# Perpetual monitor for cancelled flow runs with child tasks (find and flood pattern)
442@perpetual_service(
443 enabled_getter=lambda: (
444 get_current_settings().server.services.cancellation_cleanup.enabled
445 ),
446)
447async def monitor_cancelled_flow_runs(
448 docket: Docket = CurrentDocket(),
449 db: PrefectDBInterface = Depends(provide_database_interface),
450 perpetual: Perpetual = Perpetual(
451 automatic=True,
452 every=datetime.timedelta(
453 seconds=get_current_settings().server.services.cancellation_cleanup.loop_seconds
454 ),
455 ),
456) -> None:
457 """Monitor for cancelled flow runs and schedule child task cancellation."""
459 batch_size = 200
460 cancelled_flow_query = (
461 sa.select(db.FlowRun.id)
462 .where(
463 db.FlowRun.state_type == states.StateType.CANCELLED,
464 db.FlowRun.end_time.is_not(None),
465 db.FlowRun.end_time >= (now("UTC") - datetime.timedelta(days=1)),
466 )
467 .order_by(db.FlowRun.id)
468 .limit(batch_size)
469 )
471 async with db.session_context() as session:
472 flow_run_result = await session.execute(cancelled_flow_query)
473 flow_run_ids = flow_run_result.scalars().all()
475 for flow_run_id in flow_run_ids: 475 ↛ 476line 475 didn't jump to line 476 because the loop on line 475 never started
476 await docket.add(cancel_child_task_runs)(flow_run_id)
479# Perpetual monitor for subflow runs that need cancellation (find and flood pattern)
480@perpetual_service(
481 enabled_getter=lambda: (
482 get_current_settings().server.services.cancellation_cleanup.enabled
483 ),
484)
485async def monitor_subflow_runs(
486 docket: Docket = CurrentDocket(),
487 db: PrefectDBInterface = Depends(provide_database_interface),
488 perpetual: Perpetual = Perpetual(
489 automatic=True,
490 every=datetime.timedelta(
491 seconds=get_current_settings().server.services.cancellation_cleanup.loop_seconds
492 ),
493 ),
494) -> None:
495 """Monitor for subflow runs that need to be cancelled."""
497 batch_size = 200
498 subflow_query = (
499 sa.select(db.FlowRun.id)
500 .where(
501 sa.or_(
502 db.FlowRun.state_type == states.StateType.PENDING,
503 db.FlowRun.state_type == states.StateType.SCHEDULED,
504 db.FlowRun.state_type == states.StateType.RUNNING,
505 db.FlowRun.state_type == states.StateType.PAUSED,
506 db.FlowRun.state_type == states.StateType.CANCELLING,
507 ),
508 db.FlowRun.parent_task_run_id.is_not(None),
509 )
510 .order_by(db.FlowRun.id)
511 .limit(batch_size)
512 )
514 async with db.session_context() as session:
515 subflow_run_result = await session.execute(subflow_query)
516 subflow_run_ids = subflow_run_result.scalars().all()
518 for subflow_run_id in subflow_run_ids: 518 ↛ 519line 518 didn't jump to line 519 because the loop on line 518 never started
519 await docket.add(cancel_subflow_run)(subflow_run_id)