prefect-client 2.19.4__py3-none-any.whl → 3.0.0rc2__py3-none-any.whl

prefect/futures.py

@@ -1,375 +1,242 @@
-"""
-Futures represent the execution of a task and allow retrieval of the task run's state.
-
-This module contains the definition for futures as well as utilities for resolving
-futures in nested data structures.
-"""
-import asyncio
-import warnings
+import abc
+import concurrent.futures
+import inspect
+import uuid
 from functools import partial
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Awaitable,
-    Callable,
-    Generic,
-    Optional,
-    Set,
-    TypeVar,
-    Union,
-    cast,
-    overload,
-)
-from uuid import UUID
-
-import anyio
-
-from prefect._internal.concurrency.api import create_call, from_async, from_sync
-from prefect._internal.concurrency.event_loop import run_coroutine_in_loop_from_async
-from prefect.client.orchestration import PrefectClient
-from prefect.client.utilities import inject_client
-from prefect.states import State
+from typing import Any, Generic, Optional, Set, Union, cast
+
+from typing_extensions import TypeVar
+
+from prefect.client.orchestration import get_client
+from prefect.client.schemas.objects import TaskRun
+from prefect.exceptions import ObjectNotFound
+from prefect.logging.loggers import get_logger
+from prefect.states import Pending, State
+from prefect.task_runs import TaskRunWaiter
 from prefect.utilities.annotations import quote
-from prefect.utilities.asyncutils import A, Async, Sync, sync
+from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import StopVisiting, visit_collection
 
-if TYPE_CHECKING:
-    from prefect.task_runners import BaseTaskRunner
+F = TypeVar("F")
 
+logger = get_logger(__name__)
 
-R = TypeVar("R")
 
-
-class PrefectFuture(Generic[R, A]):
+class PrefectFuture(abc.ABC):
     """
-    Represents the result of a computation happening in a task runner.
-
-    When tasks are called, they are submitted to a task runner which creates a future
-    for access to the state and result of the task.
-
-    Examples:
-        Define a task that returns a string
-
-        >>> from prefect import flow, task
-        >>> @task
-        >>> def my_task() -> str:
-        >>>     return "hello"
-
-        Calls of this task in a flow will return a future
-
-        >>> @flow
-        >>> def my_flow():
-        >>>     future = my_task.submit()  # PrefectFuture[str, Sync] includes result type
-        >>>     future.task_run.id  # UUID for the task run
-
-        Wait for the task to complete
-
-        >>> @flow
-        >>> def my_flow():
-        >>>     future = my_task.submit()
-        >>>     final_state = future.wait()
-
-        Wait N seconds for the task to complete
-
-        >>> @flow
-        >>> def my_flow():
-        >>>     future = my_task.submit()
-        >>>     final_state = future.wait(0.1)
-        >>>     if final_state:
-        >>>         ... # Task done
-        >>>     else:
-        >>>         ... # Task not done yet
-
-        Wait for a task to complete and retrieve its result
-
-        >>> @flow
-        >>> def my_flow():
-        >>>     future = my_task.submit()
-        >>>     result = future.result()
-        >>>     assert result == "hello"
-
-        Wait N seconds for a task to complete and retrieve its result
-
-        >>> @flow
-        >>> def my_flow():
-        >>>     future = my_task.submit()
-        >>>     result = future.result(timeout=5)
-        >>>     assert result == "hello"
-
-        Retrieve the state of a task without waiting for completion
-
-        >>> @flow
-        >>> def my_flow():
-        >>>     future = my_task.submit()
-        >>>     state = future.get_state()
+    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.
     """
 
