Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/api/concurrency_limits_v2.py: 74%
162 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1from datetime import datetime, timedelta, timezone
2from typing import List, Literal, Optional, Union
3from uuid import UUID
5import sqlalchemy as sa
6from fastapi import Body, Depends, HTTPException, Path
7from sqlalchemy.ext.asyncio import AsyncSession
9import prefect.server.models as models
10import prefect.server.schemas as schemas
11from prefect._internal.compatibility.starlette import status
12from prefect.server.api.dependencies import LimitBody
13from prefect.server.concurrency.lease_storage import (
14 ConcurrencyLimitLeaseMetadata,
15 get_concurrency_lease_storage,
16)
17from prefect.server.database import PrefectDBInterface, provide_database_interface
18from prefect.server.schemas import actions
19from prefect.server.utilities.schemas import PrefectBaseModel
20from prefect.server.utilities.server import PrefectRouter
21from prefect.settings.context import get_current_settings
22from prefect.types._concurrency import ConcurrencyLeaseHolder
23from prefect.utilities.math import clamped_poisson_interval
25router: PrefectRouter = PrefectRouter(
26 prefix="/v2/concurrency_limits", tags=["Concurrency Limits V2"]
27)
30def _global_concurrency_limit_response(
31 model: object, active_slots: float | int
32) -> schemas.responses.GlobalConcurrencyLimitResponse:
33 return schemas.responses.GlobalConcurrencyLimitResponse.model_validate(
34 model
35 ).model_copy(update={"active_slots": int(active_slots)})
38@router.post("/", status_code=status.HTTP_201_CREATED)
39async def create_concurrency_limit_v2(
40 concurrency_limit: actions.ConcurrencyLimitV2Create,
41 db: PrefectDBInterface = Depends(provide_database_interface),
42) -> schemas.core.ConcurrencyLimitV2:
43 """
44 Create a task run concurrency limit.
46 For more information, see https://docs.prefect.io/v3/how-to-guides/workflows/global-concurrency-limits.
47 """
48 async with db.session_context(begin_transaction=True) as session:
49 model = await models.concurrency_limits_v2.create_concurrency_limit(
50 session=session, concurrency_limit=concurrency_limit
51 )
53 return schemas.core.ConcurrencyLimitV2.model_validate(model)
56@router.get("/{id_or_name}")
57async def read_concurrency_limit_v2(
58 id_or_name: Union[UUID, str] = Path(
59 ..., description="The ID or name of the concurrency limit", alias="id_or_name"
60 ),
61 db: PrefectDBInterface = Depends(provide_database_interface),
62) -> schemas.responses.GlobalConcurrencyLimitResponse:
63 if isinstance(id_or_name, str): # TODO: this seems like it shouldn't be necessary 63 ↛ 68line 63 didn't jump to line 68 because the condition on line 63 was always true
64 try:
65 id_or_name = UUID(id_or_name)
66 except ValueError:
67 pass
68 async with db.session_context() as session:
69 if isinstance(id_or_name, UUID):
70 where = db.ConcurrencyLimitV2.id == id_or_name
71 else:
72 where = db.ConcurrencyLimitV2.name == id_or_name
73 result = await session.execute(
74 sa.select(
75 db.ConcurrencyLimitV2,
76 models.concurrency_limits_v2.active_slots_after_decay(db).label(
77 "active_slots"
78 ),
79 ).where(where)
80 )
81 row = result.first()
83 if not row:
84 raise HTTPException(
85 status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency Limit not found"
86 )
88 return _global_concurrency_limit_response(row[0], row.active_slots)
91@router.post("/filter")
92async def read_all_concurrency_limits_v2(
93 limit: int = LimitBody(),
94 offset: int = Body(0, ge=0),
95 db: PrefectDBInterface = Depends(provide_database_interface),
96) -> List[schemas.responses.GlobalConcurrencyLimitResponse]:
97 async with db.session_context() as session:
98 query = sa.select(
99 db.ConcurrencyLimitV2,
100 models.concurrency_limits_v2.active_slots_after_decay(db).label(
101 "active_slots"
102 ),
103 ).order_by(db.ConcurrencyLimitV2.name)
105 if offset is not None: 105 ↛ 107line 105 didn't jump to line 107 because the condition on line 105 was always true
106 query = query.offset(offset)
107 if limit is not None: 107 ↛ 110line 107 didn't jump to line 110 because the condition on line 107 was always true
108 query = query.limit(limit)
110 result = await session.execute(query)
111 rows = result.all()
113 return [
114 _global_concurrency_limit_response(row[0], row.active_slots) for row in rows
115 ]
118@router.patch("/{id_or_name}", status_code=status.HTTP_204_NO_CONTENT)
119async def update_concurrency_limit_v2(
120 concurrency_limit: actions.ConcurrencyLimitV2Update,
121 id_or_name: Union[UUID, str] = Path(
122 ..., description="The ID or name of the concurrency limit", alias="id_or_name"
123 ),
124 db: PrefectDBInterface = Depends(provide_database_interface),
125) -> None:
126 if isinstance(id_or_name, str): # TODO: this seems like it shouldn't be necessary 126 ↛ 131line 126 didn't jump to line 131 because the condition on line 126 was always true
127 try:
128 id_or_name = UUID(id_or_name)
129 except ValueError:
130 pass
131 async with db.session_context(begin_transaction=True) as session:
132 if isinstance(id_or_name, UUID):
133 updated = await models.concurrency_limits_v2.update_concurrency_limit(
134 session,
135 concurrency_limit_id=id_or_name,
136 concurrency_limit=concurrency_limit,
137 )
138 else:
139 updated = await models.concurrency_limits_v2.update_concurrency_limit(
140 session, name=id_or_name, concurrency_limit=concurrency_limit
141 )
143 if not updated:
144 raise HTTPException(
145 status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency Limit not found"
146 )
149@router.delete("/{id_or_name}", status_code=status.HTTP_204_NO_CONTENT)
150async def delete_concurrency_limit_v2(
151 id_or_name: Union[UUID, str] = Path(
152 ..., description="The ID or name of the concurrency limit", alias="id_or_name"
153 ),
154 db: PrefectDBInterface = Depends(provide_database_interface),
155) -> None:
156 if isinstance(id_or_name, str): # TODO: this seems like it shouldn't be necessary 156 ↛ 161line 156 didn't jump to line 161 because the condition on line 156 was always true
157 try:
158 id_or_name = UUID(id_or_name)
159 except ValueError:
160 pass
161 async with db.session_context(begin_transaction=True) as session:
162 deleted = False
163 if isinstance(id_or_name, UUID):
164 deleted = await models.concurrency_limits_v2.delete_concurrency_limit(
165 session, concurrency_limit_id=id_or_name
166 )
167 else:
168 deleted = await models.concurrency_limits_v2.delete_concurrency_limit(
169 session, name=id_or_name
170 )
172 if not deleted:
173 raise HTTPException(
174 status_code=status.HTTP_404_NOT_FOUND, detail="Concurrency Limit not found"
175 )
178class MinimalConcurrencyLimitResponse(PrefectBaseModel):
179 id: UUID
180 name: str
181 limit: int
184class ConcurrencyLimitWithLeaseResponse(PrefectBaseModel):
185 lease_id: UUID
186 limits: list[MinimalConcurrencyLimitResponse]
189async def _acquire_concurrency_slots(
190 session: AsyncSession,
191 names: List[str],
192 slots: int,
193 mode: Literal["concurrency", "rate_limit"],
194) -> tuple[list[schemas.core.ConcurrencyLimitV2], bool]:
195 limits = [
196 schemas.core.ConcurrencyLimitV2.model_validate(limit)
197 for limit in (
198 await models.concurrency_limits_v2.bulk_read_concurrency_limits(
199 session=session, names=names
200 )
201 )
202 ]
204 active_limits = [limit for limit in limits if bool(limit.active)]
206 if any(limit.limit < slots for limit in active_limits): 206 ↛ anywhereline 206 didn't jump anywhere: it always raised an exception.
207 raise HTTPException(
208 status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
209 detail="Slots requested is greater than the limit",
210 )
212 non_decaying = [
213 str(limit.name) for limit in active_limits if limit.slot_decay_per_second == 0.0
214 ]
216 if mode == "rate_limit" and non_decaying:
217 raise HTTPException(
218 status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
219 detail=(
220 "Only concurrency limits with slot decay can be used for "
221 "rate limiting. The following limits do not have a decay "
222 f"configured: {','.join(non_decaying)!r}"
223 ),
224 )
225 acquired = await models.concurrency_limits_v2.bulk_increment_active_slots(
226 session=session,
227 concurrency_limit_ids=[limit.id for limit in active_limits],
228 slots=slots,
229 )
231 if not acquired:
232 await session.rollback()
234 return limits, acquired
237async def _generate_concurrency_locked_response(
238 session: AsyncSession,
239 limits: list[schemas.core.ConcurrencyLimitV2],
240 slots: int,
241) -> HTTPException:
242 """
243 Generate a 423 Locked response when concurrency slots cannot be acquired.
245 Calculates an appropriate Retry-After header value based on the blocking limit's
246 characteristics. For limits without slot decay, caps avg_slot_occupancy_seconds
247 at a configured maximum to prevent excessive retry delays from long-running tasks:
249 - Tag-based limits (name starts with "tag:"): Capped at tag_concurrency_slot_wait_seconds
250 (default 30s) to restore V1 behavior that users relied on
251 - Other limits: Capped at maximum_concurrency_slot_wait_seconds (default 30s) to allow
252 more uniform queues while still preventing astronomical delays
254 Low average occupancies are always respected (e.g., 2s stays 2s, not forced higher).
255 Limits with slot decay use the decay rate directly without capping.
257 The final retry value includes jitter via clamped_poisson_interval to prevent
258 thundering herd when many tasks retry simultaneously.
259 """
260 active_limits = [limit for limit in limits if bool(limit.active)]
262 await models.concurrency_limits_v2.bulk_update_denied_slots(
263 session=session,
264 concurrency_limit_ids=[limit.id for limit in active_limits],
265 slots=slots,
266 )
268 def num_blocking_slots(limit: schemas.core.ConcurrencyLimitV2) -> float:
269 if limit.slot_decay_per_second > 0:
270 return slots + limit.denied_slots
271 else:
272 return (slots + limit.denied_slots) / limit.limit
274 blocking_limit = max((limit for limit in active_limits), key=num_blocking_slots)
275 blocking_slots = num_blocking_slots(blocking_limit)
277 if blocking_limit.slot_decay_per_second == 0.0: 277 ↛ 289line 277 didn't jump to line 289 because the condition on line 277 was always true
278 settings = get_current_settings()
279 max_wait = (
280 settings.server.tasks.tag_concurrency_slot_wait_seconds
281 if blocking_limit.name.startswith("tag:")
282 else settings.server.concurrency.maximum_concurrency_slot_wait_seconds
283 )
284 wait_time_per_slot = min(blocking_limit.avg_slot_occupancy_seconds, max_wait)
285 # Cap the total wait time at max_wait to prevent excessive retry delays
286 # when denied_slots accumulates from burst traffic
287 average_interval = min(wait_time_per_slot * blocking_slots, max_wait)
288 else:
289 wait_time_per_slot = 1.0 / blocking_limit.slot_decay_per_second
290 average_interval = wait_time_per_slot * blocking_slots
292 retry_after = clamped_poisson_interval(average_interval=average_interval)
294 return HTTPException(
295 status_code=status.HTTP_423_LOCKED,
296 headers={
297 "Retry-After": str(retry_after),
298 },
299 )
302@router.post("/increment", status_code=status.HTTP_200_OK)
303async def bulk_increment_active_slots(
304 slots: int = Body(..., gt=0),
305 names: List[str] = Body(..., min_items=1),
306 mode: Literal["concurrency", "rate_limit"] = Body("concurrency"),
307 create_if_missing: Optional[bool] = Body(
308 None,
309 deprecated="Limits must be explicitly created before acquiring concurrency slots.",
310 ),
311 db: PrefectDBInterface = Depends(provide_database_interface),
312) -> List[MinimalConcurrencyLimitResponse]:
313 async with db.session_context(begin_transaction=True) as session:
314 acquired_limits, acquired = await _acquire_concurrency_slots(
315 session=session,
316 names=names,
317 slots=slots,
318 mode=mode,
319 )
321 if acquired: 321 ↛ 329line 321 didn't jump to line 329 because the condition on line 321 was always true
322 return [
323 MinimalConcurrencyLimitResponse(
324 id=limit.id, name=str(limit.name), limit=limit.limit
325 )
326 for limit in acquired_limits
327 ]
328 else:
329 async with db.session_context(begin_transaction=True) as session:
330 raise await _generate_concurrency_locked_response(
331 session=session,
332 limits=acquired_limits,
333 slots=slots,
334 )
337@router.post("/increment-with-lease", status_code=status.HTTP_200_OK)
338async def bulk_increment_active_slots_with_lease(
339 slots: int = Body(..., gt=0),
340 names: List[str] = Body(..., min_items=1),
341 mode: Literal["concurrency", "rate_limit"] = Body("concurrency"),
342 lease_duration: float = Body(
343 300, # 5 minutes
344 ge=60, # 1 minute
345 le=60 * 60 * 24, # 1 day
346 description="The duration of the lease in seconds.",
347 ),
348 holder: Optional[ConcurrencyLeaseHolder] = Body(
349 None,
350 description="The holder of the lease with type (flow_run, task_run, or deployment) and id.",
351 ),
352 db: PrefectDBInterface = Depends(provide_database_interface),
353) -> ConcurrencyLimitWithLeaseResponse:
354 async with db.session_context(begin_transaction=True) as session:
355 acquired_limits, acquired = await _acquire_concurrency_slots(
356 session=session,
357 names=names,
358 slots=slots,
359 mode=mode,
360 )
362 if acquired: 362 ↛ 383line 362 didn't jump to line 383 because the condition on line 362 was always true
363 lease_storage = get_concurrency_lease_storage()
364 lease = await lease_storage.create_lease(
365 resource_ids=[limit.id for limit in acquired_limits],
366 ttl=timedelta(seconds=lease_duration),
367 metadata=ConcurrencyLimitLeaseMetadata(
368 slots=slots,
369 holder=holder,
370 ),
371 )
372 return ConcurrencyLimitWithLeaseResponse(
373 lease_id=lease.id,
374 limits=[
375 MinimalConcurrencyLimitResponse(
376 id=limit.id, name=str(limit.name), limit=limit.limit
377 )
378 for limit in acquired_limits
379 ],
380 )
382 else:
383 async with db.session_context(begin_transaction=True) as session:
384 raise await _generate_concurrency_locked_response(
385 session=session,
386 limits=acquired_limits,
387 slots=slots,
388 )
391@router.post("/decrement", status_code=status.HTTP_200_OK)
392async def bulk_decrement_active_slots(
393 slots: int = Body(..., gt=0),
394 names: List[str] = Body(..., min_items=1),
395 occupancy_seconds: Optional[float] = Body(None, gt=0.0),
396 create_if_missing: bool = Body(
397 None,
398 deprecated="Limits must be explicitly created before decrementing active slots.",
399 ),
400 db: PrefectDBInterface = Depends(provide_database_interface),
401) -> List[MinimalConcurrencyLimitResponse]:
402 async with db.session_context(begin_transaction=True) as session:
403 limits = await models.concurrency_limits_v2.bulk_read_concurrency_limits(
404 session=session, names=names
405 )
407 if not limits:
408 return []
410 await models.concurrency_limits_v2.bulk_decrement_active_slots(
411 session=session,
412 concurrency_limit_ids=[limit.id for limit in limits if bool(limit.active)],
413 slots=slots,
414 occupancy_seconds=occupancy_seconds,
415 )
417 return [
418 MinimalConcurrencyLimitResponse(
419 id=limit.id, name=str(limit.name), limit=limit.limit
420 )
421 for limit in limits
422 ]
425@router.post("/decrement-with-lease", status_code=status.HTTP_204_NO_CONTENT)
426async def bulk_decrement_active_slots_with_lease(
427 lease_id: UUID = Body(
428 ...,
429 description="The ID of the lease corresponding to the concurrency limits to decrement.",
430 embed=True,
431 ),
432 db: PrefectDBInterface = Depends(provide_database_interface),
433) -> None:
434 lease_storage = get_concurrency_lease_storage()
435 lease = await lease_storage.read_lease(lease_id)
436 if not lease:
437 return
439 occupancy_seconds = (datetime.now(timezone.utc) - lease.created_at).total_seconds()
441 async with db.session_context(begin_transaction=True) as session:
442 await models.concurrency_limits_v2.bulk_decrement_active_slots(
443 session=session,
444 concurrency_limit_ids=lease.resource_ids,
445 slots=lease.metadata.slots if lease.metadata else 0,
446 occupancy_seconds=occupancy_seconds,
447 )
448 await lease_storage.revoke_lease(lease_id)
451@router.post("/leases/{lease_id}/renew", status_code=status.HTTP_204_NO_CONTENT)
452async def renew_concurrency_lease(
453 lease_id: UUID = Path(..., description="The ID of the lease to renew"),
454 lease_duration: float = Body(
455 300, # 5 minutes
456 ge=60, # 1 minute
457 le=60 * 60 * 24, # 1 day
458 description="The duration of the lease in seconds.",
459 embed=True,
460 ),
461) -> None:
462 lease_storage = get_concurrency_lease_storage()
464 # Atomically renew the lease (checks existence and updates index in single operation)
465 renewed = await lease_storage.renew_lease(
466 lease_id=lease_id,
467 ttl=timedelta(seconds=lease_duration),
468 )
470 # Handle the three possible return values:
471 # - True: lease successfully renewed
472 # - False: lease not found or expired
473 # - None: legacy implementation (check lease existence to determine success)
474 lease = None
475 if renewed is None: 475 ↛ 477line 475 didn't jump to line 477 because the condition on line 475 was never true
476 # Legacy implementation returned None - check if lease actually exists
477 lease = await lease_storage.read_lease(lease_id)
479 if renewed is False or (renewed is None and lease is None): 479 ↛ exitline 479 didn't return from function 'renew_concurrency_lease' because the condition on line 479 was always true
480 raise HTTPException(
481 status_code=status.HTTP_410_GONE,
482 detail="Lease not found - it may have expired or been revoked",
483 )