prefect-client 2.20.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/client/schemas/responses.py

@@ -2,20 +2,13 @@ import datetime
 from typing import Any, Dict, List, Optional, TypeVar, Union
 from uuid import UUID
 
-from prefect._internal.compatibility.deprecated import DeprecatedInfraOverridesField
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field
-else:
-    from pydantic import Field
-
+from pydantic import ConfigDict, Field
+from pydantic_extra_types.pendulum_dt import DateTime
 from typing_extensions import Literal
 
 import prefect.client.schemas.objects as objects
 from prefect._internal.schemas.bases import ObjectBaseModel, PrefectBaseModel
-from prefect._internal.schemas.fields import CreatedBy, DateTimeTZ, UpdatedBy
-from prefect.client.schemas.schedules import SCHEDULE_TYPES
+from prefect._internal.schemas.fields import CreatedBy, UpdatedBy
 from prefect.utilities.collections import AutoEnum
 from prefect.utilities.names import generate_slug
 
@@ -120,10 +113,10 @@ class HistoryResponseState(PrefectBaseModel):
 class HistoryResponse(PrefectBaseModel):
     """Represents a history of aggregation states over an interval"""
 
-    interval_start: DateTimeTZ = Field(
+    interval_start: DateTime = Field(
         default=..., description="The start date of the interval."
     )
-    interval_end: DateTimeTZ = Field(
+    interval_end: DateTime = Field(
         default=..., description="The end date of the interval."
     )
     states: List[HistoryResponseState] = Field(
@@ -147,8 +140,7 @@ class OrchestrationResult(PrefectBaseModel):
 
 
 class WorkerFlowRunResponse(PrefectBaseModel):
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     work_pool_id: UUID
     work_queue_id: UUID
@@ -219,18 +211,18 @@ class FlowRunResponse(ObjectBaseModel):
     run_count: int = Field(
         default=0, description="The number of times the flow run was executed."
     )
-    expected_start_time: Optional[DateTimeTZ] = Field(
+    expected_start_time: Optional[DateTime] = Field(
         default=None,
         description="The flow run's expected start time.",
     )
-    next_scheduled_start_time: Optional[DateTimeTZ] = Field(
+    next_scheduled_start_time: Optional[DateTime] = Field(
         default=None,
         description="The next time the flow run is scheduled to start.",
     )
-    start_time: Optional[DateTimeTZ] = Field(
+    start_time: Optional[DateTime] = Field(
         default=None, description="The actual start time."
     )
-    end_time: Optional[DateTimeTZ] = Field(
+    end_time: Optional[DateTime] = Field(
         default=None, description="The actual end time."
     )
     total_run_time: datetime.timedelta = Field(
@@ -279,7 +271,7 @@ class FlowRunResponse(ObjectBaseModel):
     state: Optional[objects.State] = Field(
         default=None,
         description="The state of the flow run.",
-        examples=[objects.State(type=objects.StateType.COMPLETED)],
+        examples=["objects.State(type=objects.StateType.COMPLETED)"],
     )
     job_variables: Optional[dict] = Field(
         default=None, description="Job variables for the flow run."
@@ -304,13 +296,13 @@ class FlowRunResponse(ObjectBaseModel):
         """
         if isinstance(other, objects.FlowRun):
             exclude_fields = {"estimated_run_time", "estimated_start_time_delta"}
-            return self.dict(exclude=exclude_fields) == other.dict(
+            return self.model_dump(exclude=exclude_fields) == other.model_dump(
                 exclude=exclude_fields
             )
         return super().__eq__(other)
 
 
-class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
+class DeploymentResponse(ObjectBaseModel):
     name: str = Field(default=..., description="The name of the deployment.")
     version: Optional[str] = Field(
         default=None, description="An optional version for the deployment."
@@ -321,11 +313,8 @@ class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
     flow_id: UUID = Field(
         default=..., description="The flow id associated with the deployment."
     )
-    schedule: Optional[SCHEDULE_TYPES] = Field(
-        default=None, description="A schedule for the deployment."
-    )
-    is_schedule_active: bool = Field(
-        default=True, description="Whether or not the deployment schedule is active."
+    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."
@@ -357,7 +346,7 @@ class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
             " be scheduled."
         ),
     )
-    last_polled: Optional[DateTimeTZ] = Field(
+    last_polled: Optional[DateTime] = Field(
         default=None,
         description="The last time the deployment was polled for status updates.",
     )
@@ -378,12 +367,6 @@ class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
             "The path to the entrypoint for the workflow, relative to the `path`."
         ),
     )
-    manifest_path: Optional[str] = Field(
-        default=None,
-        description=(
-            "The path to the flow's manifest file, relative to the chosen storage."
-        ),
-    )
     storage_document_id: Optional[UUID] = Field(
         default=None,
         description="The block document defining storage used for this flow.",
@@ -407,7 +390,7 @@ class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
         ),
     )
     enforce_parameter_schema: bool = Field(
-        default=False,
+        default=True,
         description=(
             "Whether or not the deployment should enforce the parameter schema."
         ),
@@ -423,8 +406,7 @@ class DeploymentResponse(DeprecatedInfraOverridesField, ObjectBaseModel):
 
 
 class MinimalConcurrencyLimitResponse(PrefectBaseModel):
-    class Config:
-        extra = "ignore"  # 2024/4/1
+    model_config = ConfigDict(extra="ignore")
 
     id: UUID
     name: str

prefect/client/schemas/schedules.py

@@ -3,28 +3,22 @@ Schedule schemas
 """
 
 import datetime
-from typing import Optional, Union
+from typing import Annotated, Any, Optional, Union
 
 import dateutil
 import dateutil.rrule
 import pendulum
+from pydantic import AfterValidator, ConfigDict, Field, field_validator, model_validator
+from pydantic_extra_types.pendulum_dt import DateTime
+from typing_extensions import TypeAlias, TypeGuard
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
 from prefect._internal.schemas.bases import PrefectBaseModel
-from prefect._internal.schemas.fields import DateTimeTZ
 from prefect._internal.schemas.validators import (
     default_anchor_date,
     default_timezone,
     validate_cron_string,
     validate_rrule_string,
-    validate_rrule_timezone,
 )
-from prefect.types import PositiveDuration
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import Field, validator
-else:
-    from pydantic import Field, validator
 
 MAX_ITERATIONS = 1000
 # approx. 1 years worth of RDATEs + buffer
@@ -55,26 +49,24 @@ class IntervalSchedule(PrefectBaseModel):
 
     Args:
         interval (datetime.timedelta): an interval to schedule on
-        anchor_date (DateTimeTZ, optional): an anchor date to schedule increments against;
+        anchor_date (DateTime, optional): an anchor date to schedule increments against;
             if not provided, the current timestamp will be used
         timezone (str, optional): a valid timezone string
     """
 
-    class Config:
-        extra = "forbid"
-        exclude_none = True
+    model_config = ConfigDict(extra="forbid", exclude_none=True)
 
-    interval: PositiveDuration
-    anchor_date: Optional[DateTimeTZ] = None
+    interval: datetime.timedelta = Field(gt=datetime.timedelta(0))
+    anchor_date: Annotated[DateTime, AfterValidator(default_anchor_date)] = Field(
+        default_factory=lambda: pendulum.now("UTC"),
+        examples=["2020-01-01T00:00:00Z"],
+    )
     timezone: Optional[str] = Field(default=None, examples=["America/New_York"])
 
-    @validator("anchor_date", always=True)
-    def validate_anchor_date(cls, v):
-        return default_anchor_date(v)
-
-    @validator("timezone", always=True)
-    def validate_default_timezone(cls, v, values):
-        return default_timezone(v, values=values)
+    @model_validator(mode="after")
+    def validate_timezone(self):
+        self.timezone = default_timezone(self.timezone, self.model_dump())
+        return self
 
 
 class CronSchedule(PrefectBaseModel):
@@ -102,8 +94,7 @@ class CronSchedule(PrefectBaseModel):
 
     """
 
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
     cron: str = Field(default=..., examples=["0 0 * * *"])
     timezone: Optional[str] = Field(default=None, examples=["America/New_York"])
@@ -114,11 +105,13 @@ class CronSchedule(PrefectBaseModel):
         ),
     )
 
-    @validator("timezone")
+    @field_validator("timezone")
+    @classmethod
     def valid_timezone(cls, v):
         return default_timezone(v)
 
-    @validator("cron")
+    @field_validator("cron")
+    @classmethod
     def valid_cron_string(cls, v):
         return validate_cron_string(v)
 
@@ -146,13 +139,15 @@ class RRuleSchedule(PrefectBaseModel):
         timezone (str, optional): a valid timezone string
     """
 
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
     rrule: str
-    timezone: Optional[str] = Field(default=None, examples=["America/New_York"])
+    timezone: Optional[str] = Field(
+        default="UTC", examples=["America/New_York"], validate_default=True
+    )
 
-    @validator("rrule")
+    @field_validator("rrule")
+    @classmethod
     def validate_rrule_str(cls, v):
         return validate_rrule_string(v)
 
@@ -259,17 +254,39 @@ class RRuleSchedule(PrefectBaseModel):
 
             return rrule
 
-    @validator("timezone", always=True)
+    @field_validator("timezone")
     def valid_timezone(cls, v):
-        return validate_rrule_timezone(v)
+        """
+        Validate that the provided timezone is a valid IANA timezone.
+
+        Unfortunately this list is slightly different from the list of valid
+        timezones in pendulum that we use for cron and interval timezone validation.
+        """
+        from prefect._internal.pytz import HAS_PYTZ
+
+        if HAS_PYTZ:
+            import pytz
+        else:
+            from prefect._internal import pytz
+
+        if v and v not in pytz.all_timezones_set:
+            raise ValueError(f'Invalid timezone: "{v}"')
+        elif v is None:
+            return "UTC"
+        return v
 
 
 class NoSchedule(PrefectBaseModel):
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
+
+
+SCHEDULE_TYPES: TypeAlias = Union[
+    IntervalSchedule, CronSchedule, RRuleSchedule, NoSchedule
+]
 
 
-SCHEDULE_TYPES = Union[IntervalSchedule, CronSchedule, RRuleSchedule, NoSchedule]
+def is_schedule_type(obj: Any) -> TypeGuard[SCHEDULE_TYPES]:
+    return isinstance(obj, (IntervalSchedule, CronSchedule, RRuleSchedule, NoSchedule))
 
 
 def construct_schedule(
@@ -308,6 +325,8 @@ def construct_schedule(
     if interval:
         if isinstance(interval, (int, float)):
             interval = datetime.timedelta(seconds=interval)
+        if not anchor_date:
+            anchor_date = DateTime.now()
         schedule = IntervalSchedule(
             interval=interval, anchor_date=anchor_date, timezone=timezone
         )

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/subscriptions.py

@@ -1,15 +1,15 @@
 import asyncio
-from typing import Any, Dict, Generic, List, Optional, Type, TypeVar
+from typing import Any, Dict, Generic, Iterable, Optional, Type, TypeVar
 
 import orjson
 import websockets
 import websockets.exceptions
-from prefect._vendor.starlette.status import WS_1008_POLICY_VIOLATION
+from starlette.status import WS_1008_POLICY_VIOLATION
 from typing_extensions import Self
 
 from prefect._internal.schemas.bases import IDBaseModel
 from prefect.logging import get_logger
-from prefect.settings import PREFECT_API_KEY, PREFECT_API_URL
+from prefect.settings import PREFECT_API_KEY
 
 logger = get_logger(__name__)
 
@@ -21,15 +21,16 @@ class Subscription(Generic[S]):
         self,
         model: Type[S],
         path: str,
-        keys: List[str],
+        keys: Iterable[str],
         client_id: Optional[str] = None,
+        base_url: Optional[str] = None,
     ):
         self.model = model
         self.client_id = client_id
-        base_url = PREFECT_API_URL.value().replace("http", "ws", 1)
+        base_url = base_url.replace("http", "ws", 1)
         self.subscription_url = f"{base_url}{path}"
 
-        self.keys = keys
+        self.keys = list(keys)
 
         self._connect = websockets.connect(
             self.subscription_url,
@@ -48,7 +49,7 @@ class Subscription(Generic[S]):
 
                 await self._websocket.send(orjson.dumps({"type": "ack"}).decode())
 
-                return self.model.parse_raw(message)
+                return self.model.model_validate_json(message)
             except (
                 ConnectionRefusedError,
                 websockets.exceptions.ConnectionClosedError,

prefect/client/types/flexible_schedule_list.py

@@ -0,0 +1,11 @@
+from typing import TYPE_CHECKING, Any, Sequence, Union
+
+from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    from prefect.client.schemas.actions import DeploymentScheduleCreate
+    from prefect.client.schemas.schedules import SCHEDULE_TYPES
+
+FlexibleScheduleList: TypeAlias = Sequence[
+    Union[DeploymentScheduleCreate, dict[str, Any], "SCHEDULE_TYPES"]
+]

prefect/client/utilities.py

@@ -42,12 +42,15 @@ def get_or_create_client(
     if client is not None:
         return client, True
     from prefect._internal.concurrency.event_loop import get_running_loop
-    from prefect.context import FlowRunContext, TaskRunContext
+    from prefect.context import AsyncClientContext, FlowRunContext, TaskRunContext
 
+    async_client_context = AsyncClientContext.get()
     flow_run_context = FlowRunContext.get()
     task_run_context = TaskRunContext.get()
 
-    if (
+    if async_client_context and async_client_context.client._loop == get_running_loop():
+        return async_client_context.client, True
+    elif (
         flow_run_context
         and getattr(flow_run_context.client, "_loop", None) == get_running_loop()
     ):
@@ -75,10 +78,10 @@ def client_injector(
 
 
 def inject_client(
-    fn: Callable[P, Coroutine[Any, Any, Any]],
-) -> Callable[P, Coroutine[Any, Any, Any]]:
+    fn: Callable[P, Coroutine[Any, Any, R]],
+) -> Callable[P, Coroutine[Any, Any, R]]:
     """
-    Simple helper to provide a context managed client to a asynchronous function.
+    Simple helper to provide a context managed client to an asynchronous function.
 
     The decorated function _must_ take a `client` kwarg and if a client is passed when
     called it will be used instead of creating a new one, but it will not be context
@@ -86,7 +89,7 @@ def inject_client(
     """
 
     @wraps(fn)
-    async def with_injected_client(*args: P.args, **kwargs: P.kwargs) -> Any:
+    async def with_injected_client(*args: P.args, **kwargs: P.kwargs) -> R:
         client = cast(Optional["PrefectClient"], kwargs.pop("client", None))
         client, inferred = get_or_create_client(client)
         if not inferred:

prefect/concurrency/asyncio.py

@@ -1,7 +1,8 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import List, Literal, Optional, Union, cast
+from typing import AsyncGenerator, List, Literal, Optional, Union, cast
 
+import anyio
 import httpx
 import pendulum
 
@@ -11,9 +12,10 @@ except ImportError:
     # pendulum < 3
     from pendulum.period import Period as Interval  # type: ignore
 
-from prefect import get_client
+from prefect.client.orchestration import get_client
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
 
+from .context import ConcurrencyContext
 from .events import (
     _emit_concurrency_acquisition_events,
     _emit_concurrency_release_events,
@@ -34,7 +36,9 @@ async def concurrency(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
-):
+    create_if_missing: bool = True,
+    max_retries: Optional[int] = None,
+) -> AsyncGenerator[None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
 
@@ -43,6 +47,8 @@ async def concurrency(
         occupy: The number of slots to acquire and hold from each limit.
         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.
+        create_if_missing: Whether to create the concurrency limits if they do not exist.
+        max_retries: The maximum number of retries to acquire the concurrency slots.
 
     Raises:
         TimeoutError: If the slots are not acquired within the given timeout.
@@ -60,9 +66,18 @@ async def concurrency(
         await resource_heavy()
     ```
     """
+    if not names:
+        yield
+        return
+
     names = names if isinstance(names, list) else [names]
+
     limits = await _acquire_concurrency_slots(
-        names, occupy, timeout_seconds=timeout_seconds
+        names,
+        occupy,
+        timeout_seconds=timeout_seconds,
+        create_if_missing=create_if_missing,
+        max_retries=max_retries,
     )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
@@ -71,13 +86,28 @@ async def concurrency(
         yield
     finally:
         occupancy_period = cast(Interval, (pendulum.now("UTC") - acquisition_time))
-        await _release_concurrency_slots(
-            names, occupy, occupancy_period.total_seconds()
-        )
+        try:
+            await _release_concurrency_slots(
+                names, occupy, 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, occupy, occupancy_period.total_seconds())
+                )
+
         _emit_concurrency_release_events(limits, occupy, emitted_events)
 
 
-async def rate_limit(names: Union[str, List[str]], occupy: int = 1):
+async def rate_limit(
+    names: Union[str, List[str]],
+    occupy: int = 1,
+    timeout_seconds: Optional[float] = None,
+    create_if_missing: Optional[bool] = True,
+) -> None:
     """Block execution until an `occupy` number of slots of the concurrency
     limits given in `names` are acquired. Requires that all given concurrency
     limits have a slot decay.
@@ -85,9 +115,22 @@ async def rate_limit(names: Union[str, List[str]], occupy: int = 1):
     Args:
         names: The names of the concurrency limits to acquire slots from.
         occupy: The number of slots to acquire and hold from each limit.
+        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.
+        create_if_missing: Whether to create the concurrency limits if they do not exist.
     """
+    if not names:
+        return
+
     names = names if isinstance(names, list) else [names]
-    limits = await _acquire_concurrency_slots(names, occupy, mode="rate_limit")
+
+    limits = await _acquire_concurrency_slots(
+        names,
+        occupy,
+        mode="rate_limit",
+        timeout_seconds=timeout_seconds,
+        create_if_missing=create_if_missing,
+    )
     _emit_concurrency_acquisition_events(limits, occupy)
 
 
@@ -96,9 +139,13 @@ async def _acquire_concurrency_slots(
     slots: int,
     mode: Union[Literal["concurrency"], Literal["rate_limit"]] = "concurrency",
     timeout_seconds: Optional[float] = None,
+    create_if_missing: Optional[bool] = True,
+    max_retries: Optional[int] = None,
 ) -> List[MinimalConcurrencyLimitResponse]:
     service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
-    future = service.send((slots, mode, timeout_seconds))
+    future = service.send(
+        (slots, mode, timeout_seconds, create_if_missing, max_retries)
+    )
     response_or_exception = await asyncio.wrap_future(future)
 
     if isinstance(response_or_exception, Exception):
@@ -127,4 +174,6 @@ async def _release_concurrency_slots(
 def _response_to_minimal_concurrency_limit_response(
     response: httpx.Response,
 ) -> List[MinimalConcurrencyLimitResponse]:
-    return [MinimalConcurrencyLimitResponse.parse_obj(obj_) for obj_ in response.json()]
+    return [
+        MinimalConcurrencyLimitResponse.model_validate(obj_) for obj_ in response.json()
+    ]

prefect/concurrency/context.py

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

prefect/concurrency/events.py

@@ -20,7 +20,7 @@ def _emit_concurrency_event(
     }
 
     related = [
-        RelatedResource.parse_obj(
+        RelatedResource.model_validate(
             {
                 "prefect.resource.id": f"prefect.concurrency-limit.{limit.id}",
                 "prefect.resource.role": "concurrency-limit",
@@ -54,6 +54,6 @@ def _emit_concurrency_release_events(
     limits: List[MinimalConcurrencyLimitResponse],
     occupy: int,
     events: Dict[UUID, Optional[Event]],
-):
+) -> None:
     for limit in limits:
         _emit_concurrency_event("released", limit, limits, occupy, events[limit.id])