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

prefect/states.py

@@ -1,3 +1,4 @@
+import asyncio
 import datetime
 import sys
 import traceback
@@ -13,28 +14,31 @@ from typing_extensions import TypeGuard
 
 from prefect.client.schemas import State as State
 from prefect.client.schemas import StateDetails, StateType
-from prefect.deprecated.data_documents import (
-    DataDocument,
-    result_from_state_with_data_document,
-)
 from prefect.exceptions import (
     CancelledRun,
     CrashedRun,
     FailedRun,
+    MissingContextError,
     MissingResult,
     PausedRun,
     TerminationSignal,
     UnfinishedRun,
 )
-from prefect.results import BaseResult, R, ResultFactory
+from prefect.logging.loggers import get_logger, get_run_logger
+from prefect.results import BaseResult, R, ResultStore
 from prefect.settings import PREFECT_ASYNC_FETCH_STATE_RESULT
 from prefect.utilities.annotations import BaseAnnotation
 from prefect.utilities.asyncutils import in_async_main_thread, sync_compatible
 from prefect.utilities.collections import ensure_iterable
 
+logger = get_logger("states")
+
 
 def get_state_result(
-    state: State[R], raise_on_failure: bool = True, fetch: Optional[bool] = None
+    state: State[R],
+    raise_on_failure: bool = True,
+    fetch: Optional[bool] = None,
+    retry_result_failure: bool = True,
 ) -> R:
     """
     Get the result from a state.
@@ -47,7 +51,6 @@ def get_state_result(
     ):
         # Fetch defaults to `True` for sync users or async users who have opted in
         fetch = True
-
     if not fetch:
         if fetch is None and in_async_main_thread():
             warnings.warn(
@@ -60,19 +63,53 @@ def get_state_result(
                 DeprecationWarning,
                 stacklevel=2,
             )
-        # Backwards compatibility
-        if isinstance(state.data, DataDocument):
-            return result_from_state_with_data_document(
-                state, raise_on_failure=raise_on_failure
-            )
-        else:
-            return state.data
+
+        return state.data
     else:
-        return _get_state_result(state, raise_on_failure=raise_on_failure)
+        return _get_state_result(
+            state,
+            raise_on_failure=raise_on_failure,
+            retry_result_failure=retry_result_failure,
+        )
+
+
+RESULT_READ_MAXIMUM_ATTEMPTS = 10
+RESULT_READ_RETRY_DELAY = 0.25
+
+
+async def _get_state_result_data_with_retries(
+    state: State[R], retry_result_failure: bool = True
+) -> R:
+    # Results may be written asynchronously, possibly after their corresponding
+    # state has been written and events have been emitted, so we should give some
+    # grace here about missing results.  The exception below could come in the form
+    # of a missing file, a short read, or other types of errors depending on the
+    # result storage backend.
+    if retry_result_failure is False:
+        max_attempts = 1
+    else:
+        max_attempts = RESULT_READ_MAXIMUM_ATTEMPTS
+
+    for i in range(1, max_attempts + 1):
+        try:
+            return await state.data.get()
+        except Exception as e:
+            if i == max_attempts:
+                raise
+            logger.debug(
+                "Exception %r while reading result, retry %s/%s in %ss...",
+                e,
+                i,
+                max_attempts,
+                RESULT_READ_RETRY_DELAY,
+            )
+            await asyncio.sleep(RESULT_READ_RETRY_DELAY)
 
 
 @sync_compatible
-async def _get_state_result(state: State[R], raise_on_failure: bool) -> R:
+async def _get_state_result(
+    state: State[R], raise_on_failure: bool, retry_result_failure: bool = True
+) -> R:
     """
     Internal implementation for `get_state_result` without async backwards compatibility
     """
@@ -90,12 +127,11 @@ async def _get_state_result(state: State[R], raise_on_failure: bool) -> R:
     ):
         raise await get_state_exception(state)
 
-    if isinstance(state.data, DataDocument):
-        result = result_from_state_with_data_document(
-            state, raise_on_failure=raise_on_failure
+    if isinstance(state.data, BaseResult):
+        result = await _get_state_result_data_with_retries(
+            state, retry_result_failure=retry_result_failure
         )
-    elif isinstance(state.data, BaseResult):
-        result = await state.data.get()
+
     elif state.data is None:
         if state.is_failed() or state.is_crashed() or state.is_cancelled():
             return await get_state_exception(state)
@@ -131,7 +167,7 @@ def format_exception(exc: BaseException, tb: TracebackType = None) -> str:
 
 async def exception_to_crashed_state(
     exc: BaseException,
-    result_factory: Optional[ResultFactory] = None,
+    result_store: Optional[ResultStore] = None,
 ) -> State:
     """
     Takes an exception that occurs _outside_ of user code and converts it to a
@@ -170,8 +206,8 @@ async def exception_to_crashed_state(
             f" {format_exception(exc)}"
         )
 
-    if result_factory:
-        data = await result_factory.create_result(exc)
+    if result_store:
+        data = await result_store.create_result(exc)
     else:
         # Attach the exception for local usage, will not be available when retrieved
         # from the API
@@ -182,12 +218,18 @@ async def exception_to_crashed_state(
 
 async def exception_to_failed_state(
     exc: Optional[BaseException] = None,
-    result_factory: Optional[ResultFactory] = None,
+    result_store: Optional[ResultStore] = None,
+    write_result: bool = False,
     **kwargs,
 ) -> State:
     """
     Convenience function for creating `Failed` states from exceptions
     """
+    try:
+        local_logger = get_run_logger()
+    except MissingContextError:
+        local_logger = logger
+
     if not exc:
         _, exc, _ = sys.exc_info()
         if exc is None:
@@ -197,8 +239,16 @@ async def exception_to_failed_state(
     else:
         pass
 
-    if result_factory:
-        data = await result_factory.create_result(exc)
+    if result_store:
+        data = await result_store.create_result(exc)
+        if write_result:
+            try:
+                await data.write()
+            except Exception as exc:
+                local_logger.warning(
+                    "Failed to write result: %s Execution will continue, but the result has not been written",
+                    exc,
+                )
     else:
         # Attach the exception for local usage, will not be available when retrieved
         # from the API
@@ -212,10 +262,19 @@ async def exception_to_failed_state(
     #       excluded from messages for now
     message = existing_message + format_exception(exc)
 
-    return Failed(data=data, message=message, **kwargs)
+    state = Failed(data=data, message=message, **kwargs)
+    state.state_details.retriable = False
+
+    return state
 
 
-async def return_value_to_state(retval: R, result_factory: ResultFactory) -> State[R]:
+async def return_value_to_state(
+    retval: R,
+    result_store: ResultStore,
+    key: Optional[str] = None,
+    expiration: Optional[datetime.datetime] = None,
+    write_result: bool = False,
+) -> State[R]:
     """
     Given a return value from a user's function, create a `State` the run should
     be placed in.
@@ -236,28 +295,39 @@ async def return_value_to_state(retval: R, result_factory: ResultFactory) -> Sta
     Callers should resolve all futures into states before passing return values to this
     function.
     """
+    try:
+        local_logger = get_run_logger()
+    except MissingContextError:
+        local_logger = logger
 
     if (
-        is_state(retval)
+        isinstance(retval, State)
         # Check for manual creation
         and not retval.state_details.flow_run_id
         and not retval.state_details.task_run_id
     ):
         state = retval
-
-        # Do not modify states with data documents attached; backwards compatibility
-        if isinstance(state.data, DataDocument):
-            return state
-
-        # Unless the user has already constructed a result explicitly, use the factory
+        # Unless the user has already constructed a result explicitly, use the store
         # to update the data to the correct type
         if not isinstance(state.data, BaseResult):
-            state.data = await result_factory.create_result(state.data)
-
+            result = await result_store.create_result(
+                state.data,
+                key=key,
+                expiration=expiration,
+            )
+            if write_result:
+                try:
+                    await result.write()
+                except Exception as exc:
+                    local_logger.warning(
+                        "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
+                        exc,
+                    )
+            state.data = result
         return state
 
     # Determine a new state from the aggregate of contained states
-    if is_state(retval) or is_state_iterable(retval):
+    if isinstance(retval, State) or is_state_iterable(retval):
         states = StateGroup(ensure_iterable(retval))
 
         # Determine the new state type
@@ -289,10 +359,23 @@ async def return_value_to_state(retval: R, result_factory: ResultFactory) -> Sta
         # TODO: We may actually want to set the data to a `StateGroup` object and just
         #       allow it to be unpacked into a tuple and such so users can interact with
         #       it
+        result = await result_store.create_result(
+            retval,
+            key=key,
+            expiration=expiration,
+        )
+        if write_result:
+            try:
+                await result.write()
+            except Exception as exc:
+                local_logger.warning(
+                    "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
+                    exc,
+                )
         return State(
             type=new_state_type,
             message=message,
-            data=await result_factory.create_result(retval),
+            data=result,
         )
 
     # Generators aren't portable, implicitly convert them to a list.
@@ -302,7 +385,23 @@ async def return_value_to_state(retval: R, result_factory: ResultFactory) -> Sta
         data = retval
 
     # Otherwise, they just gave data and this is a completed retval
-    return Completed(data=await result_factory.create_result(data))
+    if isinstance(data, BaseResult):
+        return Completed(data=data)
+    else:
+        result = await result_store.create_result(
+            data,
+            key=key,
+            expiration=expiration,
+        )
+        if write_result:
+            try:
+                await result.write()
+            except Exception as exc:
+                local_logger.warning(
+                    "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
+                    exc,
+                )
+        return Completed(data=result)
 
 
 @sync_compatible
@@ -342,7 +441,7 @@ async def get_state_exception(state: State) -> BaseException:
         raise ValueError(f"Expected failed or crashed state got {state!r}.")
 
     if isinstance(state.data, BaseResult):
-        result = await state.data.get()
+        result = await _get_state_result_data_with_retries(state)
     elif state.data is None:
         result = None
     else:
@@ -360,7 +459,7 @@ async def get_state_exception(state: State) -> BaseException:
     elif isinstance(result, str):
         return wrapper(result)
 
-    elif is_state(result):
+    elif isinstance(result, State):
         # Return the exception from the inner state
         return await get_state_exception(result)
 
@@ -392,23 +491,6 @@ async def raise_state_exception(state: State) -> None:
     raise await get_state_exception(state)
 
 
-def is_state(obj: Any) -> TypeGuard[State]:
-    """
-    Check if the given object is a state instance
-    """
-    # We may want to narrow this to client-side state types but for now this provides
-    # backwards compatibility
-    try:
-        from prefect.server.schemas.states import State as State_
-
-        classes_ = (State, State_)
-    except ImportError:
-        classes_ = State
-
-    # return isinstance(obj, (State, State_))
-    return isinstance(obj, classes_)
-
-
 def is_state_iterable(obj: Any) -> TypeGuard[Iterable[State]]:
     """
     Check if a the given object is an iterable of states types
@@ -427,7 +509,7 @@ def is_state_iterable(obj: Any) -> TypeGuard[Iterable[State]]:
         and isinstance(obj, (list, set, tuple))
         and obj
     ):
-        return all([is_state(o) for o in obj])
+        return all([isinstance(o, State) for o in obj])
     else:
         return False
 
@@ -495,7 +577,7 @@ def Scheduled(
     Returns:
         State: a Scheduled state
     """
-    state_details = StateDetails.parse_obj(kwargs.pop("state_details", {}))
+    state_details = StateDetails.model_validate(kwargs.pop("state_details", {}))
     if scheduled_time is None:
         scheduled_time = pendulum.now("UTC")
     elif state_details.scheduled_time:
@@ -581,7 +663,7 @@ def Paused(
     Returns:
         State: a Paused state
     """
-    state_details = StateDetails.parse_obj(kwargs.pop("state_details", {}))
+    state_details = StateDetails.model_validate(kwargs.pop("state_details", {}))
 
     if state_details.pause_timeout:
         raise ValueError("An extra pause timeout was provided in state_details")
@@ -642,6 +724,21 @@ def AwaitingRetry(
     )
 
 
+def AwaitingConcurrencySlot(
+    cls: Type[State[R]] = State,
+    scheduled_time: Optional[datetime.datetime] = None,
+    **kwargs: Any,
+) -> State[R]:
+    """Convenience function for creating `AwaitingConcurrencySlot` states.
+
+    Returns:
+        State: a AwaitingConcurrencySlot state
+    """
+    return Scheduled(
+        cls=cls, scheduled_time=scheduled_time, name="AwaitingConcurrencySlot", **kwargs
+    )
+
+
 def Retrying(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     """Convenience function for creating `Retrying` states.