prefect-client 3.1.15__py3-none-any.whl → 3.2.0__py3-none-any.whl

prefect/automations.py

@@ -179,13 +179,11 @@ class Automation(AutomationCore):
 
     @overload
     @classmethod
-    async def aread(cls, id: UUID, name: Optional[str] = ...) -> Self:
-        ...
+    async def aread(cls, id: UUID, name: Optional[str] = ...) -> Self: ...
 
     @overload
     @classmethod
-    async def aread(cls, id: None = None, name: str = ...) -> Self:
-        ...
+    async def aread(cls, id: None = None, name: str = ...) -> Self: ...
 
     @classmethod
     async def aread(cls, id: Optional[UUID] = None, name: Optional[str] = None) -> Self:
@@ -227,13 +225,11 @@ class Automation(AutomationCore):
 
     @overload
     @classmethod
-    async def read(cls, id: UUID, name: Optional[str] = ...) -> Self:
-        ...
+    async def read(cls, id: UUID, name: Optional[str] = ...) -> Self: ...
 
     @overload
     @classmethod
-    async def read(cls, id: None = None, name: str = ...) -> Self:
-        ...
+    async def read(cls, id: None = None, name: str = ...) -> Self: ...
 
     @classmethod
     @async_dispatch(aread)

prefect/blocks/notifications.py

@@ -22,14 +22,14 @@ class AbstractAppriseNotificationBlock(NotificationBlock, ABC):
     An abstract class for sending notifications using Apprise.
     """
 
-    notify_type: Literal[
-        "prefect_default", "info", "success", "warning", "failure"
-    ] = Field(
-        default=PREFECT_NOTIFY_TYPE_DEFAULT,
-        description=(
-            "The type of notification being performed; the prefect_default "
-            "is a plain notification that does not attach an image."
-        ),
+    notify_type: Literal["prefect_default", "info", "success", "warning", "failure"] = (
+        Field(
+            default=PREFECT_NOTIFY_TYPE_DEFAULT,
+            description=(
+                "The type of notification being performed; the prefect_default "
+                "is a plain notification that does not attach an image."
+            ),
+        )
     )
 
     def __init__(self, *args: Any, **kwargs: Any):

prefect/cache_policies.py

@@ -252,6 +252,8 @@ class _None(CachePolicy):
 class TaskSource(CachePolicy):
     """
     Policy for computing a cache key based on the source code of the task.
