prefect-client 3.0.0rc17__py3-none-any.whl → 3.0.0rc19__py3-none-any.whl

prefect/_internal/concurrency/services.py

@@ -39,6 +39,7 @@ class QueueService(abc.ABC, Generic[T]):
             daemon=True,
             name=f"{type(self).__name__}Thread",
         )
+        self._logger = logging.getLogger(f"{type(self).__name__}")
 
     def start(self):
         logger.debug("Starting service %r", self)
@@ -144,11 +145,24 @@ class QueueService(abc.ABC, Generic[T]):
             self._done_event.set()
 
     async def _main_loop(self):
+        last_log_time = 0
+        log_interval = 4  # log every 4 seconds
+
         while True:
             item: T = await self._queue_get_thread.submit(
                 create_call(self._queue.get)
             ).aresult()
 
+            if self._stopped:
+                current_time = asyncio.get_event_loop().time()
+                queue_size = self._queue.qsize()
+
+                if current_time - last_log_time >= log_interval and queue_size > 0:
+                    self._logger.warning(
+                        f"Still processing items: {queue_size} items remaining..."
+                    )
+                    last_log_time = current_time
+
             if item is None:
                 logger.debug("Exiting service %r", self)
                 self._queue.task_done()

prefect/_internal/schemas/bases.py

