prefect-client 3.0.0rc17__tar.gz → 3.0.0rc19__tar.gz

{prefect-client-3.0.0rc17/src/prefect_client.egg-info → prefect-client-3.0.0rc19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.0.0rc17
+Version: 3.0.0rc19
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.

{prefect-client-3.0.0rc17 → prefect-client-3.0.0rc19}/requirements-client.txt

@@ -8,11 +8,10 @@ exceptiongroup >= 1.0.0
 fastapi >= 0.111.0, < 1.0.0
 fsspec >= 2022.5.0
 graphviz >= 0.20.1
-griffe >= 0.20.0, <0.48.0
+griffe >= 0.49.0, <2.0.0
 httpcore >=1.0.5, < 2.0.0
 httpx[http2] >= 0.23, != 0.23.2
 importlib_metadata >= 4.4; python_version < '3.10'
-importlib-resources >= 6.1.3, < 6.2.0
 jsonpatch >= 1.32, < 2.0
 jsonschema >= 4.0.0, < 5.0.0
 orjson >= 3.7, < 4.0
@@ -35,4 +34,4 @@ toml >= 0.10.0
 typing_extensions >= 4.5.0, < 5.0.0
 ujson >= 5.8.0, < 6.0.0
 uvicorn >=0.14.0, !=0.29.0
-websockets >= 10.4, < 13.0
+websockets >= 10.4, < 14.0

{prefect-client-3.0.0rc17 → prefect-client-3.0.0rc19}/requirements.txt

@@ -1,18 +1,17 @@
 -r requirements-client.txt
 
-aiosqlite >= 0.17.0
+aiosqlite >= 0.17.0, < 1.0.0
 alembic >= 1.7.5, < 2.0.0
 apprise >= 1.1.0, < 2.0.0
-asyncpg >= 0.23
+asyncpg >= 0.23, < 1.0.0
 click >= 8.0, < 8.2
 cryptography >= 36.0.1
 dateparser >= 1.1.1, < 2.0.0
-docker >= 4.0
+docker >= 4.0, < 8.0
 graphviz >= 0.20.1
-griffe >= 0.20.0
 jinja2 >= 3.0.0, < 4.0.0
 jinja2-humanize-extension >= 0.4.0
-humanize >= 4.9.0
+humanize >= 4.9.0, < 5.0.0
 pytz >= 2021.1, < 2025
 readchar >= 4.0.0, < 5.0.0
 sqlalchemy[asyncio] >= 2.0, < 3.0.0

{prefect-client-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-client-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-client-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-3.0.0rc17 → prefect-client-3.0.0rc19}/src/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-client-3.0.0rc19/src/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-client-3.0.0rc19/src/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)