+
+    This policy only considers raw lines of code in the task, and not the source code of nested tasks.
     """
 
     def compute_key(

prefect/client/base.py

@@ -55,8 +55,7 @@ Send = Callable[[Message], Awaitable[None]]
 
 @runtime_checkable
 class ASGIApp(Protocol):
-    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        ...
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: ...
 
 
 @asynccontextmanager
@@ -301,9 +300,9 @@ class PrefectHttpxAsyncClient(httpx.AsyncClient):
             )
             await anyio.sleep(retry_seconds)
 
-        assert (
-            response is not None
-        ), "Retry handling ended without response or exception"
+        assert response is not None, (
+            "Retry handling ended without response or exception"
+        )
 
         # We ran out of retries, return the failed response
         return response
@@ -520,9 +519,9 @@ class PrefectHttpxSyncClient(httpx.Client):
             )
             time.sleep(retry_seconds)
 
-        assert (
-            response is not None
-        ), "Retry handling ended without response or exception"
+        assert response is not None, (
+            "Retry handling ended without response or exception"
+        )
 
         # We ran out of retries, return the failed response
         return response

prefect/client/collections.py

@@ -7,14 +7,11 @@ from prefect.client.orchestration import ServerType, get_client
 
 
 class CollectionsMetadataClient(Protocol):
-    async def read_worker_metadata(self) -> Dict[str, Any]:
-        ...
+    async def read_worker_metadata(self) -> Dict[str, Any]: ...
 
-    async def __aenter__(self) -> "CollectionsMetadataClient":
-        ...
+    async def __aenter__(self) -> "CollectionsMetadataClient": ...
 
-    async def __aexit__(self, *exc_info: Any) -> Any:
-        ...
+    async def __aexit__(self, *exc_info: Any) -> Any: ...
 
 
 def get_collections_metadata_client(

prefect/client/orchestration/__init__.py

@@ -2,7 +2,6 @@ import asyncio
 import base64
 import datetime
 import ssl
-import warnings
 from collections.abc import Iterable
 from contextlib import AsyncExitStack
 from logging import Logger
@@ -48,6 +47,12 @@ from prefect.client.orchestration._automations.client import (
     AutomationClient,
     AutomationAsyncClient,
 )
+
+from prefect.client.orchestration._work_pools.client import (
+    WorkPoolClient,
+    WorkPoolAsyncClient,
+)
+
 from prefect._experimental.sla.client import SlaClient, SlaAsyncClient
 
 from prefect.client.orchestration._flows.client import (
@@ -85,8 +90,6 @@ from prefect.client.schemas.actions import (
     FlowRunNotificationPolicyUpdate,
     TaskRunCreate,
     TaskRunUpdate,
-    WorkPoolCreate,
-    WorkPoolUpdate,
     WorkQueueCreate,
     WorkQueueUpdate,
 )
@@ -96,8 +99,6 @@ from prefect.client.schemas.filters import (
     FlowRunFilter,
     FlowRunNotificationPolicyFilter,
     TaskRunFilter,
-    WorkerFilter,
-    WorkPoolFilter,
     WorkQueueFilter,
     WorkQueueFilterName,
 )
@@ -107,15 +108,10 @@ from prefect.client.schemas.objects import (
     Parameter,
     TaskRunPolicy,
     TaskRunResult,
-    Worker,
-    WorkerMetadata,
-    WorkPool,
     WorkQueue,
     WorkQueueStatusDetail,
 )
-from prefect.client.schemas.responses import (
-    WorkerFlowRunResponse,
-)
+
 from prefect.client.schemas.sorting import (
     TaskRunSort,
 )
@@ -133,7 +129,6 @@ from prefect.settings import (
     PREFECT_CLOUD_API_URL,
     PREFECT_SERVER_ALLOW_EPHEMERAL_MODE,
     PREFECT_TESTING_UNIT_TEST_MODE,
-    get_current_settings,
 )
 
 if TYPE_CHECKING:
@@ -157,15 +152,13 @@ def get_client(
     *,
     httpx_settings: Optional[dict[str, Any]] = ...,
     sync_client: Literal[False] = False,
-) -> "PrefectClient":
-    ...
+) -> "PrefectClient": ...
 
 
 @overload
 def get_client(
     *, httpx_settings: Optional[dict[str, Any]] = ..., sync_client: Literal[True] = ...
-) -> "SyncPrefectClient":
-    ...
+) -> "SyncPrefectClient": ...
 
 
 def get_client(
@@ -263,6 +256,7 @@ class PrefectClient(
     BlocksDocumentAsyncClient,
     BlocksSchemaAsyncClient,
     BlocksTypeAsyncClient,
+    WorkPoolAsyncClient,
 ):
     """
     An asynchronous client for interacting with the [Prefect REST API](/api-ref/rest-api/).
@@ -814,7 +808,7 @@ class PrefectClient(
                 retry_delay=retry_delay,
                 retry_jitter_factor=task.retry_jitter_factor,
             ),
-            state=state.to_state_create(),
+            state=prefect.states.to_state_create(state),
             task_inputs=task_inputs or {},
         )
         content = task_run_data.model_dump_json(exclude={"id"} if id is None else None)
