prefect-client 3.0.0rc18__py3-none-any.whl → 3.0.0rc20__py3-none-any.whl

prefect/events/related.py

@@ -21,6 +21,7 @@ from .schemas.events import RelatedResource
 
 if TYPE_CHECKING:
     from prefect._internal.schemas.bases import ObjectBaseModel
+    from prefect.client.orchestration import PrefectClient
 
 ResourceCacheEntry = Dict[str, Union[str, "ObjectBaseModel", None]]
 RelatedResourceCache = Dict[str, Tuple[ResourceCacheEntry, DateTime]]
@@ -54,9 +55,9 @@ def object_as_related_resource(kind: str, role: str, object: Any) -> RelatedReso
 
 
 async def related_resources_from_run_context(
+    client: "PrefectClient",
     exclude: Optional[Set[str]] = None,
 ) -> List[RelatedResource]:
-    from prefect.client.orchestration import get_client
     from prefect.client.schemas.objects import FlowRun
     from prefect.context import FlowRunContext, TaskRunContext
 
@@ -77,86 +78,84 @@ async def related_resources_from_run_context(
 
     related_objects: List[ResourceCacheEntry] = []
 
-    async with get_client() as client:
+    async def dummy_read():
+        return {}
 
-        async def dummy_read():
-            return {}
-
-        if flow_run_context:
-            related_objects.append(
-                {
-                    "kind": "flow-run",
-                    "role": "flow-run",
-                    "object": flow_run_context.flow_run,
-                },
-            )
-        else:
-            related_objects.append(
-                await _get_and_cache_related_object(
-                    kind="flow-run",
-                    role="flow-run",
-                    client_method=client.read_flow_run,
-                    obj_id=flow_run_id,
-                    cache=RESOURCE_CACHE,
-                )
+    if flow_run_context:
+        related_objects.append(
+            {
+                "kind": "flow-run",
+                "role": "flow-run",
+                "object": flow_run_context.flow_run,
+            },
+        )
+    else:
+        related_objects.append(
+            await _get_and_cache_related_object(
+                kind="flow-run",
+                role="flow-run",
+                client_method=client.read_flow_run,
+                obj_id=flow_run_id,
+                cache=RESOURCE_CACHE,
             )
+        )
 
-        if task_run_context:
-            related_objects.append(
-                {
-                    "kind": "task-run",
-                    "role": "task-run",
-                    "object": task_run_context.task_run,
-                },
-            )
+    if task_run_context:
+        related_objects.append(
+            {
+                "kind": "task-run",
+                "role": "task-run",
+                "object": task_run_context.task_run,
+            },
+        )
 
-        flow_run = related_objects[0]["object"]
+    flow_run = related_objects[0]["object"]
 
-        if isinstance(flow_run, FlowRun):
-            related_objects += list(
-                await asyncio.gather(
+    if isinstance(flow_run, FlowRun):
+        related_objects += list(
+            await asyncio.gather(
+                _get_and_cache_related_object(
+                    kind="flow",
+                    role="flow",
+                    client_method=client.read_flow,
+                    obj_id=flow_run.flow_id,
+                    cache=RESOURCE_CACHE,
+                ),
+                (
+                    _get_and_cache_related_object(
+                        kind="deployment",
+                        role="deployment",
+                        client_method=client.read_deployment,
+                        obj_id=flow_run.deployment_id,
+                        cache=RESOURCE_CACHE,
+                    )
+                    if flow_run.deployment_id
+                    else dummy_read()
+                ),
+                (
+                    _get_and_cache_related_object(
+                        kind="work-queue",
+                        role="work-queue",
+                        client_method=client.read_work_queue,
+                        obj_id=flow_run.work_queue_id,
+                        cache=RESOURCE_CACHE,
+                    )
+                    if flow_run.work_queue_id
+                    else dummy_read()
+                ),
+                (
                     _get_and_cache_related_object(
-                        kind="flow",
-                        role="flow",
-                        client_method=client.read_flow,
-                        obj_id=flow_run.flow_id,
+                        kind="work-pool",
+                        role="work-pool",
+                        client_method=client.read_work_pool,
+                        obj_id=flow_run.work_pool_name,
                         cache=RESOURCE_CACHE,
-                    ),
-                    (
-                        _get_and_cache_related_object(
-                            kind="deployment",
-                            role="deployment",
-                            client_method=client.read_deployment,
-                            obj_id=flow_run.deployment_id,
-                            cache=RESOURCE_CACHE,
-                        )
-                        if flow_run.deployment_id
-                        else dummy_read()
-                    ),
-                    (
-                        _get_and_cache_related_object(
-                            kind="work-queue",
-                            role="work-queue",
-                            client_method=client.read_work_queue,
-                            obj_id=flow_run.work_queue_id,
-                            cache=RESOURCE_CACHE,
-                        )
-                        if flow_run.work_queue_id
-                        else dummy_read()
-                    ),
-                    (
-                        _get_and_cache_related_object(
-                            kind="work-pool",
-                            role="work-pool",
-                            client_method=client.read_work_pool,
-                            obj_id=flow_run.work_pool_name,
-                            cache=RESOURCE_CACHE,
-                        )
-                        if flow_run.work_pool_name
-                        else dummy_read()
-                    ),
-                )
+                    )
+                    if flow_run.work_pool_name
+                    else dummy_read()
+                ),
             )
+        )
 
     related = []
     tags = set()

prefect/events/utilities.py

@@ -7,6 +7,7 @@ from pydantic_extra_types.pendulum_dt import DateTime
 
 from .clients import (
     AssertingEventsClient,
+    AssertingPassthroughEventsClient,
     PrefectCloudEventsClient,
     PrefectEventsClient,
 )
@@ -49,6 +50,7 @@ def emit_event(
         return None
 
     operational_clients = [
+        AssertingPassthroughEventsClient,
         AssertingEventsClient,
         PrefectCloudEventsClient,
         PrefectEventsClient,

prefect/events/worker.py

@@ -1,6 +1,6 @@
 from contextlib import asynccontextmanager
 from contextvars import Context, copy_context
-from typing import Any, Dict, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type
 from uuid import UUID
 
 from typing_extensions import Self
@@ -22,6 +22,9 @@ from .clients import (
 from .related import related_resources_from_run_context
 from .schemas.events import Event
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 
 def should_emit_events() -> bool:
     return (
@@ -55,14 +58,18 @@ class EventsWorker(QueueService[Event]):
         self.client_type = client_type
         self.client_options = client_options
         self._client: EventsClient
+        self._orchestration_client: "PrefectClient"
         self._context_cache: Dict[UUID, Context] = {}
 
     @asynccontextmanager
     async def _lifespan(self):
         self._client = self.client_type(**{k: v for k, v in self.client_options})
+        from prefect.client.orchestration import get_client
 
+        self._orchestration_client = get_client()
         async with self._client:
-            yield
+            async with self._orchestration_client:
+                yield
 
     def _prepare_item(self, event: Event) -> Event:
         self._context_cache[event.id] = copy_context()
@@ -77,7 +84,9 @@ class EventsWorker(QueueService[Event]):
 
     async def attach_related_resources_from_context(self, event: Event):
         exclude = {resource.id for resource in event.involved_resources}
-        event.related += await related_resources_from_run_context(exclude=exclude)
+        event.related += await related_resources_from_run_context(
+            client=self._orchestration_client, exclude=exclude
+        )
 
     @classmethod
     def instance(

prefect/exceptions.py

@@ -412,3 +412,9 @@ class PrefectImportError(ImportError):
 
     def __init__(self, message: str) -> None:
         super().__init__(message)
+
+
+class SerializationError(PrefectException):
+    """
+    Raised when an object cannot be serialized.
+    """

prefect/flow_engine.py

@@ -30,6 +30,7 @@ from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.filters import FlowRunFilter
 from prefect.client.schemas.sorting import FlowRunSort
 from prefect.concurrency.context import ConcurrencyContext
+from prefect.concurrency.v1.context import ConcurrencyContext as ConcurrencyContextV1
 from prefect.context import FlowRunContext, SyncClientContext, TagsContext
 from prefect.exceptions import (
     Abort,
@@ -204,6 +205,7 @@ class FlowRunEngine(Generic[P, R]):
                     result_factory=run_coro_as_sync(ResultFactory.from_flow(self.flow)),
                 )
                 self.short_circuit = True
+                self.call_hooks()
 
         new_state = Running()
         state = self.set_state(new_state)
@@ -267,6 +269,7 @@ class FlowRunEngine(Generic[P, R]):
             return_value_to_state(
                 resolved_result,
                 result_factory=result_factory,
+                write_result=True,
             )
         )
         self.set_state(terminal_state)
@@ -286,6 +289,7 @@ class FlowRunEngine(Generic[P, R]):
                 message=msg or "Flow run encountered an exception:",
                 result_factory=result_factory
                 or getattr(context, "result_factory", None),
+                write_result=True,
             )
         )
         state = self.set_state(terminal_state)
@@ -506,6 +510,7 @@ class FlowRunEngine(Generic[P, R]):
                     task_runner=task_runner,
                 )
             )
+            stack.enter_context(ConcurrencyContextV1())
             stack.enter_context(ConcurrencyContext())
 
             # set the logger to the flow run logger

prefect/flows.py

@@ -642,7 +642,8 @@ class Flow(Generic[P, R]):
         cron: Optional[Union[Iterable[str], str]] = None,
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
-        schedules: Optional[List["FlexibleScheduleList"]] = None,
+        schedules: Optional["FlexibleScheduleList"] = None,
+        concurrency_limit: Optional[int] = None,
         parameters: Optional[dict] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         description: Optional[str] = None,
@@ -666,6 +667,7 @@ class Flow(Generic[P, R]):
             paused: Whether or not to set this deployment as paused.
             schedules: A list of schedule objects defining when to execute runs of this deployment.
                 Used to define multiple schedules or additional scheduling options such as `timezone`.
+            concurrency_limit: The maximum number of runs of this deployment that can run at the same time.
             parameters: A dictionary of default parameter values to pass to runs of this deployment.
             triggers: A list of triggers that will kick off runs of this deployment.
             description: A description for the created deployment. Defaults to the flow's
@@ -718,6 +720,7 @@ class Flow(Generic[P, R]):
                 rrule=rrule,
                 paused=paused,
                 schedules=schedules,
+                concurrency_limit=concurrency_limit,
                 tags=tags,
                 triggers=triggers,
                 parameters=parameters or {},
@@ -727,7 +730,7 @@ class Flow(Generic[P, R]):
                 work_pool_name=work_pool_name,
                 work_queue_name=work_queue_name,
                 job_variables=job_variables,
-            )
+            )  # type: ignore # TODO: remove sync_compatible
         else:
             return RunnerDeployment.from_flow(
                 self,
@@ -737,6 +740,7 @@ class Flow(Generic[P, R]):
                 rrule=rrule,
                 paused=paused,
                 schedules=schedules,
+                concurrency_limit=concurrency_limit,
                 tags=tags,
                 triggers=triggers,
                 parameters=parameters or {},
@@ -1055,6 +1059,7 @@ class Flow(Generic[P, R]):
         rrule: Optional[str] = None,
         paused: Optional[bool] = None,
         schedules: Optional[List[DeploymentScheduleCreate]] = None,
+        concurrency_limit: Optional[int] = None,
         triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
         parameters: Optional[dict] = None,
         description: Optional[str] = None,
@@ -1101,6 +1106,7 @@ class Flow(Generic[P, R]):
             paused: Whether or not to set this deployment as paused.
             schedules: A list of schedule objects defining when to execute runs of this deployment.
                 Used to define multiple schedules or additional scheduling options like `timezone`.
+            concurrency_limit: The maximum number of runs that can be executed concurrently.
             parameters: A dictionary of default parameter values to pass to runs of this deployment.
             description: A description for the created deployment. Defaults to the flow's
                 description if not provided.
@@ -1175,6 +1181,7 @@ class Flow(Generic[P, R]):
             cron=cron,
             rrule=rrule,
             schedules=schedules,
+            concurrency_limit=concurrency_limit,
             paused=paused,
             triggers=triggers,
             parameters=parameters,

prefect/logging/handlers.py

@@ -157,7 +157,10 @@ class APILogHandler(logging.Handler):
             if log_handling_when_missing_flow == "warn":
                 # Warn when a logger is used outside of a run context, the stack level here
                 # gets us to the user logging call
-                warnings.warn(str(exc), stacklevel=8)
+                warnings.warn(
+                    f"{exc} Set PREFECT_LOGGING_TO_API_WHEN_MISSING_FLOW=ignore to suppress this warning.",
+                    stacklevel=8,
+                )
                 return
             elif log_handling_when_missing_flow == "ignore":
                 return

prefect/main.py

@@ -6,7 +6,6 @@ from prefect.flows import flow, Flow, serve
 from prefect.transactions import Transaction
 from prefect.tasks import task, Task
 from prefect.context import tags
-from prefect.manifests import Manifest
 from prefect.utilities.annotations import unmapped, allow_failure
 from prefect.results import BaseResult
 from prefect.flow_runs import pause_flow_run, resume_flow_run, suspend_flow_run
@@ -26,10 +25,14 @@ import prefect.context
 # Perform any forward-ref updates needed for Pydantic models
 import prefect.client.schemas
 
-prefect.context.FlowRunContext.model_rebuild()
-prefect.context.TaskRunContext.model_rebuild()
-prefect.client.schemas.State.model_rebuild()
-prefect.client.schemas.StateCreate.model_rebuild()
+prefect.context.FlowRunContext.model_rebuild(
+    _types_namespace={"Flow": Flow, "BaseResult": BaseResult}
+)
+prefect.context.TaskRunContext.model_rebuild(_types_namespace={"Task": Task})
+prefect.client.schemas.State.model_rebuild(_types_namespace={"BaseResult": BaseResult})
+prefect.client.schemas.StateCreate.model_rebuild(
+    _types_namespace={"BaseResult": BaseResult}
+)
 Transaction.model_rebuild()
 
 # Configure logging
@@ -55,7 +58,6 @@ __all__ = [
     "Flow",
     "get_client",
     "get_run_logger",
-    "Manifest",
     "State",
     "tags",
     "task",

prefect/records/base.py

@@ -1,22 +1,88 @@
+import abc
 import os
 import socket
 import threading
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
-from prefect.results import BaseResult
+if TYPE_CHECKING:
+    from prefect.results import BaseResult
+    from prefect.transactions import IsolationLevel
 
 
-class RecordStore:
-    def read(self, key: str):
-        raise NotImplementedError
+@dataclass
+class TransactionRecord:
+    """
+    A dataclass representation of a transaction record.
+    """
 
-    def write(self, key: str, value: dict):
-        raise NotImplementedError
+    key: str
+    result: "BaseResult"
+
+
+class RecordStore(abc.ABC):
+    @abc.abstractmethod
+    def read(
+        self, key: str, holder: Optional[str] = None
+    ) -> Optional[TransactionRecord]:
+        """
+        Read the transaction record with the given key.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock. If a lock exists on
+                the record being written, the read will be blocked until the lock is
+                released if the provided holder does not match the holder of the lock.
+                If not provided, a default holder based on the current host, process,
+                and thread will be used.
+
+        Returns:
+            TransactionRecord: The transaction record with the given key.
+        """
+        ...
+
+    @abc.abstractmethod
+    def write(self, key: str, result: "BaseResult", holder: Optional[str] = None):
+        """
+        Write the transaction record with the given key.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            record: The transaction record to write.
+            holder: Unique identifier for the holder of the lock. If a lock exists on
+                the record being written, the write will be rejected if the provided
+                holder does not match the holder of the lock. If not provided,
+                a default holder based on the current host, process, and thread will
+                be used.
+        """
+        ...
 
+    @abc.abstractmethod
     def exists(self, key: str) -> bool:
-        return False
+        """
+        Check if the transaction record with the given key exists.
+
+        Args:
+            key: Unique identifier for the transaction record.
+
+        Returns:
+            bool: True if the record exists; False otherwise.
+        """
+        ...
+
+    @abc.abstractmethod
+    def supports_isolation_level(self, isolation_level: "IsolationLevel") -> bool:
+        """
+        Check if the record store supports the given isolation level.
+
+        Args:
+            isolation_level: The isolation level to check.
+
+        Returns:
+            bool: True if the record store supports the isolation level; False otherwise.
+        """
+        ...
 
     def acquire_lock(
         self,
@@ -155,13 +221,3 @@ class RecordStore:
             yield
         finally:
             self.release_lock(key=key, holder=holder)
-
-
-@dataclass
-class TransactionRecord:
-    """
-    A dataclass representation of a transaction record.
-    """
-
-    key: str
-    result: Optional[BaseResult] = None