-    def __init__(
-        self,
-        name: str,
-        key: UUID,
-        task_runner: "BaseTaskRunner",
-        asynchronous: A = True,
-        _final_state: State[R] = None,  # Exposed for testing
-    ) -> None:
-        self.key = key
-        self.name = name
-        self.asynchronous = asynchronous
-        self.task_run = None
-        self._final_state = _final_state
-        self._exception: Optional[Exception] = None
-        self._task_runner = task_runner
-        self._submitted = anyio.Event()
-
-        self._loop = asyncio.get_running_loop()
-
-    @overload
-    def wait(
-        self: "PrefectFuture[R, Async]", timeout: None = None
-    ) -> Awaitable[State[R]]:
-        ...
-
-    @overload
-    def wait(self: "PrefectFuture[R, Sync]", timeout: None = None) -> State[R]:
-        ...
+    def __init__(self, task_run_id: uuid.UUID):
+        self._task_run_id = task_run_id
+        self._final_state = None
 
-    @overload
-    def wait(
-        self: "PrefectFuture[R, Async]", timeout: float
-    ) -> Awaitable[Optional[State[R]]]:
-        ...
+    @property
+    def task_run_id(self) -> uuid.UUID:
+        """The ID of the task run associated with this future"""
+        return self._task_run_id
 
-    @overload
-    def wait(self: "PrefectFuture[R, Sync]", timeout: float) -> Optional[State[R]]:
+    @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 = cast(TaskRun, 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()
+
+    @abc.abstractmethod
+    def wait(self, timeout: Optional[float] = None) -> None:
         ...
-
-    def wait(self, timeout=None):
         """
-        Wait for the run to finish and return the final state
+        Wait for the task run to complete. 
 
-        If the timeout is reached before the run reaches a final state,
-        `None` is returned.
-        """
-        wait = create_call(self._wait, timeout=timeout)
-        if self.asynchronous:
-            return from_async.call_soon_in_loop_thread(wait).aresult()
-        else:
-            # type checking cannot handle the overloaded timeout passing
-            return from_sync.call_soon_in_loop_thread(wait).result()  # type: ignore
+        If the task run has already completed, this method will return immediately.
 
-    @overload
-    async def _wait(self, timeout: None = None) -> State[R]:
-        ...
-
-    @overload
-    async def _wait(self, timeout: float) -> Optional[State[R]]:
-        ...
-
-    async def _wait(self, timeout=None):
-        """
-        Async implementation for `wait`
+        Args:
+            - timeout: The maximum number of seconds to wait for the task run to complete.
+              If the task run has not completed after the timeout has elapsed, this method will return.
         """
-        await self._wait_for_submission()
-
-        if self._final_state:
-            return self._final_state
 
-        self._final_state = await self._task_runner.wait(self.key, timeout)
-        return self._final_state
-
-    @overload
+    @abc.abstractmethod
     def result(
-        self: "PrefectFuture[R, Sync]",
-        timeout: float = None,
+        self,
+        timeout: Optional[float] = None,
         raise_on_failure: bool = True,
-    ) -> R:
-        ...
-
-    @overload
-    def result(
-        self: "PrefectFuture[R, Sync]",
-        timeout: float = None,
-        raise_on_failure: bool = False,
-    ) -> Union[R, Exception]:
+    ) -> Any:
         ...
+        """
+        Get the result of the task run associated with this future.
 
-    @overload
-    def result(
-        self: "PrefectFuture[R, Async]",
-        timeout: float = None,
-        raise_on_failure: bool = True,
-    ) -> Awaitable[R]:
-        ...
+        If the task run has not completed, this method will wait for the task run to complete.
 
-    @overload
-    def result(
-        self: "PrefectFuture[R, Async]",
-        timeout: float = None,
-        raise_on_failure: bool = False,
-    ) -> Awaitable[Union[R, Exception]]:
-        ...
+        Args:
+            - timeout: The maximum number of seconds to wait for the task run to complete.
+            If the task run has not completed after the timeout has elapsed, this method will return.
+            - raise_on_failure: If `True`, an exception will be raised if the task run fails.
 
-    def result(self, timeout: float = None, raise_on_failure: bool = True):
+        Returns:
+            The result of the task run.
         """
-        Wait for the run to finish and return the final state.
 
-        If the timeout is reached before the run reaches a final state, a `TimeoutError`
-        will be raised.
 
