prefect-client 3.2.7__py3-none-any.whl → 3.2.8__py3-none-any.whl

prefect/flow_runs.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 from typing import (
     TYPE_CHECKING,
     Any,
-    Optional,
     Type,
     TypeVar,
     overload,
@@ -13,6 +14,7 @@ import anyio
 from prefect.client.orchestration import PrefectClient, get_client
 from prefect.client.schemas import FlowRun
 from prefect.client.schemas.objects import (
+    State,
     StateType,
 )
 from prefect.client.schemas.responses import SetStateStatus
@@ -21,6 +23,8 @@ from prefect.context import (
     FlowRunContext,
     TaskRunContext,
 )
+from prefect.events.clients import get_events_subscriber
+from prefect.events.filters import EventFilter, EventNameFilter, EventResourceFilter
 from prefect.exceptions import (
     Abort,
     FlowPauseTimeout,
@@ -52,9 +56,9 @@ if TYPE_CHECKING:
 @inject_client
 async def wait_for_flow_run(
     flow_run_id: UUID,
-    timeout: Optional[int] = 10800,
-    poll_interval: int = 5,
-    client: Optional["PrefectClient"] = None,
+    timeout: int | None = 10800,
+    poll_interval: int | None = None,
+    client: "PrefectClient | None" = None,
     log_states: bool = False,
 ) -> FlowRun:
     """
@@ -63,7 +67,9 @@ async def wait_for_flow_run(
     Args:
         flow_run_id: The flow run ID for the flow run to wait for.
         timeout: The wait timeout in seconds. Defaults to 10800 (3 hours).
-        poll_interval: The poll interval in seconds. Defaults to 5.
+        poll_interval: Deprecated; polling is no longer used to wait for flow runs.
+        client: Optional Prefect client. If not provided, one will be injected.
+        log_states: If True, log state changes. Defaults to False.
 
     Returns:
         FlowRun: The finished flow run.
@@ -113,17 +119,37 @@ async def wait_for_flow_run(
 
             ```
     """
+    if poll_interval is not None:
+        get_logger().warning(
+            "The `poll_interval` argument is deprecated and will be removed in a future release. "
+        )
+
     assert client is not None, "Client injection failed"
     logger = get_logger()
+
+    event_filter = EventFilter(
+        event=EventNameFilter(prefix=["prefect.flow-run"]),
+        resource=EventResourceFilter(id=[f"prefect.flow-run.{flow_run_id}"]),
+    )
+
     with anyio.move_on_after(timeout):
-        while True:
+        async with get_events_subscriber(filter=event_filter) as subscriber:
             flow_run = await client.read_flow_run(flow_run_id)
-            flow_state = flow_run.state
-            if log_states:
-                logger.info(f"Flow run is in state {flow_run.state.name!r}")
-            if flow_state and flow_state.is_final():
+            if flow_run.state and flow_run.state.is_final():
+                if log_states:
+                    logger.info(f"Flow run is in state {flow_run.state.name!r}")
                 return flow_run
-            await anyio.sleep(poll_interval)
+
+            async for event in subscriber:
+                state_type = StateType(event.resource["prefect.state-type"])
+                state = State(type=state_type)
+
+                if log_states:
+                    logger.info(f"Flow run is in state {state.name!r}")
+
+                if state.is_final():
+                    return await client.read_flow_run(flow_run_id)
+
     raise FlowRunWaitTimeout(
         f"Flow run with ID {flow_run_id} exceeded watch timeout of {timeout} seconds"
     )
@@ -138,7 +164,7 @@ async def pause_flow_run(
     wait_for_input: None = None,
     timeout: int = 3600,
     poll_interval: int = 10,
-    key: Optional[str] = None,
+    key: str | None = None,
 ) -> None: ...
 
 
@@ -147,17 +173,17 @@ async def pause_flow_run(
     wait_for_input: Type[T],
     timeout: int = 3600,
     poll_interval: int = 10,
-    key: Optional[str] = None,
+    key: str | None = None,
 ) -> T: ...
 
 
 @sync_compatible
 async def pause_flow_run(
-    wait_for_input: Optional[Type[T]] = None,
+    wait_for_input: Type[T] | None = None,
     timeout: int = 3600,
     poll_interval: int = 10,
-    key: Optional[str] = None,
-) -> Optional[T]:
+    key: str | None = None,
+) -> T | None:
     """
     Pauses the current flow run by blocking execution until resumed.
 
@@ -213,10 +239,13 @@ async def pause_flow_run(
 async def _in_process_pause(
     timeout: int = 3600,
     poll_interval: int = 10,
-    key: Optional[str] = None,
-    client=None,
-    wait_for_input: Optional[T] = None,
-) -> Optional[T]:
+    key: str | None = None,
+    client: "PrefectClient | None" = None,
+    wait_for_input: Type[T] | None = None,
+) -> T | None:
+    if TYPE_CHECKING:
+        assert client is not None
+
     if TaskRunContext.get():
         raise RuntimeError("Cannot pause task runs.")
 
@@ -302,32 +331,32 @@ async def _in_process_pause(
 @overload
 async def suspend_flow_run(
     wait_for_input: None = None,
-    flow_run_id: Optional[UUID] = None,
-    timeout: Optional[int] = 3600,
-    key: Optional[str] = None,
-    client: Optional[PrefectClient] = None,
+    flow_run_id: UUID | None = None,
+    timeout: int | None = 3600,
+    key: str | None = None,
+    client: "PrefectClient | None" = None,
 ) -> None: ...
 
 
 @overload
 async def suspend_flow_run(
     wait_for_input: Type[T],
-    flow_run_id: Optional[UUID] = None,
-    timeout: Optional[int] = 3600,
-    key: Optional[str] = None,
-    client: Optional[PrefectClient] = None,
+    flow_run_id: UUID | None = None,
+    timeout: int | None = 3600,
+    key: str | None = None,
+    client: "PrefectClient | None" = None,
 ) -> T: ...
 
 
 @sync_compatible
 @inject_client
 async def suspend_flow_run(
-    wait_for_input: Optional[Type[T]] = None,
-    flow_run_id: Optional[UUID] = None,
-    timeout: Optional[int] = 3600,
-    key: Optional[str] = None,
-    client: Optional[PrefectClient] = None,
-) -> Optional[T]:
+    wait_for_input: Type[T] | None = None,
+    flow_run_id: UUID | None = None,
+    timeout: int | None = 3600,
+    key: str | None = None,
+    client: "PrefectClient | None" = None,
+) -> T | None:
     """
     Suspends a flow run by stopping code execution until resumed.
 
@@ -357,6 +386,9 @@ async def suspend_flow_run(
             resumed with the input, the flow will resume and the input will be
             loaded and returned from this function.
     """
+    if TYPE_CHECKING:
+        assert client is not None
+
     context = FlowRunContext.get()
 
     if flow_run_id is None:
@@ -427,7 +459,7 @@ async def suspend_flow_run(
 
 @sync_compatible
 async def resume_flow_run(
-    flow_run_id: UUID, run_input: Optional[dict[str, Any]] = None
+    flow_run_id: UUID, run_input: dict[str, Any] | None = None
 ) -> None:
     """
     Resumes a paused flow.

prefect/flows.py

@@ -2353,7 +2353,9 @@ async def load_flow_from_flow_run(
         from prefect.deployments.steps.core import StepExecutionError, run_steps
 
         try:
-            output = await run_steps(deployment.pull_steps)
+            output = await run_steps(
+                deployment.pull_steps, print_function=run_logger.info
+            )
         except StepExecutionError as e:
             e = e.__cause__ or e
             run_logger.error(str(e))

prefect/futures.py

@@ -1,14 +1,18 @@
+from __future__ import annotations
+
 import abc
 import asyncio
 import concurrent.futures
 import threading
 import uuid
+import warnings
 from collections.abc import Generator, Iterator
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Generic
 
 from typing_extensions import NamedTuple, Self, TypeVar
 
+from prefect._waiters import FlowRunWaiter
 from prefect.client.orchestration import get_client
 from prefect.exceptions import ObjectNotFound
 from prefect.logging.loggers import get_logger, get_run_logger
@@ -31,22 +35,42 @@ logger: "logging.Logger" = get_logger(__name__)
 class PrefectFuture(abc.ABC, Generic[R]):
     """
     Abstract base class for Prefect futures. A Prefect future is a handle to the
-    asynchronous execution of a task run. It provides methods to wait for the task
-    to complete and to retrieve the result of the task run.
+    asynchronous execution of a run. It provides methods to wait for the
+    to complete and to retrieve the result of the run.
     """
 
     def __init__(self, task_run_id: uuid.UUID):
+        warnings.warn(
+            "The __init__ method of PrefectFuture is deprecated and will be removed in a future release. "
+            "If you are subclassing PrefectFuture, please implement the __init__ method in your subclass or "
+            "subclass PrefectTaskRunFuture instead.",
+            DeprecationWarning,
+        )
         self._task_run_id = task_run_id
-        self._final_state: Optional[State[R]] = None
+        self._final_state: State[R] | None = None
 
     @property
     def task_run_id(self) -> uuid.UUID:
         """The ID of the task run associated with this future"""
+        warnings.warn(
+            "The task_run_id property of PrefectFuture is deprecated and will be removed in a future release. "
+            "If you are subclassing PrefectFuture, please implement the task_run_id property in your subclass or "
+            "subclass PrefectTaskRunFuture instead.",
+            DeprecationWarning,
+        )
+
         return self._task_run_id
 
     @property
     def state(self) -> State:
         """The current state of the task run associated with this future"""
+        warnings.warn(
+            "The state property of PrefectFuture is deprecated and will be removed in a future release. "
+            "If you are subclassing PrefectFuture, please implement the state property in your subclass or "
+            "subclass PrefectTaskRunFuture instead.",
+            DeprecationWarning,
+        )
+
         if self._final_state:
             return self._final_state
         client = get_client(sync_client=True)
@@ -59,7 +83,7 @@ class PrefectFuture(abc.ABC, Generic[R]):
         return task_run.state or Pending()
 
     @abc.abstractmethod
-    def wait(self, timeout: Optional[float] = None) -> None:
+    def wait(self, timeout: float | None = None) -> None:
         ...
         """
         Wait for the task run to complete.
@@ -74,7 +98,7 @@ class PrefectFuture(abc.ABC, Generic[R]):
     @abc.abstractmethod
     def result(
         self,
-        timeout: Optional[float] = None,
+        timeout: float | None = None,
         raise_on_failure: bool = True,
     ) -> R:
         ...
@@ -103,7 +127,36 @@ class PrefectFuture(abc.ABC, Generic[R]):
         ...
 
 
-class PrefectWrappedFuture(PrefectFuture[R], abc.ABC, Generic[R, F]):
+class PrefectTaskRunFuture(PrefectFuture[R]):
+    """
+    A Prefect future that represents the eventual execution of a task run.
+    """
+
+    def __init__(self, task_run_id: uuid.UUID):
+        self._task_run_id = task_run_id
+        self._final_state: State[R] | None = None
+
+    @property
+    def task_run_id(self) -> uuid.UUID:
+        """The ID of the task run associated with this future"""
+        return self._task_run_id
+
+    @property
+    def state(self) -> State:
+        """The current state of the task run associated with this future"""
+        if self._final_state:
+            return self._final_state
+        client = get_client(sync_client=True)
+        try:
+            task_run = client.read_task_run(task_run_id=self.task_run_id)
+        except ObjectNotFound:
+            # We'll be optimistic and assume this task will eventually start
+            # TODO: Consider using task run events to wait for the task to start
+            return Pending()
+        return task_run.state or Pending()
+
+
+class PrefectWrappedFuture(PrefectTaskRunFuture[R], abc.ABC, Generic[R, F]):
     """
     A Prefect future that wraps another future object.
 
@@ -140,7 +193,7 @@ class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future[
     when the task run is submitted to a ThreadPoolExecutor.
     """
 
-    def wait(self, timeout: Optional[float] = None) -> None:
+    def wait(self, timeout: float | None = None) -> None:
         try:
             result = self._wrapped_future.result(timeout=timeout)
         except concurrent.futures.TimeoutError:
@@ -150,7 +203,7 @@ class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future[
 
     def result(
         self,
-        timeout: Optional[float] = None,
+        timeout: float | None = None,
         raise_on_failure: bool = True,
     ) -> R:
         if not self._final_state:
@@ -190,7 +243,7 @@ class PrefectConcurrentFuture(PrefectWrappedFuture[R, concurrent.futures.Future[
         )
 
 
-class PrefectDistributedFuture(PrefectFuture[R]):
+class PrefectDistributedFuture(PrefectTaskRunFuture[R]):
     """
     Represents the result of a computation happening anywhere.
 
@@ -202,10 +255,10 @@ class PrefectDistributedFuture(PrefectFuture[R]):
     done_callbacks: list[Callable[[PrefectFuture[R]], None]] = []
     waiter = None
 
-    def wait(self, timeout: Optional[float] = None) -> None:
+    def wait(self, timeout: float | None = None) -> None:
         return run_coro_as_sync(self.wait_async(timeout=timeout))
 
-    async def wait_async(self, timeout: Optional[float] = None) -> None:
+    async def wait_async(self, timeout: float | None = None) -> None:
         if self._final_state:
             logger.debug(
                 "Final state already set for %s. Returning...", self.task_run_id
@@ -244,7 +297,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
 
     def result(
         self,
-        timeout: Optional[float] = None,
+        timeout: float | None = None,
         raise_on_failure: bool = True,
     ) -> R:
         return run_coro_as_sync(
@@ -253,7 +306,7 @@ class PrefectDistributedFuture(PrefectFuture[R]):
 
     async def result_async(
         self,
-        timeout: Optional[float] = None,
+        timeout: float | None = None,
         raise_on_failure: bool = True,
     ) -> R:
         if not self._final_state:
@@ -289,6 +342,123 @@ class PrefectDistributedFuture(PrefectFuture[R]):
         return hash(self.task_run_id)
 
 
+class PrefectFlowRunFuture(PrefectFuture[R]):
+    """
+    A Prefect future that represents the eventual execution of a flow run.
+    """
+
+    def __init__(self, flow_run_id: uuid.UUID):
+        self._flow_run_id = flow_run_id
+        self._final_state: State[R] | None = None
+
+    @property
+    def flow_run_id(self) -> uuid.UUID:
+        """The ID of the flow run associated with this future"""
+        return self._flow_run_id
+
+    @property
+    def state(self) -> State:
+        """The current state of the flow run associated with this future"""
+        if self._final_state:
+            return self._final_state
+        client = get_client(sync_client=True)
+        state = Pending()
+        try:
+            flow_run = client.read_flow_run(flow_run_id=self.flow_run_id)
+            if flow_run.state:
+                state = flow_run.state
+        except ObjectNotFound:
+            # We'll be optimistic and assume this flow run will eventually start
+            pass
+        return state
+
+    def wait(self, timeout: float | None = None) -> None:
+        return run_coro_as_sync(self.wait_async(timeout=timeout))
+
+    async def wait_async(self, timeout: float | None = None) -> None:
+        if self._final_state:
+            logger.debug(
+                "Final state already set for %s. Returning...", self.task_run_id
+            )
+            return
+
+        # Ask for the instance of FlowRunWaiter _now_ so that it's already running and
+        # can catch the completion event if it happens before we start listening for it.
+        FlowRunWaiter.instance()
+
+        # Read task run to see if it is still running
+        async with get_client() as client:
+            flow_run = await client.read_flow_run(flow_run_id=self._flow_run_id)
+            if flow_run.state is None:
+                raise RuntimeError(
+                    f"Flow run {self.flow_run_id} has no state which means it hasn't started yet."
+                )
+            if flow_run.state and flow_run.state.is_final():
+                logger.debug(
+                    "Flow run %s already finished. Returning...",
+                    self.flow_run_id,
+                )
+                self._final_state = flow_run.state
+                return
+
+            # If still running, wait for a completed event from the server
+            logger.debug(
+                "Waiting for completed event for flow run %s...",
+                self.flow_run_id,
+            )
+            await FlowRunWaiter.wait_for_flow_run(self._flow_run_id, timeout=timeout)
+            flow_run = await client.read_flow_run(flow_run_id=self._flow_run_id)
+            if flow_run.state and flow_run.state.is_final():
+                self._final_state = flow_run.state
+            return
+
+    def result(
+        self,
+        timeout: float | None = None,
+        raise_on_failure: bool = True,
+    ) -> R:
+        return run_coro_as_sync(
+            self.aresult(timeout=timeout, raise_on_failure=raise_on_failure)
+        )
+
+    async def aresult(
+        self,
+        timeout: float | None = None,
+        raise_on_failure: bool = True,
+    ) -> R:
+        if not self._final_state:
+            await self.wait_async(timeout=timeout)
+            if not self._final_state:
+                raise TimeoutError(
+                    f"Task run {self.task_run_id} did not complete within {timeout} seconds"
+                )
+
+        return await self._final_state.result(
+            raise_on_failure=raise_on_failure, fetch=True
+        )
+
+    def add_done_callback(self, fn: Callable[[PrefectFuture[R]], None]) -> None:
+        if self._final_state:
+            fn(self)
+            return
+        FlowRunWaiter.instance()
+        with get_client(sync_client=True) as client:
+            flow_run = client.read_flow_run(flow_run_id=self._flow_run_id)
+            if flow_run.state and flow_run.state.is_final():
+                self._final_state = flow_run.state
+                fn(self)
+                return
+            FlowRunWaiter.add_done_callback(self._flow_run_id, partial(fn, self))
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, PrefectFlowRunFuture):
+            return False
+        return self.flow_run_id == other.flow_run_id
+
+    def __hash__(self) -> int:
+        return hash(self.flow_run_id)
+
+
 class PrefectFutureList(list[PrefectFuture[R]], Iterator[PrefectFuture[R]]):
     """
     A list of Prefect futures.
@@ -297,7 +467,7 @@ class PrefectFutureList(list[PrefectFuture[R]], Iterator[PrefectFuture[R]]):
     in the list to complete and to retrieve the results of all task runs.
     """
 
-    def wait(self, timeout: Optional[float] = None) -> None:
+    def wait(self, timeout: float | None = None) -> None:
         """
         Wait for all futures in the list to complete.
 
@@ -309,7 +479,7 @@ class PrefectFutureList(list[PrefectFuture[R]], Iterator[PrefectFuture[R]]):
 
     def result(
         self: Self,
-        timeout: Optional[float] = None,
+        timeout: float | None = None,
         raise_on_failure: bool = True,
     ) -> list[R]:
         """
@@ -341,7 +511,7 @@ class PrefectFutureList(list[PrefectFuture[R]], Iterator[PrefectFuture[R]]):
 
 
 def as_completed(
-    futures: list[PrefectFuture[R]], timeout: Optional[float] = None
+    futures: list[PrefectFuture[R]], timeout: float | None = None
 ) -> Generator[PrefectFuture[R], None]:
     unique_futures: set[PrefectFuture[R]] = set(futures)
     total_futures = len(unique_futures)
@@ -392,7 +562,7 @@ class DoneAndNotDoneFutures(NamedTuple, Generic[R]):
 
 
 def wait(
-    futures: list[PrefectFuture[R]], timeout: Optional[float] = None
+    futures: list[PrefectFuture[R]], timeout: float | None = None
 ) -> DoneAndNotDoneFutures[R]:
     """
     Wait for the futures in the given sequence to complete.
@@ -442,8 +612,8 @@ def wait(
 
 
 def resolve_futures_to_states(
-    expr: Union[PrefectFuture[R], Any],
-) -> Union[State, Any]:
+    expr: PrefectFuture[R] | Any,
+) -> PrefectFuture[R] | Any:
     """
     Given a Python built-in collection, recursively find `PrefectFutures` and build a
     new collection with the same structure with futures resolved to their final states.
@@ -454,8 +624,8 @@ def resolve_futures_to_states(
     futures: set[PrefectFuture[R]] = set()
 
     def _collect_futures(
-        futures: set[PrefectFuture[R]], expr: Any, context: Any
-    ) -> Union[PrefectFuture[R], Any]:
+        futures: set[PrefectFuture[R]], expr: Any | PrefectFuture[R], context: Any
+    ) -> Any | PrefectFuture[R]:
         # Expressions inside quotes should not be traversed
         if isinstance(context.get("annotation"), quote):
             raise StopVisiting()