@@ -925,7 +919,7 @@ class PrefectClient(
         Returns:
             an OrchestrationResult model representation of state orchestration output
         """
-        state_create = state.to_state_create()
+        state_create = prefect.states.to_state_create(state)
         state_create.state_details.task_run_id = task_run_id
         response = await self._client.post(
             f"/task_runs/{task_run_id}/set_state",
@@ -1100,214 +1094,6 @@ class PrefectClient(
             response.json()
         )
 
-    async def send_worker_heartbeat(
-        self,
-        work_pool_name: str,
-        worker_name: str,
-        heartbeat_interval_seconds: Optional[float] = None,
-        get_worker_id: bool = False,
-        worker_metadata: Optional[WorkerMetadata] = None,
-    ) -> Optional[UUID]:
-        """
-        Sends a worker heartbeat for a given work pool.
-
-        Args:
-            work_pool_name: The name of the work pool to heartbeat against.
-            worker_name: The name of the worker sending the heartbeat.
-            return_id: Whether to return the worker ID. Note: will return `None` if the connected server does not support returning worker IDs, even if `return_id` is `True`.
-            worker_metadata: Metadata about the worker to send to the server.
-        """
-        params: dict[str, Any] = {
-            "name": worker_name,
-            "heartbeat_interval_seconds": heartbeat_interval_seconds,
-        }
-        if worker_metadata:
-            params["metadata"] = worker_metadata.model_dump(mode="json")
-        if get_worker_id:
-            params["return_id"] = get_worker_id
-
-        resp = await self._client.post(
-            f"/work_pools/{work_pool_name}/workers/heartbeat",
-            json=params,
-        )
-
-        if (
-            (
-                self.server_type == ServerType.CLOUD
-                or get_current_settings().testing.test_mode
-            )
-            and get_worker_id
-            and resp.status_code == 200
-        ):
-            return UUID(resp.text)
-        else:
-            return None
-
-    async def read_workers_for_work_pool(
-        self,
-        work_pool_name: str,
-        worker_filter: Optional[WorkerFilter] = None,
-        offset: Optional[int] = None,
-        limit: Optional[int] = None,
-    ) -> list[Worker]:
-        """
-        Reads workers for a given work pool.
-
-        Args:
-            work_pool_name: The name of the work pool for which to get
-                member workers.
-            worker_filter: Criteria by which to filter workers.
-            limit: Limit for the worker query.
-            offset: Limit for the worker query.
-        """
-        response = await self._client.post(
-            f"/work_pools/{work_pool_name}/workers/filter",
-            json={
-                "workers": (
-                    worker_filter.model_dump(mode="json", exclude_unset=True)
-                    if worker_filter
-                    else None
-                ),
-                "offset": offset,
-                "limit": limit,
-            },
-        )
-
-        return pydantic.TypeAdapter(list[Worker]).validate_python(response.json())
-
-    async def read_work_pool(self, work_pool_name: str) -> WorkPool:
-        """
-        Reads information for a given work pool
-
-        Args:
-            work_pool_name: The name of the work pool to for which to get
-                information.
-
-        Returns:
-            Information about the requested work pool.
-        """
-        try:
-            response = await self._client.get(f"/work_pools/{work_pool_name}")
-            return WorkPool.model_validate(response.json())
-        except httpx.HTTPStatusError as e:
-            if e.response.status_code == status.HTTP_404_NOT_FOUND:
-                raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
-            else:
-                raise
-
-    async def read_work_pools(
-        self,
-        limit: Optional[int] = None,
-        offset: int = 0,
-        work_pool_filter: Optional[WorkPoolFilter] = None,
-    ) -> list[WorkPool]:
-        """
-        Reads work pools.
-
-        Args:
-            limit: Limit for the work pool query.
-            offset: Offset for the work pool query.
-            work_pool_filter: Criteria by which to filter work pools.
-
-        Returns:
-            A list of work pools.
-        """
-
-        body: dict[str, Any] = {
-            "limit": limit,
-            "offset": offset,
-            "work_pools": (
-                work_pool_filter.model_dump(mode="json") if work_pool_filter else None
-            ),
-        }
-        response = await self._client.post("/work_pools/filter", json=body)
-        return pydantic.TypeAdapter(list[WorkPool]).validate_python(response.json())
-
-    async def create_work_pool(
-        self,
-        work_pool: WorkPoolCreate,
-        overwrite: bool = False,
-    ) -> WorkPool:
-        """
-        Creates a work pool with the provided configuration.
-
-        Args:
-            work_pool: Desired configuration for the new work pool.
-
-        Returns:
-            Information about the newly created work pool.
-        """
-        try:
-            response = await self._client.post(
-                "/work_pools/",
-                json=work_pool.model_dump(mode="json", exclude_unset=True),
-            )
-        except httpx.HTTPStatusError as e:
-            if e.response.status_code == status.HTTP_409_CONFLICT:
-                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
-
-        return WorkPool.model_validate(response.json())
-
-    async def update_work_pool(
-        self,
-        work_pool_name: str,
-        work_pool: WorkPoolUpdate,
-    ) -> None:
-        """
-        Updates a work pool.
-
-        Args:
-            work_pool_name: Name of the work pool to update.
-            work_pool: Fields to update in the work pool.
-        """
-        try:
-            await self._client.patch(
-                f"/work_pools/{work_pool_name}",
-                json=work_pool.model_dump(mode="json", exclude_unset=True),
-            )
-        except httpx.HTTPStatusError as e:
-            if e.response.status_code == status.HTTP_404_NOT_FOUND:
-                raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
-            else:
-                raise
-
-    async def delete_work_pool(
-        self,
-        work_pool_name: str,
-    ) -> None:
-        """
-        Deletes a work pool.
-
-        Args:
-            work_pool_name: Name of the work pool to delete.
-        """
-        try:
-            await self._client.delete(f"/work_pools/{work_pool_name}")
-        except httpx.HTTPStatusError as e:
-            if e.response.status_code == status.HTTP_404_NOT_FOUND:
-                raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
-            else:
-                raise
-
     async def read_work_queues(
         self,
         work_pool_name: Optional[str] = None,
@@ -1353,41 +1139,6 @@ class PrefectClient(
 
         return pydantic.TypeAdapter(list[WorkQueue]).validate_python(response.json())
 
-    async def get_scheduled_flow_runs_for_work_pool(
-        self,
-        work_pool_name: str,
-        work_queue_names: Optional[list[str]] = None,
-        scheduled_before: Optional[datetime.datetime] = None,
-    ) -> list[WorkerFlowRunResponse]:
-        """
-        Retrieves scheduled flow runs for the provided set of work pool queues.
-
-        Args:
-            work_pool_name: The name of the work pool that the work pool
-                queues are associated with.
-            work_queue_names: The names of the work pool queues from which
-                to get scheduled flow runs.
-            scheduled_before: Datetime used to filter returned flow runs. Flow runs
-                scheduled for after the given datetime string will not be returned.
-
-        Returns:
-            A list of worker flow run responses containing information about the
-            retrieved flow runs.
-        """
-        body: dict[str, Any] = {}
-        if work_queue_names is not None:
-            body["work_queue_names"] = list(work_queue_names)
-        if scheduled_before:
-            body["scheduled_before"] = str(scheduled_before)
-
-        response = await self._client.post(
-            f"/work_pools/{work_pool_name}/get_scheduled_flow_runs",
-            json=body,
-        )
-        return pydantic.TypeAdapter(list[WorkerFlowRunResponse]).validate_python(
-            response.json()
-        )
-
     async def read_worker_metadata(self) -> dict[str, Any]:
         """Reads worker metadata stored in Prefect collection registry."""
         response = await self._client.get("collections/views/aggregate-worker-metadata")
@@ -1504,6 +1255,7 @@ class SyncPrefectClient(
     BlocksDocumentClient,
     BlocksSchemaClient,
     BlocksTypeClient,
+    WorkPoolClient,
 ):
     """
     A synchronous client for interacting with the [Prefect REST API](/api-ref/rest-api/).
@@ -1834,7 +1586,7 @@ class SyncPrefectClient(
                 retry_delay=retry_delay,
                 retry_jitter_factor=task.retry_jitter_factor,
             ),
-            state=state.to_state_create(),
+            state=prefect.states.to_state_create(state),
             task_inputs=task_inputs or {},
         )
 