@@ -35,6 +35,7 @@ class PrefectBaseModel(BaseModel):
 
     model_config = ConfigDict(
         ser_json_timedelta="float",
+        defer_build=True,
         extra=(
             "ignore"
             if os.getenv("PREFECT_TEST_MODE", "0").lower() not in ["true", "1"]

prefect/blocks/core.py

@@ -24,9 +24,7 @@ from typing import (
 )
 from uuid import UUID, uuid4
 
-from griffe.dataclasses import Docstring
-from griffe.docstrings.dataclasses import DocstringSection, DocstringSectionKind
-from griffe.docstrings.parsers import Parser, parse
+from griffe import Docstring, DocstringSection, DocstringSectionKind, Parser, parse
 from packaging.version import InvalidVersion, Version
 from pydantic import (
     BaseModel,
@@ -130,7 +128,9 @@ def _is_subclass(cls, parent_cls) -> bool:
     Checks if a given class is a subclass of another class. Unlike issubclass,
     this will not throw an exception if cls is an instance instead of a type.
     """
-    return inspect.isclass(cls) and issubclass(cls, parent_cls)
+    # For python<=3.11 inspect.isclass() will return True for parametrized types (e.g. list[str])
+    # so we need to check for get_origin() to avoid TypeError for issubclass.
+    return inspect.isclass(cls) and not get_origin(cls) and issubclass(cls, parent_cls)
 
 
 def _collect_secret_fields(
@@ -138,12 +138,12 @@ def _collect_secret_fields(
 ) -> None:
     """
     Recursively collects all secret fields from a given type and adds them to the
-    secrets list, supporting nested Union / BaseModel fields. Also, note, this function
-    mutates the input secrets list, thus does not return anything.
+    secrets list, supporting nested Union / Dict / Tuple / List / BaseModel fields.
+    Also, note, this function mutates the input secrets list, thus does not return anything.
     """
-    if get_origin(type_) is Union:
-        for union_type in get_args(type_):
-            _collect_secret_fields(name, union_type, secrets)
+    if get_origin(type_) in (Union, dict, list, tuple):
+        for nested_type in get_args(type_):
+            _collect_secret_fields(name, nested_type, secrets)
         return
     elif _is_subclass(type_, BaseModel):
         for field_name, field in type_.model_fields.items():
@@ -232,21 +232,25 @@ def schema_extra(schema: Dict[str, Any], model: Type["Block"]):
 
     # create block schema references
     refs = schema["block_schema_references"] = {}
+
+    def collect_block_schema_references(field_name: str, annotation: type) -> None:
+        """Walk through the annotation and collect block schemas for any nested blocks."""
+        if Block.is_block_class(annotation):
+            if isinstance(refs.get(field_name), list):
+                refs[field_name].append(annotation._to_block_schema_reference_dict())
+            elif isinstance(refs.get(field_name), dict):
+                refs[field_name] = [
+                    refs[field_name],
+                    annotation._to_block_schema_reference_dict(),
+                ]
+            else:
+                refs[field_name] = annotation._to_block_schema_reference_dict()
+        if get_origin(annotation) in (Union, list, tuple, dict):
+            for type_ in get_args(annotation):
+                collect_block_schema_references(field_name, type_)
+
     for name, field in model.model_fields.items():
-        if Block.is_block_class(field.annotation):
-            refs[name] = field.annotation._to_block_schema_reference_dict()
-        if get_origin(field.annotation) in [Union, list]:
-            for type_ in get_args(field.annotation):
-                if Block.is_block_class(type_):
-                    if isinstance(refs.get(name), list):
-                        refs[name].append(type_._to_block_schema_reference_dict())
-                    elif isinstance(refs.get(name), dict):
-                        refs[name] = [
-                            refs[name],
-                            type_._to_block_schema_reference_dict(),
-                        ]
-                    else:
-                        refs[name] = type_._to_block_schema_reference_dict()
+        collect_block_schema_references(name, field.annotation)
 
 
 @register_base_type
@@ -1067,13 +1071,16 @@ class Block(BaseModel, ABC):
                 "subclass and not on a Block interface class directly."
             )
 
+        async def register_blocks_in_annotation(annotation: type) -> None:
+            """Walk through the annotation and register any nested blocks."""
+            if Block.is_block_class(annotation):
+                await annotation.register_type_and_schema(client=client)
+            elif get_origin(annotation) in (Union, tuple, list, dict):
+                for inner_annotation in get_args(annotation):
+                    await register_blocks_in_annotation(inner_annotation)
+
         for field in cls.model_fields.values():
-            if Block.is_block_class(field.annotation):
-                await field.annotation.register_type_and_schema(client=client)
-            if get_origin(field.annotation) is Union:
-                for annotation in get_args(field.annotation):
-                    if Block.is_block_class(annotation):
-                        await annotation.register_type_and_schema(client=client)
+            await register_blocks_in_annotation(field.annotation)
 
         try:
             block_type = await client.read_block_type_by_slug(

prefect/client/orchestration.py

@@ -939,6 +939,57 @@ class PrefectClient:
             else:
                 raise
 
+    async def increment_v1_concurrency_slots(
+        self,
+        names: List[str],
+        task_run_id: UUID,
+    ) -> httpx.Response:
+        """
+        Increment concurrency limit slots for the specified limits.
+
+        Args:
+            names (List[str]): A list of limit names for which to increment limits.
+            task_run_id (UUID): The task run ID incrementing the limits.
+        """
+        data = {
+            "names": names,
+            "task_run_id": str(task_run_id),
+        }
+
+        return await self._client.post(
+            "/concurrency_limits/increment",
+            json=data,
+        )
+
+    async def decrement_v1_concurrency_slots(
+        self,
+        names: List[str],
+        task_run_id: UUID,
+        occupancy_seconds: float,
+    ) -> httpx.Response:
+        """
+        Decrement concurrency limit slots for the specified limits.
+
+        Args:
+            names (List[str]): A list of limit names to decrement.
+            task_run_id (UUID): The task run ID that incremented the limits.
+            occupancy_seconds (float): The duration in seconds that the limits
+                were held.
+
+        Returns:
+            httpx.Response: The HTTP response from the server.
+        """
+        data = {
+            "names": names,
+            "task_run_id": str(task_run_id),
+            "occupancy_seconds": occupancy_seconds,
+        }
+
+        return await self._client.post(
+            "/concurrency_limits/decrement",
+            json=data,
+        )
+
     async def create_work_queue(
         self,
         name: str,
@@ -1599,6 +1650,7 @@ class PrefectClient:
         name: str,
         version: Optional[str] = None,
         schedules: Optional[List[DeploymentScheduleCreate]] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[Dict[str, Any]] = None,
         description: Optional[str] = None,
         work_queue_name: Optional[str] = None,
@@ -1656,6 +1708,7 @@ class PrefectClient:
             parameter_openapi_schema=parameter_openapi_schema,
             paused=paused,
             schedules=schedules or [],
+            concurrency_limit=concurrency_limit,
             pull_steps=pull_steps,
             enforce_parameter_schema=enforce_parameter_schema,
         )
@@ -2612,6 +2665,7 @@ class PrefectClient:
     async def create_work_pool(
         self,
         work_pool: WorkPoolCreate,
+        overwrite: bool = False,
     ) -> WorkPool:
         """
         Creates a work pool with the provided configuration.
@@ -2629,7 +2683,24 @@ class PrefectClient:
             )
         except httpx.HTTPStatusError as e:
             if e.response.status_code == status.HTTP_409_CONFLICT:
-                raise prefect.exceptions.ObjectAlreadyExists(http_exc=e) from e
+                if overwrite:
+                    existing_work_pool = await self.read_work_pool(
+                        work_pool_name=work_pool.name
+                    )
+                    if existing_work_pool.type != work_pool.type:
+                        warnings.warn(
+                            "Overwriting work pool type is not supported. Ignoring provided type.",
+                            category=UserWarning,
+                        )
+                    await self.update_work_pool(
+                        work_pool_name=work_pool.name,
+                        work_pool=WorkPoolUpdate.model_validate(
+                            work_pool.model_dump(exclude={"name", "type"})
+                        ),
+                    )
+                    response = await self._client.get(f"/work_pools/{work_pool.name}")
+                else:
+                    raise prefect.exceptions.ObjectAlreadyExists(http_exc=e) from e
             else:
                 raise
 
@@ -3156,7 +3227,7 @@ class PrefectClient:
         return pydantic.TypeAdapter(List[Automation]).validate_python(response.json())
 
     async def find_automation(
-        self, id_or_name: Union[str, UUID], exit_if_not_found: bool = True
+        self, id_or_name: Union[str, UUID]
     ) -> Optional[Automation]:
         if isinstance(id_or_name, str):
             try:
@@ -4096,3 +4167,27 @@ class SyncPrefectClient:
                 "occupancy_seconds": occupancy_seconds,
             },
         )