-        If `raise_on_failure` is `True` and the task run failed, the task run's
-        exception will be raised.
-        """
-        result = create_call(
-            self._result, timeout=timeout, raise_on_failure=raise_on_failure
-        )
-        if self.asynchronous:
-            return from_async.call_soon_in_loop_thread(result).aresult()
-        else:
-            return from_sync.call_soon_in_loop_thread(result).result()
+class PrefectWrappedFuture(PrefectFuture, abc.ABC, Generic[F]):
+    """
+    A Prefect future that wraps another future object.
+    """
 
-    async def _result(self, timeout: float = None, raise_on_failure: bool = True):
-        """
-        Async implementation of `result`
-        """
-        final_state = await self._wait(timeout=timeout)
-        if not final_state:
-            raise TimeoutError("Call timed out before task finished.")
-        return await final_state.result(raise_on_failure=raise_on_failure, fetch=True)
-
-    @overload
-    def get_state(
-        self: "PrefectFuture[R, Async]", client: PrefectClient = None
-    ) -> Awaitable[State[R]]:
-        ...
+    def __init__(self, task_run_id: uuid.UUID, wrapped_future: F):
+        self._wrapped_future = wrapped_future
+        super().__init__(task_run_id)
 
-    @overload
-    def get_state(
-        self: "PrefectFuture[R, Sync]", client: PrefectClient = None
-    ) -> State[R]:
-        ...
+    @property
+    def wrapped_future(self) -> F:
+        """The underlying future object wrapped by this Prefect future"""
+        return self._wrapped_future
 
-    def get_state(self, client: PrefectClient = None):
-        """
-        Get the current state of the task run.
-        """
-        if self.asynchronous:
-            return cast(Awaitable[State[R]], self._get_state(client=client))
-        else:
-            return cast(State[R], sync(self._get_state, client=client))
 
-    @inject_client
-    async def _get_state(self, client: PrefectClient = None) -> State[R]:
-        assert client is not None  # always injected
+class PrefectConcurrentFuture(PrefectWrappedFuture[concurrent.futures.Future]):
+    """
+    A Prefect future that wraps a concurrent.futures.Future. This future is used
+    when the task run is submitted to a ThreadPoolExecutor.
+    """
 
-        # We must wait for the task run id to be populated
-        await self._wait_for_submission()
+    def wait(self, timeout: Optional[float] = None) -> None:
+        try:
+            result = self._wrapped_future.result(timeout=timeout)
+        except concurrent.futures.TimeoutError:
+            return
+        if isinstance(result, State):
+            self._final_state = result
 
-        task_run = await client.read_task_run(self.task_run.id)
+    def result(
+        self,
+        timeout: Optional[float] = None,
+        raise_on_failure: bool = True,
+    ) -> Any:
+        if not self._final_state:
+            try:
+                future_result = self._wrapped_future.result(timeout=timeout)
+            except concurrent.futures.TimeoutError as exc:
+                raise TimeoutError(
+                    f"Task run {self.task_run_id} did not complete within {timeout} seconds"
+                ) from exc
+
+            if isinstance(future_result, State):
+                self._final_state = future_result
+            else:
+                return future_result
+
+        _result = self._final_state.result(
+            raise_on_failure=raise_on_failure, fetch=True
+        )
+        # state.result is a `sync_compatible` function that may or may not return an awaitable
+        # depending on whether the parent frame is sync or not
+        if inspect.isawaitable(_result):
+            _result = run_coro_as_sync(_result)
+        return _result
 
-        if not task_run:
-            raise RuntimeError("Future has no associated task run in the server.")
 
-        # Update the task run reference
-        self.task_run = task_run
-        return task_run.state
+class PrefectDistributedFuture(PrefectFuture):
+    """
+    Represents the result of a computation happening anywhere.
 