@@ -1928,7 +1680,7 @@ class SyncPrefectClient(
         Returns:
             an OrchestrationResult model representation of state orchestration output
         """
-        state_create = state.to_state_create()
+        state_create = prefect.states.to_state_create(state)
         state_create.state_details.task_run_id = task_run_id
         response = self._client.post(
             f"/task_runs/{task_run_id}/set_state",

prefect/client/orchestration/_deployments/client.py

@@ -165,7 +165,11 @@ class DeploymentClient(BaseClient):
             "PATCH",
             "/deployments/{id}",
             path_params={"id": deployment_id},
-            json=deployment.model_dump(mode="json", exclude_unset=True),
+            json=deployment.model_dump(
+                mode="json",
+                exclude_unset=True,
+                exclude={"name", "flow_name", "triggers"},
+            ),
         )
 
     def _create_deployment_from_schema(self, schema: "DeploymentCreate") -> "UUID":
@@ -550,7 +554,7 @@ class DeploymentClient(BaseClient):
         """
         from prefect.client.schemas.actions import DeploymentFlowRunCreate
         from prefect.client.schemas.objects import FlowRun
-        from prefect.states import Scheduled
+        from prefect.states import Scheduled, to_state_create
 
         parameters = parameters or {}
         context = context or {}
@@ -561,7 +565,7 @@ class DeploymentClient(BaseClient):
         flow_run_create = DeploymentFlowRunCreate(
             parameters=parameters,
             context=context,
-            state=state.to_state_create(),
+            state=to_state_create(state),
             tags=list(tags),
             name=name,
             idempotency_key=idempotency_key,
@@ -708,7 +712,11 @@ class DeploymentAsyncClient(BaseAsyncClient):
             "PATCH",
             "/deployments/{id}",
             path_params={"id": deployment_id},
-            json=deployment.model_dump(mode="json", exclude_unset=True),
+            json=deployment.model_dump(
+                mode="json",
+                exclude_unset=True,
+                exclude={"name", "flow_name", "triggers"},
+            ),
         )
 
     async def _create_deployment_from_schema(
@@ -1095,7 +1103,7 @@ class DeploymentAsyncClient(BaseAsyncClient):
         """
         from prefect.client.schemas.actions import DeploymentFlowRunCreate
         from prefect.client.schemas.objects import FlowRun
-        from prefect.states import Scheduled
+        from prefect.states import Scheduled, to_state_create
 
         parameters = parameters or {}
         context = context or {}
@@ -1106,7 +1114,7 @@ class DeploymentAsyncClient(BaseAsyncClient):
         flow_run_create = DeploymentFlowRunCreate(
             parameters=parameters,
             context=context,
-            state=state.to_state_create(),
+            state=to_state_create(state),
             tags=list(tags),
             name=name,
             idempotency_key=idempotency_key,

prefect/client/orchestration/_flow_runs/client.py

@@ -69,7 +69,7 @@ class FlowRunClient(BaseClient):
         """
         from prefect.client.schemas.actions import FlowCreate, FlowRunCreate
         from prefect.client.schemas.objects import Flow, FlowRun, FlowRunPolicy
-        from prefect.states import Pending
+        from prefect.states import Pending, to_state_create
 
         parameters = parameters or {}
         context = context or {}
@@ -93,7 +93,7 @@ class FlowRunClient(BaseClient):
             context=context,
             tags=list(tags or []),
             parent_task_run_id=parent_task_run_id,
-            state=state.to_state_create(),
+            state=to_state_create(state),
             empirical_policy=FlowRunPolicy(
                 retries=flow.retries,
                 retry_delay=int(flow.retry_delay_seconds or 0),
@@ -322,10 +322,12 @@ class FlowRunClient(BaseClient):
         """
         from uuid import UUID, uuid4
 
+        from prefect.states import to_state_create
+
         flow_run_id = (
             flow_run_id if isinstance(flow_run_id, UUID) else UUID(flow_run_id)
         )
-        state_create = state.to_state_create()
+        state_create = to_state_create(state)
         state_create.state_details.flow_run_id = flow_run_id
         state_create.state_details.transition_id = uuid4()
         try:
@@ -501,7 +503,7 @@ class FlowRunAsyncClient(BaseAsyncClient):
         """
         from prefect.client.schemas.actions import FlowCreate, FlowRunCreate
         from prefect.client.schemas.objects import Flow, FlowRun, FlowRunPolicy
-        from prefect.states import Pending
+        from prefect.states import Pending, to_state_create
 
         parameters = parameters or {}
         context = context or {}
@@ -525,7 +527,7 @@ class FlowRunAsyncClient(BaseAsyncClient):
             context=context,
             tags=list(tags or []),
             parent_task_run_id=parent_task_run_id,
-            state=state.to_state_create(),
+            state=to_state_create(state),
             empirical_policy=FlowRunPolicy(
                 retries=flow.retries,
                 retry_delay=int(flow.retry_delay_seconds or 0),
@@ -755,10 +757,12 @@ class FlowRunAsyncClient(BaseAsyncClient):
         """
         from uuid import UUID, uuid4
 
+        from prefect.states import to_state_create
+
         flow_run_id = (
             flow_run_id if isinstance(flow_run_id, UUID) else UUID(flow_run_id)
         )
-        state_create = state.to_state_create()
+        state_create = to_state_create(state)
         state_create.state_details.flow_run_id = flow_run_id
         state_create.state_details.transition_id = uuid4()
         try: