Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/api/workers.py: 43%
354 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1"""
2Routes for interacting with work queue objects.
3"""
5from dataclasses import dataclass
6from logging import Logger
7from typing import TYPE_CHECKING, Any, List, Optional
8from uuid import UUID
10import sqlalchemy as sa
11from fastapi import (
12 Body,
13 Depends,
14 HTTPException,
15 Path,
16 WebSocket,
17 status,
18)
19from packaging.version import Version
20from pydantic import ValidationError
21from sqlalchemy.ext.asyncio import AsyncSession
23import prefect.server.api.dependencies as dependencies
24import prefect.server.models as models
25import prefect.server.schemas as schemas
26from prefect._internal.uuid7 import uuid7
27from prefect.client.schemas.worker_channel import (
28 CLEANUP_DELIVERY_CAPABILITY,
29 WORK_POOL_SNAPSHOT_CAPABILITY,
30 WORKER_HEARTBEAT_CAPABILITY,
31 WorkerChannelCapability,
32 WorkerChannelCloseReason,
33 WorkerChannelProtocolError,
34 WorkerHelloFrame,
35 WorkerReadyFrame,
36 WorkPoolSnapshotPayload,
37 select_worker_channel_version,
38 validate_worker_channel_frame,
39)
40from prefect.logging import get_logger
41from prefect.server.api.validation import validate_job_variable_defaults_for_work_pool
42from prefect.server.database import PrefectDBInterface, provide_database_interface
43from prefect.server.models.deployments import mark_deployments_ready
44from prefect.server.models.work_queues import (
45 emit_work_queue_status_event,
46 mark_work_queues_ready,
47)
48from prefect.server.models.workers import emit_work_pool_status_event
49from prefect.server.schemas.statuses import WorkQueueStatus
50from prefect.server.utilities import subscriptions
51from prefect.server.utilities import worker_channel as worker_channel_utils
52from prefect.server.utilities.server import PrefectRouter
53from prefect.server.worker_communication.cleanup_queue import (
54 WorkerCleanupQueue,
55 get_worker_cleanup_queue,
56)
57from prefect.types import DateTime
58from prefect.types._datetime import now
60if TYPE_CHECKING: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true
61 from prefect.server.database.orm_models import WorkPool as ORMWorkPool
62 from prefect.server.database.orm_models import WorkQueue as ORMWorkQueue
64router: PrefectRouter = PrefectRouter(
65 prefix="/work_pools",
66 tags=["Work Pools"],
67)
68logger: Logger = get_logger("prefect.server.api.workers")
70_OSS_WORKER_CHANNEL_REQUIRED_CAPABILITIES: list[WorkerChannelCapability] = [
71 WORKER_HEARTBEAT_CAPABILITY,
72 WORK_POOL_SNAPSHOT_CAPABILITY,
73]
76# -----------------------------------------------------
77# --
78# --
79# -- Utility functions & dependencies
80# --
81# --
82# -----------------------------------------------------
85class WorkerLookups:
86 async def _get_work_pool_id_from_name(
87 self, session: AsyncSession, work_pool_name: str
88 ) -> UUID:
89 """
90 Given a work pool name, return its ID. Used for translating
91 user-facing APIs (which are name-based) to internal ones (which are
92 id-based).
93 """
94 work_pool = await models.workers.read_work_pool_by_name(
95 session=session,
96 work_pool_name=work_pool_name,
97 )
98 if not work_pool:
99 raise HTTPException(
100 status_code=status.HTTP_404_NOT_FOUND,
101 detail=f'Work pool "{work_pool_name}" not found.',
102 )
104 return work_pool.id
106 async def _get_default_work_queue_id_from_work_pool_name(
107 self, session: AsyncSession, work_pool_name: str
108 ):
109 """
110 Given a work pool name, return the ID of its default queue.
111 Used for translating user-facing APIs (which are name-based)
112 to internal ones (which are id-based).
113 """
114 work_pool = await models.workers.read_work_pool_by_name(
115 session=session,
116 work_pool_name=work_pool_name,
117 )
118 if not work_pool:
119 raise HTTPException(
120 status_code=status.HTTP_404_NOT_FOUND,
121 detail=f'Work pool "{work_pool_name}" not found.',
122 )
124 return work_pool.default_queue_id
126 async def _get_work_queue_from_name(
127 self,
128 session: AsyncSession,
129 work_pool_name: str,
130 work_queue_name: str,
131 create_queue_if_not_found: bool = False,
132 ) -> "ORMWorkQueue":
133 """
134 Given a work pool name and work pool queue name, return the ID of the
135 queue. Used for translating user-facing APIs (which are name-based) to
136 internal ones (which are id-based).
137 """
138 work_queue = await models.workers.read_work_queue_by_name(
139 session=session,
140 work_pool_name=work_pool_name,
141 work_queue_name=work_queue_name,
142 )
143 if not work_queue:
144 if not create_queue_if_not_found:
145 raise HTTPException(
146 status_code=status.HTTP_404_NOT_FOUND,
147 detail=(
148 f"Work pool queue '{work_pool_name}/{work_queue_name}' not"
149 " found."
150 ),
151 )
152 work_pool_id = await self._get_work_pool_id_from_name(
153 session=session, work_pool_name=work_pool_name
154 )
155 work_queue = await models.workers.create_work_queue(
156 session=session,
157 work_pool_id=work_pool_id,
158 work_queue=schemas.actions.WorkQueueCreate(name=work_queue_name),
159 )
161 return work_queue
163 async def _get_work_queue_id_from_name(
164 self,
165 session: AsyncSession,
166 work_pool_name: str,
167 work_queue_name: str,
168 create_queue_if_not_found: bool = False,
169 ) -> UUID:
170 queue = await self._get_work_queue_from_name(
171 session=session,
172 work_pool_name=work_pool_name,
173 work_queue_name=work_queue_name,
174 create_queue_if_not_found=create_queue_if_not_found,
175 )
176 return queue.id
179class WorkerChannelSetupError(Exception):
180 def __init__(self, close_reason: WorkerChannelCloseReason, detail: str):
181 super().__init__(detail)
182 self.close_reason = close_reason
183 self.detail = detail
186@dataclass(frozen=True)
187class WorkerChannelWorkPoolUpdateEvent:
188 work_pool_id: UUID
189 changed_fields: dict[str, dict[str, Any]]
192def _worker_requested_cleanup_delivery(hello: WorkerHelloFrame) -> bool:
193 return (
194 CLEANUP_DELIVERY_CAPABILITY in hello.payload.requested_capabilities
195 and bool(hello.payload.handled_cleanup_kinds)
196 and hello.payload.max_cleanup_concurrency > 0
197 )
200def _accepted_worker_channel_capabilities(
201 hello: WorkerHelloFrame,
202 *,
203 cleanup_queue_available: bool,
204) -> list[WorkerChannelCapability]:
205 accepted = list(_OSS_WORKER_CHANNEL_REQUIRED_CAPABILITIES)
206 if cleanup_queue_available and _worker_requested_cleanup_delivery(hello):
207 accepted.append(CLEANUP_DELIVERY_CAPABILITY)
209 return accepted
212async def _receive_worker_hello(websocket: WebSocket) -> WorkerHelloFrame:
213 try:
214 message = await websocket.receive_json()
215 frame = validate_worker_channel_frame(message)
216 except ValidationError as exc:
217 raise WorkerChannelSetupError(
218 WorkerChannelCloseReason.PROTOCOL_ERROR,
219 "Worker channel received a malformed hello frame",
220 ) from exc
221 except ValueError as exc:
222 raise WorkerChannelSetupError(
223 WorkerChannelCloseReason.PROTOCOL_ERROR,
224 "Worker channel received invalid JSON during setup",
225 ) from exc
227 if not isinstance(frame, WorkerHelloFrame):
228 raise WorkerChannelSetupError(
229 WorkerChannelCloseReason.PROTOCOL_ERROR,
230 "Expected worker.hello.v1 during worker channel setup",
231 )
233 return frame
236async def _resolve_worker_channel_work_pool(
237 session: AsyncSession,
238 work_pool_name: str,
239 hello: WorkerHelloFrame,
240) -> "ORMWorkPool":
241 work_pool = await models.workers.read_work_pool_by_name(
242 session=session,
243 work_pool_name=work_pool_name,
244 )
246 default_base_job_template = hello.payload.default_base_job_template
247 if work_pool is None:
248 if not hello.payload.create_pool_if_not_found:
249 raise WorkerChannelSetupError(
250 WorkerChannelCloseReason.AUTHORIZATION_FAILED,
251 "work_pool_not_found",
252 )
254 if work_pool_name.lower().startswith("prefect"):
255 raise WorkerChannelSetupError(
256 WorkerChannelCloseReason.AUTHORIZATION_FAILED,
257 "work_pool_creation_unauthorized",
258 )
260 await validate_job_variable_defaults_for_work_pool(
261 session, work_pool_name, default_base_job_template
262 )
263 try:
264 async with session.begin_nested():
265 work_pool = await models.workers.create_work_pool(
266 session=session,
267 work_pool=schemas.actions.WorkPoolCreate(
268 name=work_pool_name,
269 type=hello.payload.worker_type,
270 base_job_template=default_base_job_template,
271 ),
272 )
273 except sa.exc.IntegrityError:
274 work_pool = await models.workers.read_work_pool_by_name(
275 session=session,
276 work_pool_name=work_pool_name,
277 )
278 if work_pool is None:
279 raise
280 return work_pool
282 return work_pool
285async def _resolve_worker_channel_work_queues(
286 session: AsyncSession,
287 work_pool_id: UUID,
288 work_pool_name: str,
289 work_queue_names: list[str],
290) -> list["ORMWorkQueue"]:
291 if not work_queue_names:
292 return list(
293 await models.workers.read_work_queues(
294 session=session, work_pool_id=work_pool_id
295 )
296 )
298 work_queues = []
299 for work_queue_name in dict.fromkeys(work_queue_names):
300 work_queue = await models.workers.read_work_queue_by_name(
301 session=session,
302 work_pool_name=work_pool_name,
303 work_queue_name=work_queue_name,
304 )
305 if work_queue is None:
306 raise WorkerChannelSetupError(
307 WorkerChannelCloseReason.AUTHORIZATION_FAILED,
308 "work_queue_not_found",
309 )
310 work_queues.append(work_queue)
312 return work_queues
315async def _build_worker_ready_frame(
316 session: AsyncSession,
317 work_pool_name: str,
318 hello: WorkerHelloFrame,
319 cleanup_queue_available: bool,
320) -> tuple[WorkerReadyFrame, WorkerChannelWorkPoolUpdateEvent | None]:
321 try:
322 selected_channel_version = select_worker_channel_version(
323 hello.payload.supported_channel_versions
324 )
325 except WorkerChannelProtocolError as exc:
326 raise WorkerChannelSetupError(exc.close_reason, str(exc)) from exc
328 work_pool = await _resolve_worker_channel_work_pool(
329 session=session,
330 work_pool_name=work_pool_name,
331 hello=hello,
332 )
333 work_queues = await _resolve_worker_channel_work_queues(
334 session=session,
335 work_pool_id=work_pool.id,
336 work_pool_name=work_pool_name,
337 work_queue_names=hello.payload.work_queue_names,
338 )
339 default_base_job_template = hello.payload.default_base_job_template
340 work_pool_update_event = None
341 if not work_pool.base_job_template and default_base_job_template:
342 previous_base_job_template = work_pool.base_job_template
343 await validate_job_variable_defaults_for_work_pool(
344 session, work_pool_name, default_base_job_template
345 )
346 updated = await models.workers.update_work_pool(
347 session=session,
348 work_pool_id=work_pool.id,
349 work_pool=schemas.actions.WorkPoolUpdate(
350 base_job_template=default_base_job_template
351 ),
352 emit_update_event=False,
353 emit_status_change=emit_work_pool_status_event,
354 )
355 if updated:
356 work_pool_update_event = WorkerChannelWorkPoolUpdateEvent(
357 work_pool_id=work_pool.id,
358 changed_fields={
359 "base_job_template": {
360 "from": previous_base_job_template,
361 "to": default_base_job_template,
362 }
363 },
364 )
365 refreshed = await models.workers.read_work_pool(
366 session=session, work_pool_id=work_pool.id
367 )
368 assert refreshed is not None
369 work_pool = refreshed
371 try:
372 worker = await models.workers.record_worker_heartbeat(
373 session=session,
374 work_pool=work_pool,
375 worker_name=hello.payload.worker_name,
376 heartbeat_interval_seconds=hello.payload.heartbeat_interval_seconds,
377 emit_status_change=emit_work_pool_status_event,
378 return_worker=True,
379 )
380 except Exception as exc:
381 raise WorkerChannelSetupError(
382 WorkerChannelCloseReason.HEARTBEAT_PERSISTENCE_FAILED,
383 "worker_channel_initial_heartbeat_failed",
384 ) from exc
385 assert worker is not None
387 refreshed_work_pool = await models.workers.read_work_pool(
388 session=session, work_pool_id=work_pool.id
389 )
390 assert refreshed_work_pool is not None
391 initial_snapshot = WorkPoolSnapshotPayload(
392 snapshot_sequence=1,
393 reason="initial",
394 work_pool=await worker_channel_utils.build_worker_channel_work_pool_snapshot(
395 session=session,
396 work_pool=refreshed_work_pool,
397 ),
398 )
400 requested_capabilities = list(dict.fromkeys(hello.payload.requested_capabilities))
401 accepted = _accepted_worker_channel_capabilities(
402 hello,
403 cleanup_queue_available=cleanup_queue_available,
404 )
405 accepted_set = set(accepted)
406 rejected = [
407 capability
408 for capability in requested_capabilities
409 if capability not in accepted_set
410 ]
412 if rejected:
413 logger.debug(
414 "Worker channel capabilities rejected: "
415 "work_pool=%s worker_name=%s rejected=%s",
416 work_pool_name,
417 hello.payload.worker_name,
418 rejected,
419 )
421 return (
422 WorkerReadyFrame(
423 type="worker.ready.v1",
424 id=uuid7(),
425 sent_at=now("UTC"),
426 payload={
427 "consumer_id": hello.payload.consumer_id,
428 "worker_id": None,
429 "selected_channel_version": selected_channel_version,
430 "effective_heartbeat_interval_seconds": (
431 hello.payload.heartbeat_interval_seconds
432 ),
433 "accepted_capabilities": accepted,
434 "rejected_capabilities": rejected,
435 "effective_max_cleanup_concurrency": (
436 hello.payload.max_cleanup_concurrency
437 if CLEANUP_DELIVERY_CAPABILITY in accepted_set
438 else 0
439 ),
440 "resolved_work_queues": [
441 {"id": work_queue.id, "name": work_queue.name}
442 for work_queue in work_queues
443 ],
444 "initial_snapshot": initial_snapshot,
445 },
446 ),
447 work_pool_update_event,
448 )
451# -----------------------------------------------------
452# --
453# --
454# -- Worker Pools
455# --
456# --
457# -----------------------------------------------------
460@router.post("/", status_code=status.HTTP_201_CREATED)
461async def create_work_pool(
462 work_pool: schemas.actions.WorkPoolCreate,
463 db: PrefectDBInterface = Depends(provide_database_interface),
464 prefect_client_version: Optional[str] = Depends(
465 dependencies.get_prefect_client_version
466 ),
467) -> schemas.responses.WorkPoolResponse:
468 """
469 Creates a new work pool. If a work pool with the same
470 name already exists, an error will be raised.
472 For more information, see https://docs.prefect.io/v3/concepts/work-pools.
473 """
474 if work_pool.name.lower().startswith("prefect"): 474 ↛ 475line 474 didn't jump to line 475 because the condition on line 474 was never true
475 raise HTTPException(
476 status_code=status.HTTP_403_FORBIDDEN,
477 detail="Work pools starting with 'Prefect' are reserved for internal use.",
478 )
480 try:
481 async with db.session_context(begin_transaction=True) as session:
482 await validate_job_variable_defaults_for_work_pool(
483 session, work_pool.name, work_pool.base_job_template
484 )
485 model = await models.workers.create_work_pool(
486 session=session, work_pool=work_pool
487 )
489 await emit_work_pool_status_event(
490 event_id=uuid7(),
491 occurred=now("UTC"),
492 pre_update_work_pool=None,
493 work_pool=model,
494 )
496 ret = schemas.responses.WorkPoolResponse.model_validate(
497 model, from_attributes=True
498 )
499 if ret.concurrency_limit is not None:
500 ret.active_slots = 0
501 if prefect_client_version and Version(prefect_client_version) <= Version(
502 "3.3.7"
503 ):
504 # Client versions 3.3.7 and below do not support the default_result_storage_block_id field and will error
505 # when receiving it.
506 del ret.storage_configuration.default_result_storage_block_id
507 return ret
509 except sa.exc.IntegrityError:
510 raise HTTPException(
511 status_code=status.HTTP_409_CONFLICT,
512 detail="A work pool with this name already exists.",
513 )
516@router.get("/{name}")
517async def read_work_pool(
518 work_pool_name: str = Path(..., description="The work pool name", alias="name"),
519 worker_lookups: WorkerLookups = Depends(WorkerLookups),
520 db: PrefectDBInterface = Depends(provide_database_interface),
521 prefect_client_version: Optional[str] = Depends(
522 dependencies.get_prefect_client_version
523 ),
524) -> schemas.responses.WorkPoolResponse:
525 """
526 Read a work pool by name
527 """
529 async with db.session_context() as session:
530 work_pool_id = await worker_lookups._get_work_pool_id_from_name(
531 session=session, work_pool_name=work_pool_name
532 )
533 orm_work_pool = await models.workers.read_work_pool(
534 session=session, work_pool_id=work_pool_id
535 )
536 work_pool = schemas.responses.WorkPoolResponse.model_validate(
537 orm_work_pool, from_attributes=True
538 )
540 if work_pool.concurrency_limit is not None:
541 work_pool.active_slots = await models.workers.count_work_pool_active_slots(
542 session=session, work_pool_id=work_pool_id
543 )
545 if prefect_client_version and Version(prefect_client_version) <= Version(
546 "3.3.7"
547 ):
548 # Client versions 3.3.7 and below do not support the default_result_storage_block_id field and will error
549 # when receiving it.
550 del work_pool.storage_configuration.default_result_storage_block_id
552 return work_pool
555@router.post("/filter")
556async def read_work_pools(
557 work_pools: Optional[schemas.filters.WorkPoolFilter] = None,
558 limit: int = dependencies.LimitBody(),
559 offset: int = Body(0, ge=0),
560 db: PrefectDBInterface = Depends(provide_database_interface),
561 prefect_client_version: Optional[str] = Depends(
562 dependencies.get_prefect_client_version
563 ),
564) -> List[schemas.responses.WorkPoolResponse]:
565 """
566 Read multiple work pools
567 """
568 async with db.session_context() as session:
569 orm_work_pools = await models.workers.read_work_pools(
570 session=session,
571 work_pool_filter=work_pools,
572 offset=offset,
573 limit=limit,
574 )
575 ret = [
576 schemas.responses.WorkPoolResponse.model_validate(w, from_attributes=True)
577 for w in orm_work_pools
578 ]
579 pools_with_limit = [wp for wp in ret if wp.concurrency_limit is not None]
580 if pools_with_limit:
581 slot_counts = await models.workers.count_work_pool_active_slots_bulk(
582 session=session,
583 work_pool_ids=[wp.id for wp in pools_with_limit],
584 )
585 for work_pool in pools_with_limit:
586 work_pool.active_slots = slot_counts.get(work_pool.id, 0)
587 if prefect_client_version and Version(prefect_client_version) <= Version(
588 "3.3.7"
589 ):
590 # Client versions 3.3.7 and below do not support the default_result_storage_block_id field and will error
591 # when receiving it.
592 for work_pool in ret:
593 del work_pool.storage_configuration.default_result_storage_block_id
594 return ret
597@router.post("/count")
598async def count_work_pools(
599 work_pools: Optional[schemas.filters.WorkPoolFilter] = Body(None, embed=True),
600 db: PrefectDBInterface = Depends(provide_database_interface),
601) -> int:
602 """
603 Count work pools
604 """
605 async with db.session_context() as session:
606 return await models.workers.count_work_pools(
607 session=session, work_pool_filter=work_pools
608 )
611@router.patch("/{name}", status_code=status.HTTP_204_NO_CONTENT)
612async def update_work_pool(
613 work_pool: schemas.actions.WorkPoolUpdate,
614 work_pool_name: str = Path(..., description="The work pool name", alias="name"),
615 worker_lookups: WorkerLookups = Depends(WorkerLookups),
616 db: PrefectDBInterface = Depends(provide_database_interface),
617) -> None:
618 """
619 Update a work pool
620 """
622 # Reserved pools can only updated pause / concurrency
623 update_values = work_pool.model_dump(exclude_unset=True)
624 if work_pool_name.lower().startswith("prefect") and ( 624 ↛ 627line 624 didn't jump to line 627 because the condition on line 624 was never true
625 set(update_values).difference({"is_paused", "concurrency_limit"})
626 ):
627 raise HTTPException(
628 status_code=status.HTTP_403_FORBIDDEN,
629 detail=(
630 "Work pools starting with 'Prefect' are reserved for internal use "
631 "and can only be updated to set concurrency limits or pause."
632 ),
633 )
635 async with db.session_context(begin_transaction=True) as session:
636 work_pool_id = await worker_lookups._get_work_pool_id_from_name(
637 session=session, work_pool_name=work_pool_name
638 )
639 updated = await models.workers.update_work_pool(
640 session=session,
641 work_pool_id=work_pool_id,
642 work_pool=work_pool,
643 emit_status_change=emit_work_pool_status_event,
644 )
646 if updated and worker_channel_utils.work_pool_update_triggers_snapshot( 646 ↛ 649line 646 didn't jump to line 649 because the condition on line 646 was never true
647 update_values
648 ):
649 await worker_channel_utils.publish_snapshot_invalidation(
650 worker_channel_utils.WorkerChannelSnapshotInvalidation(
651 work_pool_id=work_pool_id,
652 reason="work_pool_updated",
653 )
654 )
657@router.delete("/{name}", status_code=status.HTTP_204_NO_CONTENT)
658async def delete_work_pool(
659 work_pool_name: str = Path(..., description="The work pool name", alias="name"),
660 worker_lookups: WorkerLookups = Depends(WorkerLookups),
661 db: PrefectDBInterface = Depends(provide_database_interface),
662) -> None:
663 """
664 Delete a work pool
665 """
667 if work_pool_name.lower().startswith("prefect"): 667 ↛ 668line 667 didn't jump to line 668 because the condition on line 667 was never true
668 raise HTTPException(
669 status_code=status.HTTP_403_FORBIDDEN,
670 detail=(
671 "Work pools starting with 'Prefect' are reserved for internal use and"
672 " can not be deleted."
673 ),
674 )
676 async with db.session_context(begin_transaction=True) as session:
677 work_pool_id = await worker_lookups._get_work_pool_id_from_name(
678 session=session, work_pool_name=work_pool_name
679 )
681 deleted = await models.workers.delete_work_pool(
682 session=session, work_pool_id=work_pool_id
683 )
685 if deleted: 685 ↛ exitline 685 didn't return from function 'delete_work_pool' because the condition on line 685 was always true
686 await worker_channel_utils.publish_snapshot_invalidation(
687 worker_channel_utils.WorkerChannelSnapshotInvalidation(
688 work_pool_id=work_pool_id,
689 reason="work_pool_deleted",
690 work_pool_deleted=True,
691 )
692 )
695@router.post("/{name}/concurrency_status")
696async def read_work_pool_concurrency_status(
697 work_pool_name: str = Path(..., description="The work pool name", alias="name"),
698 page: int = Body(1, ge=1),
699 limit: int = dependencies.LimitBody(),
700 flow_run_limit: int = Body(10, ge=0, le=200, description="Max flow runs per queue"),
701 worker_lookups: WorkerLookups = Depends(WorkerLookups),
702 db: PrefectDBInterface = Depends(provide_database_interface),
703) -> schemas.responses.WorkPoolConcurrencyStatus:
704 """
705 Read concurrency status for a work pool, including per-queue breakdown
706 with flow run summaries. Queues are paginated; flow runs per queue are
707 capped by flow_run_limit.
708 """
709 import asyncio
711 from prefect.types._datetime import now as prefect_now
713 queue_offset = (page - 1) * limit
715 async with db.session_context() as session:
716 work_pool = await models.workers.read_work_pool_by_name(
717 session=session, work_pool_name=work_pool_name
718 )
719 if not work_pool:
720 raise HTTPException(
721 status_code=status.HTTP_404_NOT_FOUND,
722 detail=f"Work pool {work_pool_name!r} not found.",
723 )
725 # Paginate queues in the DB and get total count + active slots
726 # concurrently
727 (
728 work_queues_page,
729 total_queue_count,
730 total_active,
731 counts_by_queue,
732 ) = await asyncio.gather(
733 models.workers.read_work_queues(
734 session=session,
735 work_pool_id=work_pool.id,
736 offset=queue_offset,
737 limit=limit,
738 ),
739 models.workers.count_work_queues(
740 session=session,
741 work_pool_id=work_pool.id,
742 ),
743 models.workers.count_work_pool_slot_holders(
744 session=session,
745 work_pool_id=work_pool.id,
746 ),
747 models.workers.count_work_pool_slot_holders_by_queue(
748 session=session,
749 work_pool_id=work_pool.id,
750 ),
751 )
753 # Only fetch flow run details for the queues on this page
754 page_queue_ids = [wq.id for wq in work_queues_page]
755 slot_holders = await models.workers.get_work_pool_slot_holders(
756 session=session,
757 work_pool_id=work_pool.id,
758 work_queue_ids=page_queue_ids,
759 flow_run_limit=flow_run_limit,
760 )
762 current_time = prefect_now("UTC")
764 # Group flow runs by work queue id
765 runs_by_queue: dict[UUID, list[tuple]] = {}
766 for run, slot_acquired_at in slot_holders: 766 ↛ 767line 766 didn't jump to line 767 because the loop on line 766 never started
767 queue_id = run.work_queue_id
768 if queue_id is not None:
769 runs_by_queue.setdefault(queue_id, []).append((run, slot_acquired_at))
771 def _build_summary(run, slot_acquired_at) -> schemas.responses.FlowRunSlotSummary:
772 state_ts = run.state_timestamp
773 return schemas.responses.FlowRunSlotSummary(
774 id=run.id,
775 name=run.name,
776 state_type=run.state_type if run.state_type else None,
777 state_name=run.state_name if run.state_name else None,
778 start_time=run.start_time,
779 state_timestamp=state_ts,
780 time_in_current_state=((current_time - state_ts) if state_ts else None),
781 )
783 queue_details = []
784 for wq in work_queues_page:
785 display_tuples = runs_by_queue.get(wq.id, [])
786 active = counts_by_queue.get(wq.id, 0)
787 queue_details.append(
788 schemas.responses.WorkQueueConcurrencyStatusDetail(
789 queue_id=wq.id,
790 queue_name=wq.name,
791 active_slots=active,
792 concurrency_limit=wq.concurrency_limit,
793 flow_runs=[_build_summary(r, sa) for r, sa in display_tuples],
794 flow_run_count=active,
795 )
796 )
798 return schemas.responses.WorkPoolConcurrencyStatus(
799 active_slots=total_active,
800 concurrency_limit=work_pool.concurrency_limit,
801 queues=queue_details,
802 count=total_queue_count,
803 limit=limit,
804 pages=(total_queue_count + limit - 1) // limit if limit > 0 else 0,
805 page=page,
806 )
809@router.post("/{name}/get_scheduled_flow_runs")
810async def get_scheduled_flow_runs(
811 docket: dependencies.Docket,
812 work_pool_name: str = Path(..., description="The work pool name", alias="name"),
813 work_queue_names: List[str] = Body(
814 None, description="The names of work pool queues"
815 ),
816 scheduled_before: DateTime = Body(
817 None, description="The maximum time to look for scheduled flow runs"
818 ),
819 scheduled_after: DateTime = Body(
820 None, description="The minimum time to look for scheduled flow runs"
821 ),
822 limit: int = dependencies.LimitBody(),
823 worker_lookups: WorkerLookups = Depends(WorkerLookups),
824 db: PrefectDBInterface = Depends(provide_database_interface),
825) -> List[schemas.responses.WorkerFlowRunResponse]:
826 """
827 Load scheduled runs for a worker
828 """
829 async with db.session_context() as session:
830 work_pool_id = await worker_lookups._get_work_pool_id_from_name(
831 session=session, work_pool_name=work_pool_name
832 )
834 if not work_queue_names:
835 work_queues = list(
836 await models.workers.read_work_queues(
837 session=session, work_pool_id=work_pool_id
838 )
839 )
840 # None here instructs get_scheduled_flow_runs to use the default behavior
841 # of just operating on all work queues of the pool
842 work_queue_ids = None
843 else:
844 work_queues = [
845 await worker_lookups._get_work_queue_from_name(
846 session=session,
847 work_pool_name=work_pool_name,
848 work_queue_name=name,
849 )
850 for name in work_queue_names
851 ]
852 work_queue_ids = [wq.id for wq in work_queues]
854 async with db.session_context(begin_transaction=True) as session:
855 queue_response = await models.workers.get_scheduled_flow_runs(
856 session=session,
857 work_pool_ids=[work_pool_id],
858 work_queue_ids=work_queue_ids,
859 scheduled_before=scheduled_before,
860 scheduled_after=scheduled_after,
861 limit=limit,
862 )
864 await docket.add(
865 mark_work_queues_ready,
866 key=f"mark_work_queues_ready:work_pool:{work_pool_id}",
867 )(
868 polled_work_queue_ids=[
869 wq.id for wq in work_queues if wq.status != WorkQueueStatus.NOT_READY
870 ],
871 ready_work_queue_ids=[
872 wq.id for wq in work_queues if wq.status == WorkQueueStatus.NOT_READY
873 ],
874 )
876 await docket.add(
877 mark_deployments_ready,
878 key=f"mark_deployments_ready:work_pool:{work_pool_id}",
879 )(
880 work_queue_ids=[wq.id for wq in work_queues],
881 )
883 return queue_response
886# -----------------------------------------------------
887# --
888# --
889# -- Work Pool Queues
890# --
891# --
892# -----------------------------------------------------
895@router.post("/{work_pool_name}/queues", status_code=status.HTTP_201_CREATED)
896async def create_work_queue(
897 work_queue: schemas.actions.WorkQueueCreate,
898 work_pool_name: str = Path(..., description="The work pool name"),
899 worker_lookups: WorkerLookups = Depends(WorkerLookups),
900 db: PrefectDBInterface = Depends(provide_database_interface),
901) -> schemas.responses.WorkQueueResponse:
902 """
903 Creates a new work pool queue. If a work pool queue with the same
904 name already exists, an error will be raised.
906 For more information, see https://docs.prefect.io/v3/concepts/work-pools#work-queues.
907 """
909 try:
910 async with db.session_context(begin_transaction=True) as session:
911 work_pool_id = await worker_lookups._get_work_pool_id_from_name(
912 session=session,
913 work_pool_name=work_pool_name,
914 )
916 model = await models.workers.create_work_queue(
917 session=session,
918 work_pool_id=work_pool_id,
919 work_queue=work_queue,
920 )
922 response = schemas.responses.WorkQueueResponse.model_validate(
923 model, from_attributes=True
924 )
925 if response.concurrency_limit is not None:
926 response.active_slots = 0
927 except sa.exc.IntegrityError:
928 raise HTTPException(
929 status_code=status.HTTP_409_CONFLICT,
930 detail=(
931 "A work queue with this name already exists in work pool"
932 " {work_pool_name!r}."
933 ),
934 )
936 return response
939@router.get("/{work_pool_name}/queues/{name}")
940async def read_work_queue(
941 work_pool_name: str = Path(..., description="The work pool name"),
942 work_queue_name: str = Path(
943 ..., description="The work pool queue name", alias="name"
944 ),
945 worker_lookups: WorkerLookups = Depends(WorkerLookups),
946 db: PrefectDBInterface = Depends(provide_database_interface),
947) -> schemas.responses.WorkQueueResponse:
948 """
949 Read a work pool queue
950 """
952 async with db.session_context(begin_transaction=True) as session:
953 work_queue_id = await worker_lookups._get_work_queue_id_from_name(
954 session=session,
955 work_pool_name=work_pool_name,
956 work_queue_name=work_queue_name,
957 )
959 model = await models.workers.read_work_queue(
960 session=session, work_queue_id=work_queue_id
961 )
963 response = schemas.responses.WorkQueueResponse.model_validate(
964 model, from_attributes=True
965 )
967 if response.concurrency_limit is not None:
968 response.active_slots = await models.workers.count_work_queue_active_slots(
969 session=session, work_queue_id=work_queue_id
970 )
972 return response
975@router.post("/{work_pool_name}/queues/filter")
976async def read_work_queues(
977 work_pool_name: str = Path(..., description="The work pool name"),
978 work_queues: schemas.filters.WorkQueueFilter = None,
979 limit: int = dependencies.LimitBody(),
980 offset: int = Body(0, ge=0),
981 worker_lookups: WorkerLookups = Depends(WorkerLookups),
982 db: PrefectDBInterface = Depends(provide_database_interface),
983) -> List[schemas.responses.WorkQueueResponse]:
984 """
985 Read all work pool queues
986 """
987 async with db.session_context() as session:
988 work_pool_id = await worker_lookups._get_work_pool_id_from_name(
989 session=session,
990 work_pool_name=work_pool_name,
991 )
992 wqs = await models.workers.read_work_queues(
993 session=session,
994 work_pool_id=work_pool_id,
995 work_queue_filter=work_queues,
996 limit=limit,
997 offset=offset,
998 )
1000 ret = [
1001 schemas.responses.WorkQueueResponse.model_validate(wq, from_attributes=True)
1002 for wq in wqs
1003 ]
1004 queues_with_limit = [wq for wq in ret if wq.concurrency_limit is not None]
1005 if queues_with_limit:
1006 slot_counts = await models.workers.count_work_queue_active_slots_bulk(
1007 session=session,
1008 work_queue_ids=[wq.id for wq in queues_with_limit],
1009 )
1010 for wq_response in queues_with_limit:
1011 wq_response.active_slots = slot_counts.get(wq_response.id, 0)
1013 return ret
1016@router.patch("/{work_pool_name}/queues/{name}", status_code=status.HTTP_204_NO_CONTENT)
1017async def update_work_queue(
1018 work_queue: schemas.actions.WorkQueueUpdate,
1019 work_pool_name: str = Path(..., description="The work pool name"),
1020 work_queue_name: str = Path(
1021 ..., description="The work pool queue name", alias="name"
1022 ),
1023 worker_lookups: WorkerLookups = Depends(WorkerLookups),
1024 db: PrefectDBInterface = Depends(provide_database_interface),
1025) -> None:
1026 """
1027 Update a work pool queue
1028 """
1029 async with db.session_context(begin_transaction=True) as session:
1030 await worker_lookups._get_work_pool_id_from_name(
1031 session=session,
1032 work_pool_name=work_pool_name,
1033 )
1034 work_queue_id = await worker_lookups._get_work_queue_id_from_name(
1035 work_pool_name=work_pool_name,
1036 work_queue_name=work_queue_name,
1037 session=session,
1038 )
1040 await models.workers.update_work_queue(
1041 session=session,
1042 work_queue_id=work_queue_id,
1043 work_queue=work_queue,
1044 emit_status_change=emit_work_queue_status_event,
1045 )
1048@router.delete(
1049 "/{work_pool_name}/queues/{name}", status_code=status.HTTP_204_NO_CONTENT
1050)
1051async def delete_work_queue(
1052 work_pool_name: str = Path(..., description="The work pool name"),
1053 work_queue_name: str = Path(
1054 ..., description="The work pool queue name", alias="name"
1055 ),
1056 worker_lookups: WorkerLookups = Depends(WorkerLookups),
1057 db: PrefectDBInterface = Depends(provide_database_interface),
1058) -> None:
1059 """
1060 Delete a work pool queue
1061 """
1063 async with db.session_context(begin_transaction=True) as session:
1064 await worker_lookups._get_work_pool_id_from_name(
1065 session=session,
1066 work_pool_name=work_pool_name,
1067 )
1068 work_queue_id = await worker_lookups._get_work_queue_id_from_name(
1069 session=session,
1070 work_pool_name=work_pool_name,
1071 work_queue_name=work_queue_name,
1072 )
1074 await models.workers.delete_work_queue(
1075 session=session, work_queue_id=work_queue_id
1076 )
1079# -----------------------------------------------------
1080# --
1081# --
1082# -- Workers
1083# --
1084# --
1085# -----------------------------------------------------
1088@router.websocket("/{work_pool_name}/workers/connect")
1089async def worker_channel_connect(
1090 websocket: WebSocket,
1091 work_pool_name: str = Path(..., description="The work pool name"),
1092 db: PrefectDBInterface = Depends(provide_database_interface),
1093) -> None:
1094 websocket = await subscriptions.accept_prefect_socket(
1095 websocket,
1096 require_prefect_subprotocol=True,
1097 authentication_failed_reason=WorkerChannelCloseReason.AUTHENTICATION_FAILED.value,
1098 )
1099 if not websocket:
1100 return
1102 try:
1103 hello = await _receive_worker_hello(websocket)
1104 cleanup_queue: WorkerCleanupQueue | None = None
1105 if _worker_requested_cleanup_delivery(hello):
1106 try:
1107 cleanup_queue = get_worker_cleanup_queue()
1108 except Exception:
1109 logger.exception(
1110 "Worker cleanup delivery queue initialization failed; "
1111 "rejecting cleanup delivery capability"
1112 )
1114 async with worker_channel_utils.messaging.ephemeral_subscription(
1115 worker_channel_utils.WORKER_CHANNEL_SNAPSHOT_TOPIC,
1116 ) as consumer_kwargs:
1117 async with db.session_context(begin_transaction=True) as session:
1118 ready, work_pool_update_event = await _build_worker_ready_frame(
1119 session=session,
1120 work_pool_name=work_pool_name,
1121 hello=hello,
1122 cleanup_queue_available=cleanup_queue is not None,
1123 )
1125 if work_pool_update_event is not None:
1126 async with db.session_context() as session:
1127 work_pool = await models.workers.read_work_pool(
1128 session=session,
1129 work_pool_id=work_pool_update_event.work_pool_id,
1130 )
1131 assert work_pool is not None
1132 await models.workers.emit_work_pool_updated_event(
1133 session=session,
1134 work_pool=work_pool,
1135 changed_fields=work_pool_update_event.changed_fields,
1136 )
1137 await worker_channel_utils.publish_snapshot_invalidation(
1138 worker_channel_utils.WorkerChannelSnapshotInvalidation(
1139 work_pool_id=work_pool_update_event.work_pool_id,
1140 reason="work_pool_updated",
1141 )
1142 )
1144 connection = worker_channel_utils.WorkerChannelConnection(
1145 websocket=websocket,
1146 db=db,
1147 work_pool_name=work_pool_name,
1148 work_pool_id=ready.payload.initial_snapshot.work_pool.id,
1149 consumer_id=hello.payload.consumer_id,
1150 worker_name=hello.payload.worker_name,
1151 work_pool_updated=ready.payload.initial_snapshot.work_pool.updated,
1152 cleanup_queue=(
1153 cleanup_queue
1154 if CLEANUP_DELIVERY_CAPABILITY
1155 in ready.payload.accepted_capabilities
1156 else None
1157 ),
1158 cleanup_kinds=tuple(hello.payload.handled_cleanup_kinds),
1159 cleanup_work_queue_ids=tuple(
1160 work_queue.id for work_queue in ready.payload.resolved_work_queues
1161 ),
1162 max_cleanup_concurrency=ready.payload.effective_max_cleanup_concurrency,
1163 )
1164 await connection.run(ready, consumer_kwargs)
1166 except WorkerChannelSetupError as exc:
1167 logger.debug(
1168 "Worker channel setup failed: work_pool=%s reason=%s detail=%s",
1169 work_pool_name,
1170 exc.close_reason.value,
1171 exc.detail,
1172 )
1173 if worker_channel_utils.WORKER_CHANNEL_CONNECTIONS is not None:
1174 worker_channel_utils.WORKER_CHANNEL_CONNECTIONS.labels(
1175 event="setup_failed"
1176 ).inc()
1177 await worker_channel_utils.close_worker_channel(websocket, exc.close_reason)
1178 except HTTPException as exc:
1179 logger.debug(
1180 "Worker channel setup failed HTTP validation: work_pool=%s detail=%s",
1181 work_pool_name,
1182 exc.detail,
1183 )
1184 if worker_channel_utils.WORKER_CHANNEL_CONNECTIONS is not None:
1185 worker_channel_utils.WORKER_CHANNEL_CONNECTIONS.labels(
1186 event="setup_failed"
1187 ).inc()
1188 await worker_channel_utils.close_worker_channel(
1189 websocket, WorkerChannelCloseReason.PROTOCOL_ERROR
1190 )
1191 except subscriptions.NORMAL_DISCONNECT_EXCEPTIONS:
1192 return
1193 except Exception:
1194 logger.exception(
1195 "Worker channel setup failed due to a transient server error: work_pool=%s",
1196 work_pool_name,
1197 )
1198 if worker_channel_utils.WORKER_CHANNEL_CONNECTIONS is not None:
1199 worker_channel_utils.WORKER_CHANNEL_CONNECTIONS.labels(
1200 event="setup_failed"
1201 ).inc()
1202 await worker_channel_utils.close_worker_channel(
1203 websocket, WorkerChannelCloseReason.TRANSIENT_SERVER_ERROR
1204 )
1207@router.post(
1208 "/{work_pool_name}/workers/heartbeat",
1209 status_code=status.HTTP_204_NO_CONTENT,
1210)
1211async def worker_heartbeat(
1212 work_pool_name: str = Path(..., description="The work pool name"),
1213 name: str = Body(..., description="The worker process name", embed=True),
1214 heartbeat_interval_seconds: Optional[int] = Body(
1215 None, description="The worker's heartbeat interval in seconds", embed=True
1216 ),
1217 worker_lookups: WorkerLookups = Depends(WorkerLookups),
1218 db: PrefectDBInterface = Depends(provide_database_interface),
1219) -> None:
1220 async with db.session_context(begin_transaction=True) as session:
1221 work_pool = await models.workers.read_work_pool_by_name(
1222 session=session,
1223 work_pool_name=work_pool_name,
1224 )
1225 if not work_pool:
1226 raise HTTPException(
1227 status_code=status.HTTP_404_NOT_FOUND,
1228 detail=f'Work pool "{work_pool_name}" not found.',
1229 )
1231 await models.workers.record_worker_heartbeat(
1232 session=session,
1233 work_pool=work_pool,
1234 worker_name=name,
1235 heartbeat_interval_seconds=heartbeat_interval_seconds,
1236 emit_status_change=emit_work_pool_status_event,
1237 )
1240@router.post("/{work_pool_name}/workers/filter")
1241async def read_workers(
1242 work_pool_name: str = Path(..., description="The work pool name"),
1243 workers: Optional[schemas.filters.WorkerFilter] = None,
1244 limit: int = dependencies.LimitBody(),
1245 offset: int = Body(0, ge=0),
1246 worker_lookups: WorkerLookups = Depends(WorkerLookups),
1247 db: PrefectDBInterface = Depends(provide_database_interface),
1248) -> List[schemas.responses.WorkerResponse]:
1249 """
1250 Read all worker processes
1251 """
1252 async with db.session_context() as session:
1253 work_pool_id = await worker_lookups._get_work_pool_id_from_name(
1254 session=session, work_pool_name=work_pool_name
1255 )
1256 return await models.workers.read_workers(
1257 session=session,
1258 work_pool_id=work_pool_id,
1259 worker_filter=workers,
1260 limit=limit,
1261 offset=offset,
1262 )
1265@router.delete(
1266 "/{work_pool_name}/workers/{name}", status_code=status.HTTP_204_NO_CONTENT
1267)
1268async def delete_worker(
1269 work_pool_name: str = Path(..., description="The work pool name"),
1270 worker_name: str = Path(
1271 ..., description="The work pool's worker name", alias="name"
1272 ),
1273 worker_lookups: WorkerLookups = Depends(WorkerLookups),
1274 db: PrefectDBInterface = Depends(provide_database_interface),
1275) -> None:
1276 """
1277 Delete a work pool's worker
1278 """
1280 async with db.session_context(begin_transaction=True) as session:
1281 work_pool_id = await worker_lookups._get_work_pool_id_from_name(
1282 session=session, work_pool_name=work_pool_name
1283 )
1284 deleted = await models.workers.delete_worker(
1285 session=session, work_pool_id=work_pool_id, worker_name=worker_name
1286 )
1287 if not deleted:
1288 raise HTTPException(
1289 status_code=status.HTTP_404_NOT_FOUND, detail="Worker not found."
1290 )