-    async def _wait_for_submission(self):
-        await run_coroutine_in_loop_from_async(self._loop, self._submitted.wait())
+    This class is typically used to interact with the result of a task run
+    scheduled to run in a Prefect task worker but can be used to interact with
+    any task run scheduled in Prefect's API.
+    """
 
-    def __hash__(self) -> int:
-        return hash(self.key)
+    def wait(self, timeout: Optional[float] = None) -> None:
+        return run_coro_as_sync(self.wait_async(timeout=timeout))
 
-    def __repr__(self) -> str:
-        return f"PrefectFuture({self.name!r})"
+    async def wait_async(self, timeout: Optional[float] = None):
+        if self._final_state:
+            logger.debug(
+                "Final state already set for %s. Returning...", self.task_run_id
+            )
+            return
+
+        # Read task run to see if it is still running
+        async with get_client() as client:
+            task_run = await client.read_task_run(task_run_id=self._task_run_id)
+            if task_run.state.is_final():
+                logger.debug(
+                    "Task run %s already finished. Returning...",
+                    self.task_run_id,
+                )
+                self._final_state = task_run.state
+                return
+
+            # If still running, wait for a completed event from the server
+            logger.debug(
+                "Waiting for completed event for task run %s...",
+                self.task_run_id,
+            )
+            await TaskRunWaiter.wait_for_task_run(self._task_run_id, timeout=timeout)
+            task_run = await client.read_task_run(task_run_id=self._task_run_id)
+            if task_run.state.is_final():
+                self._final_state = task_run.state
+            return
 
-    def __bool__(self) -> bool:
-        warnings.warn(
-            (
-                "A 'PrefectFuture' from a task call was cast to a boolean; "
-                "did you mean to check the result of the task instead? "
-                "e.g. `if my_task().result(): ...`"
-            ),
-            stacklevel=2,
+    def result(
+        self,
+        timeout: Optional[float] = None,
+        raise_on_failure: bool = True,
+    ) -> Any:
+        return run_coro_as_sync(
+            self.result_async(timeout=timeout, raise_on_failure=raise_on_failure)
         )
-        return True
-
 
-def _collect_futures(futures, expr, context):
-    # Expressions inside quotes should not be traversed
-    if isinstance(context.get("annotation"), quote):
-        raise StopVisiting()
-
-    if isinstance(expr, PrefectFuture):
-        futures.add(expr)
+    async def result_async(
+        self,
+        timeout: Optional[float] = None,
+        raise_on_failure: bool = True,
+    ):
+        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
+        )
 
-    return expr
+    def __eq__(self, other):
+        if not isinstance(other, PrefectDistributedFuture):
+            return False
+        return self.task_run_id == other.task_run_id
 
 
-async def resolve_futures_to_data(
-    expr: Union[PrefectFuture[R, Any], Any],
-    raise_on_failure: bool = True,
-) -> Union[R, Any]:
+def resolve_futures_to_states(
+    expr: Union[PrefectFuture, Any],
+) -> Union[State, Any]:
     """
     Given a Python built-in collection, recursively find `PrefectFutures` and build a