+
+    def decrement_v1_concurrency_slots(
+        self, names: List[str], occupancy_seconds: float, task_run_id: UUID
+    ) -> httpx.Response:
+        """
+        Release the specified concurrency limits.
+
+        Args:
+            names (List[str]): A list of limit names to decrement.
+            occupancy_seconds (float): The duration in seconds that the slots
+                were held.
+            task_run_id (UUID): The task run ID that incremented the limits.
+
+        Returns:
+            httpx.Response: The HTTP response from the server.
+        """
+        return self._client.post(
+            "/concurrency_limits/decrement",
+            json={
+                "names": names,
+                "occupancy_seconds": occupancy_seconds,
+                "task_run_id": str(task_run_id),
+            },
+        )

prefect/client/schemas/actions.py

@@ -157,6 +157,10 @@ class DeploymentCreate(ActionBaseModel):
         default_factory=list,
         description="A list of schedules for the deployment.",
     )
+    concurrency_limit: Optional[int] = Field(
+        default=None,
+        description="The concurrency limit for the deployment.",
+    )
     enforce_parameter_schema: Optional[bool] = Field(
         default=None,
         description=(
@@ -229,6 +233,10 @@ class DeploymentUpdate(ActionBaseModel):
         default=None,
         description="A list of schedules for the deployment.",
     )
+    concurrency_limit: Optional[int] = Field(
+        default=None,
+        description="The concurrency limit for the deployment.",
+    )
     tags: List[str] = Field(default_factory=list)
     work_queue_name: Optional[str] = Field(None)
     work_pool_name: Optional[str] = Field(
@@ -607,11 +615,11 @@ class WorkQueueCreate(ActionBaseModel):
         default=False,
         description="Whether the work queue is paused.",
     )
-    concurrency_limit: Optional[int] = Field(
+    concurrency_limit: Optional[NonNegativeInteger] = Field(
         default=None,
         description="A concurrency limit for the work queue.",
     )
-    priority: Optional[int] = Field(
+    priority: Optional[PositiveInteger] = Field(
         default=None,
         description=(
             "The queue's priority. Lower values are higher priority (1 is the highest)."
@@ -635,8 +643,10 @@ class WorkQueueUpdate(ActionBaseModel):
     is_paused: bool = Field(
         default=False, description="Whether or not the work queue is paused."
     )
-    concurrency_limit: Optional[int] = Field(None)
-    priority: Optional[int] = Field(None)
+    concurrency_limit: Optional[NonNegativeInteger] = Field(None)
+    priority: Optional[PositiveInteger] = Field(
+        None, description="The queue's priority."
+    )
     last_polled: Optional[DateTime] = Field(None)
 
     # DEPRECATED

prefect/client/schemas/filters.py

@@ -505,6 +505,23 @@ class DeploymentFilterTags(PrefectBaseModel, OperatorMixin):
     )
 
 
+class DeploymentFilterConcurrencyLimit(PrefectBaseModel):
+    """Filter by `Deployment.concurrency_limit`."""
+
+    ge_: Optional[int] = Field(
+        default=None,
+        description="Only include deployments with a concurrency limit greater than or equal to this value",
+    )
+    le_: Optional[int] = Field(
+        default=None,
+        description="Only include deployments with a concurrency limit less than or equal to this value",
+    )
+    is_null_: Optional[bool] = Field(
+        default=None,
+        description="If true, only include deployments without a concurrency limit",
+    )
+
+
 class DeploymentFilter(PrefectBaseModel, OperatorMixin):
     """Filter for deployments. Only deployments matching all criteria will be returned."""
 
@@ -520,6 +537,9 @@ class DeploymentFilter(PrefectBaseModel, OperatorMixin):
     work_queue_name: Optional[DeploymentFilterWorkQueueName] = Field(
         default=None, description="Filter criteria for `Deployment.work_queue_name`"
     )
+    concurrency_limit: Optional[DeploymentFilterConcurrencyLimit] = Field(
+        default=None, description="Filter criteria for `Deployment.concurrency_limit`"
+    )
 
 
 class LogFilterName(PrefectBaseModel):

prefect/client/schemas/objects.py

@@ -996,6 +996,9 @@ class Deployment(ObjectBaseModel):
     paused: bool = Field(
         default=False, description="Whether or not the deployment is paused."
     )
+    concurrency_limit: Optional[int] = Field(
+        default=None, description="The concurrency limit for the deployment."
+    )
     schedules: List[DeploymentSchedule] = Field(
         default_factory=list, description="A list of schedules for the deployment."
     )

prefect/client/schemas/responses.py

@@ -313,6 +313,9 @@ class DeploymentResponse(ObjectBaseModel):
     flow_id: UUID = Field(
         default=..., description="The flow id associated with the deployment."
     )
+    concurrency_limit: Optional[int] = Field(
+        default=None, description="The concurrency limit for the deployment."
+    )
     paused: bool = Field(
         default=False, description="Whether or not the deployment is paused."
     )

prefect/client/schemas/sorting.py

@@ -59,6 +59,8 @@ class DeploymentSort(AutoEnum):
     UPDATED_DESC = AutoEnum.auto()
     NAME_ASC = AutoEnum.auto()
     NAME_DESC = AutoEnum.auto()
+    CONCURRENCY_LIMIT_ASC = AutoEnum.auto()
+    CONCURRENCY_LIMIT_DESC = AutoEnum.auto()
 
 
 class ArtifactSort(AutoEnum):

prefect/concurrency/v1/asyncio.py

@@ -0,0 +1,143 @@
+import asyncio
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator, List, Optional, Union, cast
+from uuid import UUID
+
+import anyio
+import httpx
+import pendulum
+
+from ...client.schemas.responses import MinimalConcurrencyLimitResponse
+
+try:
+    from pendulum import Interval
+except ImportError:
+    # pendulum < 3
+    from pendulum.period import Period as Interval  # type: ignore
+
+from prefect.client.orchestration import get_client
+
+from .context import ConcurrencyContext
+from .events import (
+    _emit_concurrency_acquisition_events,
+    _emit_concurrency_release_events,
+)
+from .services import ConcurrencySlotAcquisitionService
+
+
+class ConcurrencySlotAcquisitionError(Exception):
+    """Raised when an unhandlable occurs while acquiring concurrency slots."""
+
+
+class AcquireConcurrencySlotTimeoutError(TimeoutError):
+    """Raised when acquiring a concurrency slot times out."""
+
+
+@asynccontextmanager
+async def concurrency(
+    names: Union[str, List[str]],
+    task_run_id: UUID,
+    timeout_seconds: Optional[float] = None,
+) -> AsyncGenerator[None, None]:
+    """A context manager that acquires and releases concurrency slots from the
+    given concurrency limits.
+
+    Args:
+        names: The names of the concurrency limits to acquire slots from.
+        task_run_id: The name of the task_run_id that is incrementing the slots.
+        timeout_seconds: The number of seconds to wait for the slots to be acquired before
+            raising a `TimeoutError`. A timeout of `None` will wait indefinitely.
+
+    Raises:
+        TimeoutError: If the slots are not acquired within the given timeout.
+
+    Example:
+    A simple example of using the async `concurrency` context manager:
+    ```python
+    from prefect.concurrency.v1.asyncio import concurrency
+
+    async def resource_heavy():
+        async with concurrency("test", task_run_id):
+            print("Resource heavy task")
+
+    async def main():
+        await resource_heavy()
+    ```
+    """
+    if not names:
+        yield
+        return
+
+    names_normalized: List[str] = names if isinstance(names, list) else [names]
+
+    limits = await _acquire_concurrency_slots(
+        names_normalized,
+        task_run_id=task_run_id,
+        timeout_seconds=timeout_seconds,
+    )
+    acquisition_time = pendulum.now("UTC")
+    emitted_events = _emit_concurrency_acquisition_events(limits, task_run_id)
+
+    try:
+        yield
+    finally:
+        occupancy_period = cast(Interval, (pendulum.now("UTC") - acquisition_time))
+        try:
+            await _release_concurrency_slots(
+                names_normalized, task_run_id, occupancy_period.total_seconds()
+            )
+        except anyio.get_cancelled_exc_class():
+            # The task was cancelled before it could release the slots. Add the
+            # slots to the cleanup list so they can be released when the
+            # concurrency context is exited.
+            if ctx := ConcurrencyContext.get():
+                ctx.cleanup_slots.append(
+                    (names_normalized, occupancy_period.total_seconds(), task_run_id)
+                )
+
+        _emit_concurrency_release_events(limits, emitted_events, task_run_id)
+
+
+async def _acquire_concurrency_slots(
+    names: List[str],
+    task_run_id: UUID,
+    timeout_seconds: Optional[float] = None,
+) -> List[MinimalConcurrencyLimitResponse]:
+    service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
+    future = service.send((task_run_id, timeout_seconds))
+    response_or_exception = await asyncio.wrap_future(future)
+
+    if isinstance(response_or_exception, Exception):
+        if isinstance(response_or_exception, TimeoutError):
+            raise AcquireConcurrencySlotTimeoutError(
+                f"Attempt to acquire concurrency limits timed out after {timeout_seconds} second(s)"
+            ) from response_or_exception
+
+        raise ConcurrencySlotAcquisitionError(
+            f"Unable to acquire concurrency limits {names!r}"
+        ) from response_or_exception
+
+    return _response_to_concurrency_limit_response(response_or_exception)
+
+
+async def _release_concurrency_slots(
+    names: List[str],
+    task_run_id: UUID,
+    occupancy_seconds: float,
+) -> List[MinimalConcurrencyLimitResponse]:
+    async with get_client() as client:
+        response = await client.decrement_v1_concurrency_slots(
+            names=names,
+            task_run_id=task_run_id,
+            occupancy_seconds=occupancy_seconds,
+        )
+        return _response_to_concurrency_limit_response(response)
+
+
+def _response_to_concurrency_limit_response(
+    response: httpx.Response,
+) -> List[MinimalConcurrencyLimitResponse]:
+    data = response.json() or []
+    return [
+        MinimalConcurrencyLimitResponse.model_validate(limit) for limit in data if data
+    ]

prefect/concurrency/v1/context.py

@@ -0,0 +1,27 @@
+from contextvars import ContextVar
+from typing import List, Tuple
+from uuid import UUID
+
+from prefect.client.orchestration import get_client
+from prefect.context import ContextModel, Field
+
+
+class ConcurrencyContext(ContextModel):
+    __var__: ContextVar = ContextVar("concurrency_v1")
+
+    # Track the limits that have been acquired but were not able to be released
+    # due to cancellation or some other error. These limits are released when
+    # the context manager exits.
+    cleanup_slots: List[Tuple[List[str], float, UUID]] = Field(default_factory=list)
+
+    def __exit__(self, *exc_info):
+        if self.cleanup_slots:
+            with get_client(sync_client=True) as client:
+                for names, occupancy_seconds, task_run_id in self.cleanup_slots:
+                    client.decrement_v1_concurrency_slots(
+                        names=names,
+                        occupancy_seconds=occupancy_seconds,
+                        task_run_id=task_run_id,
+                    )
+
+        return super().__exit__(*exc_info)

prefect/concurrency/v1/events.py

@@ -0,0 +1,61 @@
+from typing import Dict, List, Literal, Optional, Union
+from uuid import UUID
+
+from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
+from prefect.events import Event, RelatedResource, emit_event
+
+
+def _emit_concurrency_event(
+    phase: Union[Literal["acquired"], Literal["released"]],
+    primary_limit: MinimalConcurrencyLimitResponse,
+    related_limits: List[MinimalConcurrencyLimitResponse],
+    task_run_id: UUID,
+    follows: Union[Event, None] = None,
+) -> Union[Event, None]:
+    resource: Dict[str, str] = {
+        "prefect.resource.id": f"prefect.concurrency-limit.v1.{primary_limit.id}",
+        "prefect.resource.name": primary_limit.name,
+        "limit": str(primary_limit.limit),
+        "task_run_id": str(task_run_id),
+    }
+
+    related = [
+        RelatedResource.model_validate(
+            {
+                "prefect.resource.id": f"prefect.concurrency-limit.v1.{limit.id}",
+                "prefect.resource.role": "concurrency-limit",
+            }
+        )
+        for limit in related_limits
+        if limit.id != primary_limit.id
+    ]
+
+    return emit_event(
+        f"prefect.concurrency-limit.v1.{phase}",
+        resource=resource,
+        related=related,
+        follows=follows,
+    )
+
+
+def _emit_concurrency_acquisition_events(
+    limits: List[MinimalConcurrencyLimitResponse],
+    task_run_id: UUID,
+) -> Dict[UUID, Optional[Event]]:
+    events = {}
+    for limit in limits:
+        event = _emit_concurrency_event("acquired", limit, limits, task_run_id)
+        events[limit.id] = event
+
+    return events
+
+
+def _emit_concurrency_release_events(
+    limits: List[MinimalConcurrencyLimitResponse],
+    events: Dict[UUID, Optional[Event]],
+    task_run_id: UUID,
+) -> None:
+    for limit in limits:
+        _emit_concurrency_event(
+            "released", limit, limits, task_run_id, events[limit.id]
+        )