-    new collection with the same structure with futures resolved to their results.
-    Resolving futures to their results may wait for execution to complete and require
-    communication with the API.
+    new collection with the same structure with futures resolved to their final states.
+    Resolving futures to their final states may wait for execution to complete.
 
     Unsupported object types will be returned without modification.
     """
     futures: Set[PrefectFuture] = set()
 
-    maybe_expr = visit_collection(
-        expr,
-        visit_fn=partial(_collect_futures, futures),
-        return_data=False,
-        context={},
-    )
-    if maybe_expr is not None:
-        expr = maybe_expr
-
-    # Get results
-    results = await asyncio.gather(
-        *[
-            # We must wait for the future in the thread it was created in
-            from_async.call_soon_in_loop_thread(
-                create_call(future._result, raise_on_failure=raise_on_failure)
-            ).aresult()
-            for future in futures
-        ]
-    )
-
-    results_by_future = dict(zip(futures, results))
-
-    def replace_futures_with_results(expr, context):
-        # Expressions inside quotes should not be modified
+    def _collect_futures(futures, expr, context):
+        # Expressions inside quotes should not be traversed
         if isinstance(context.get("annotation"), quote):
             raise StopVisiting()
 
         if isinstance(expr, PrefectFuture):
-            return results_by_future[expr]
-        else:
-            return expr
-
-    return visit_collection(
-        expr,
-        visit_fn=replace_futures_with_results,
-        return_data=True,
-        context={},
-    )
-
-
-async def resolve_futures_to_states(
-    expr: Union[PrefectFuture[R, Any], Any],
-) -> Union[State[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.
-    Resolving futures to their final states may wait for execution to complete.
+            futures.add(expr)
 
-    Unsupported object types will be returned without modification.
-    """
-    futures: Set[PrefectFuture] = set()
+        return expr
 
     visit_collection(
         expr,
@@ -379,13 +246,10 @@ async def resolve_futures_to_states(
     )
 
     # Get final states for each future
-    states = await asyncio.gather(
-        *[
-            # We must wait for the future in the thread it was created in
-            from_async.call_soon_in_loop_thread(create_call(future._wait)).aresult()
-            for future in futures
-        ]
-    )
+    states = []
+    for future in futures:
+        future.wait()
+        states.append(future.state)
 
     states_by_future = dict(zip(futures, states))
 
@@ -405,25 +269,3 @@ async def resolve_futures_to_states(
         return_data=True,
         context={},
     )
-
-
-def call_repr(__fn: Callable, *args: Any, **kwargs: Any) -> str:
-    """
-    Generate a repr for a function call as "fn_name(arg_value, kwarg_name=kwarg_value)"
-    """
-
-    name = __fn.__name__
-
-    # TODO: If this computation is concerningly expensive, we can iterate checking the
-    #       length at each arg or avoid calling `repr` on args with large amounts of
-    #       data
-    call_args = ", ".join(
-        [repr(arg) for arg in args]
-        + [f"{key}={repr(val)}" for key, val in kwargs.items()]
-    )
-
-    # Enforce a maximum length
-    if len(call_args) > 100:
-        call_args = call_args[:100] + "..."
-
-    return f"{name}({call_args})"

prefect/infrastructure/__init__.py

@@ -1,27 +0,0 @@
-from prefect.infrastructure.base import Infrastructure, InfrastructureResult
-from prefect.infrastructure.container import DockerContainer, DockerContainerResult
-from prefect.infrastructure.kubernetes import (
-    KubernetesClusterConfig,
-    KubernetesImagePullPolicy,
-    KubernetesJob,
-    KubernetesJobResult,
-    KubernetesManifest,
-    KubernetesRestartPolicy,
-)
-from prefect.infrastructure.process import Process, ProcessResult
-
-# Declare API
-__all__ = [
-    "DockerContainer",
-    "DockerContainerResult",
-    "Infrastructure",
-    "InfrastructureResult",
-    "KubernetesClusterConfig",
-    "KubernetesImagePullPolicy",
-    "KubernetesJob",
-    "KubernetesJobResult",
-    "KubernetesManifest",
-    "KubernetesRestartPolicy",
-    "Process",
-    "ProcessResult",
-]

prefect/infrastructure/provisioners/__init__.py

@@ -1,12 +1,14 @@
-from typing import Any, Dict, Optional, Protocol, Type
+from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Type
 
-from prefect.client.orchestration import PrefectClient
 from prefect.infrastructure.provisioners.modal import ModalPushProvisioner
 from .cloud_run import CloudRunPushProvisioner
 from .container_instance import ContainerInstancePushProvisioner
 from .ecs import ElasticContainerServicePushProvisioner
 import rich.console
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 _provisioners = {
     "cloud-run:push": CloudRunPushProvisioner,
     "cloud-run-v2:push": CloudRunPushProvisioner,
@@ -29,7 +31,7 @@ class Provisioner(Protocol):
         self,
         work_pool_name: str,
         base_job_template: Dict[str, Any],
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ) -> Dict[str, Any]:
         ...