prefect-client 3.1.11__py3-none-any.whl → 3.1.13__py3-none-any.whl

prefect/states.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import datetime
 import sys
@@ -6,7 +8,7 @@ import uuid
 import warnings
 from collections import Counter
 from types import GeneratorType, TracebackType
-from typing import Any, Dict, Iterable, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Type
 
 import anyio
 import httpx
@@ -28,17 +30,19 @@ from prefect.exceptions import (
     UnfinishedRun,
 )
 from prefect.logging.loggers import get_logger, get_run_logger
-from prefect.results import (
-    R,
-    ResultRecord,
-    ResultRecordMetadata,
-    ResultStore,
-)
 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")
+if TYPE_CHECKING:
+    import logging
+
+    from prefect.results import (
+        R,
+        ResultStore,
+    )
+
+logger: "logging.Logger" = get_logger("states")
 
 
 @deprecated.deprecated_parameter(
@@ -49,11 +53,11 @@ logger = get_logger("states")
     help="Please ensure you are awaiting the call to `result()` when calling in an async context.",
 )
 def get_state_result(
-    state: State[R],
+    state: "State[R]",
     raise_on_failure: bool = True,
     fetch: bool = True,
     retry_result_failure: bool = True,
-) -> R:
+) -> "R":
     """
     Get the result from a state.
 
@@ -86,13 +90,18 @@ RESULT_READ_RETRY_DELAY = 0.25
 
 
 async def _get_state_result_data_with_retries(
-    state: State[R], retry_result_failure: bool = True
-) -> R:
+    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.
+    from prefect.results import (
+        ResultRecord,
+        ResultRecordMetadata,
+    )
+
     if retry_result_failure is False:
         max_attempts = 1
     else:
@@ -120,11 +129,16 @@ async def _get_state_result_data_with_retries(
 
 @sync_compatible
 async def _get_state_result(
-    state: State[R], raise_on_failure: bool, retry_result_failure: bool = True
-) -> R:
+    state: "State[R]", raise_on_failure: bool, retry_result_failure: bool = True
+) -> "R":
     """
     Internal implementation for `get_state_result` without async backwards compatibility
     """
+    from prefect.results import (
+        ResultRecord,
+        ResultRecordMetadata,
+    )
+
     if state.is_paused():
         # Paused states are not truly terminal and do not have results associated with them
         raise PausedRun("Run is paused, its result is not available.", state=state)
@@ -181,7 +195,7 @@ def format_exception(exc: BaseException, tb: TracebackType = None) -> str:
 
 async def exception_to_crashed_state(
     exc: BaseException,
-    result_store: Optional[ResultStore] = None,
+    result_store: Optional["ResultStore"] = None,
 ) -> State:
     """
     Takes an exception that occurs _outside_ of user code and converts it to a
@@ -233,10 +247,10 @@ async def exception_to_crashed_state(
 
 async def exception_to_failed_state(
     exc: Optional[BaseException] = None,
-    result_store: Optional[ResultStore] = None,
+    result_store: Optional["ResultStore"] = None,
     write_result: bool = False,
-    **kwargs,
-) -> State:
+    **kwargs: Any,
+) -> State[BaseException]:
     """
     Convenience function for creating `Failed` states from exceptions
     """
@@ -285,12 +299,12 @@ async def exception_to_failed_state(
 
 
 async def return_value_to_state(
-    retval: R,
-    result_store: ResultStore,
+    retval: "R",
+    result_store: "ResultStore",
     key: Optional[str] = None,
     expiration: Optional[datetime.datetime] = None,
     write_result: bool = False,
-) -> State[R]:
+) -> "State[R]":
     """
     Given a return value from a user's function, create a `State` the run should
     be placed in.
@@ -311,6 +325,11 @@ async def return_value_to_state(
     Callers should resolve all futures into states before passing return values to this
     function.
     """
+    from prefect.results import (
+        ResultRecord,
+        ResultRecordMetadata,
+    )
+
     try:
         local_logger = get_run_logger()
     except MissingContextError:
@@ -443,6 +462,10 @@ async def get_state_exception(state: State) -> BaseException:
         - `CrashedRun` if the state type is CRASHED.
         - `CancelledRun` if the state type is CANCELLED.
     """
+    from prefect.results import (
+        ResultRecord,
+        ResultRecordMetadata,
+    )
 
     if state.is_failed():
         wrapper = FailedRun
@@ -534,17 +557,17 @@ def is_state_iterable(obj: Any) -> TypeGuard[Iterable[State]]:
 
 
 class StateGroup:
-    def __init__(self, states: Iterable[State]) -> None:
-        self.states = states
-        self.type_counts = self._get_type_counts(states)
-        self.total_count = len(states)
-        self.cancelled_count = self.type_counts[StateType.CANCELLED]
-        self.final_count = sum(state.is_final() for state in states)
-        self.not_final_count = self.total_count - self.final_count
-        self.paused_count = self.type_counts[StateType.PAUSED]
+    def __init__(self, states: list[State]) -> None:
+        self.states: list[State] = states
+        self.type_counts: dict[StateType, int] = self._get_type_counts(states)
+        self.total_count: int = len(states)
+        self.cancelled_count: int = self.type_counts[StateType.CANCELLED]
+        self.final_count: int = sum(state.is_final() for state in states)
+        self.not_final_count: int = self.total_count - self.final_count
+        self.paused_count: int = self.type_counts[StateType.PAUSED]
 
     @property
-    def fail_count(self):
+    def fail_count(self) -> int:
         return self.type_counts[StateType.FAILED] + self.type_counts[StateType.CRASHED]
 
     def all_completed(self) -> bool:
@@ -586,7 +609,7 @@ class StateGroup:
         return f"StateGroup<{self.counts_message()}>"
 
 
-def _traced(cls: Type[State[R]], **kwargs: Any) -> State[R]:
+def _traced(cls: Type["State[R]"], **kwargs: Any) -> "State[R]":
     state_details = StateDetails.model_validate(kwargs.pop("state_details", {}))
 
     carrier = {}
@@ -597,10 +620,10 @@ def _traced(cls: Type[State[R]], **kwargs: Any) -> State[R]:
 
 
 def Scheduled(
-    cls: Type[State[R]] = State,
+    cls: Type["State[R]"] = State,
     scheduled_time: Optional[datetime.datetime] = None,
     **kwargs: Any,
-) -> State[R]:
+) -> "State[R]":
     """Convenience function for creating `Scheduled` states.
 
     Returns:
@@ -616,7 +639,7 @@ def Scheduled(
     return _traced(cls, type=StateType.SCHEDULED, state_details=state_details, **kwargs)
 
 
-def Completed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
+def Completed(cls: Type["State[R]"] = State, **kwargs: Any) -> "State[R]":
     """Convenience function for creating `Completed` states.
 
     Returns:
@@ -626,7 +649,7 @@ def Completed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     return _traced(cls, type=StateType.COMPLETED, **kwargs)
 
 
-def Running(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
+def Running(cls: Type["State[R]"] = State, **kwargs: Any) -> "State[R]":
     """Convenience function for creating `Running` states.
 
     Returns:
@@ -635,7 +658,7 @@ def Running(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     return _traced(cls, type=StateType.RUNNING, **kwargs)
 
 
-def Failed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
+def Failed(cls: Type["State[R]"] = State, **kwargs: Any) -> "State[R]":
     """Convenience function for creating `Failed` states.
 
     Returns:
@@ -644,7 +667,7 @@ def Failed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     return _traced(cls, type=StateType.FAILED, **kwargs)
 
 
-def Crashed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
+def Crashed(cls: Type["State[R]"] = State, **kwargs: Any) -> "State[R]":
     """Convenience function for creating `Crashed` states.
 
     Returns:
@@ -653,7 +676,7 @@ def Crashed(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     return _traced(cls, type=StateType.CRASHED, **kwargs)
 
 
-def Cancelling(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
+def Cancelling(cls: Type["State[R]"] = State, **kwargs: Any) -> "State[R]":
     """Convenience function for creating `Cancelling` states.
 
     Returns:
@@ -662,7 +685,7 @@ def Cancelling(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     return _traced(cls, type=StateType.CANCELLING, **kwargs)
 
 
-def Cancelled(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
+def Cancelled(cls: Type["State[R]"] = State, **kwargs: Any) -> "State[R]":
     """Convenience function for creating `Cancelled` states.
 
     Returns:
@@ -671,7 +694,7 @@ def Cancelled(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
     return _traced(cls, type=StateType.CANCELLED, **kwargs)
 
 
-def Pending(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
+def Pending(cls: Type["State[R]"] = State, **kwargs: Any) -> "State[R]":
     """Convenience function for creating `Pending` states.
 
     Returns:
@@ -681,13 +704,13 @@ def Pending(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
 
 
 def Paused(
-    cls: Type[State[R]] = State,
+    cls: Type["State[R]"] = State,
     timeout_seconds: Optional[int] = None,
     pause_expiration_time: Optional[datetime.datetime] = None,
     reschedule: bool = False,
     pause_key: Optional[str] = None,
     **kwargs: Any,
-) -> State[R]:
+) -> "State[R]":
     """Convenience function for creating `Paused` states.
 
     Returns:
@@ -717,12 +740,12 @@ def Paused(
 
 
 def Suspended(
-    cls: Type[State[R]] = State,
+    cls: Type["State[R]"] = State,
     timeout_seconds: Optional[int] = None,
     pause_expiration_time: Optional[datetime.datetime] = None,
     pause_key: Optional[str] = None,
     **kwargs: Any,
-):
+) -> "State[R]":
     """Convenience function for creating `Suspended` states.
 
     Returns:
@@ -740,10 +763,10 @@ def Suspended(
 
 
 def AwaitingRetry(
-    cls: Type[State[R]] = State,
+    cls: Type["State[R]"] = State,
     scheduled_time: Optional[datetime.datetime] = None,
     **kwargs: Any,
-) -> State[R]:
+) -> "State[R]":
     """Convenience function for creating `AwaitingRetry` states.
 
     Returns:
@@ -755,10 +778,10 @@ def AwaitingRetry(
 
 
 def AwaitingConcurrencySlot(
-    cls: Type[State[R]] = State,
+    cls: Type["State[R]"] = State,
     scheduled_time: Optional[datetime.datetime] = None,
     **kwargs: Any,
-) -> State[R]:
+) -> "State[R]":
     """Convenience function for creating `AwaitingConcurrencySlot` states.
 
     Returns:
@@ -769,7 +792,7 @@ def AwaitingConcurrencySlot(
     )
 
 
-def Retrying(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
+def Retrying(cls: Type["State[R]"] = State, **kwargs: Any) -> "State[R]":
     """Convenience function for creating `Retrying` states.
 
     Returns:
@@ -779,10 +802,10 @@ def Retrying(cls: Type[State[R]] = State, **kwargs: Any) -> State[R]:
 
 
 def Late(
-    cls: Type[State[R]] = State,
+    cls: Type["State[R]"] = State,
     scheduled_time: Optional[datetime.datetime] = None,
     **kwargs: Any,
-) -> State[R]:
+) -> "State[R]":
     """Convenience function for creating `Late` states.
 
     Returns:

prefect/task_engine.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import inspect
 import logging
@@ -28,8 +30,9 @@ from uuid import UUID
 import anyio
 import pendulum
 from opentelemetry import trace
-from typing_extensions import ParamSpec
+from typing_extensions import ParamSpec, Self
 
+from prefect.cache_policies import CachePolicy
 from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client
 from prefect.client.schemas import TaskRun
 from prefect.client.schemas.objects import State, TaskRunInput
@@ -55,7 +58,7 @@ from prefect.exceptions import (
 from prefect.logging.loggers import get_logger, patch_print, task_run_logger
 from prefect.results import (
     ResultRecord,
-    _format_user_supplied_storage_key,
+    _format_user_supplied_storage_key,  # type: ignore[reportPrivateUsage]
     get_result_store,
     should_persist_result,
 )
@@ -115,20 +118,20 @@ class BaseTaskRunEngine(Generic[P, R]):
     # holds the return value from the user code
     _return_value: Union[R, Type[NotSet]] = NotSet
     # holds the exception raised by the user code, if any
-    _raised: Union[Exception, Type[NotSet]] = NotSet
+    _raised: Union[Exception, BaseException, Type[NotSet]] = NotSet
     _initial_run_context: Optional[TaskRunContext] = None
     _is_started: bool = False
     _task_name_set: bool = False
     _last_event: Optional[PrefectEvent] = None
     _telemetry: RunTelemetry = field(default_factory=RunTelemetry)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.parameters is None:
             self.parameters = {}
 
     @property
     def state(self) -> State:
-        if not self.task_run:
+        if not self.task_run or not self.task_run.state:
             raise ValueError("Task run is not set")
         return self.task_run.state
 
@@ -142,8 +145,8 @@ class BaseTaskRunEngine(Generic[P, R]):
         return False
 
     def compute_transaction_key(self) -> Optional[str]:
-        key = None
-        if self.task.cache_policy:
+        key: Optional[str] = None
+        if self.task.cache_policy and isinstance(self.task.cache_policy, CachePolicy):
             flow_run_context = FlowRunContext.get()
             task_run_context = TaskRunContext.get()
 
@@ -153,10 +156,12 @@ class BaseTaskRunEngine(Generic[P, R]):
                 parameters = None
 
             try:
+                if not task_run_context:
+                    raise ValueError("Task run context is not set")
                 key = self.task.cache_policy.compute_key(
                     task_ctx=task_run_context,
-                    inputs=self.parameters,
-                    flow_parameters=parameters,
+                    inputs=self.parameters or {},
+                    flow_parameters=parameters or {},
                 )
             except Exception:
                 self.logger.exception(
@@ -169,7 +174,7 @@ class BaseTaskRunEngine(Generic[P, R]):
 
     def _resolve_parameters(self):
         if not self.parameters:
-            return {}
+            return None
 
         resolved_parameters = {}
         for parameter, value in self.parameters.items():
@@ -227,10 +232,8 @@ class BaseTaskRunEngine(Generic[P, R]):
         if self.task_run and self.task_run.start_time and not self.task_run.end_time:
             self.task_run.end_time = state.timestamp
 
-            if self.task_run.state.is_running():
-                self.task_run.total_run_time += (
-                    state.timestamp - self.task_run.state.timestamp
-                )
+            if self.state.is_running():
+                self.task_run.total_run_time += state.timestamp - self.state.timestamp
 
     def is_running(self) -> bool:
         """Whether or not the engine is currently running a task."""
@@ -238,7 +241,7 @@ class BaseTaskRunEngine(Generic[P, R]):
             return False
         return task_run.state.is_running() or task_run.state.is_scheduled()
 
-    def log_finished_message(self):
+    def log_finished_message(self) -> None:
         if not self.task_run:
             return
 
@@ -294,6 +297,7 @@ class BaseTaskRunEngine(Generic[P, R]):
 
 @dataclass
 class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
+    task_run: Optional[TaskRun] = None
     _client: Optional[SyncPrefectClient] = None
 
     @property
@@ -336,7 +340,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
             return False
 
-    def call_hooks(self, state: Optional[State] = None):
+    def call_hooks(self, state: Optional[State] = None) -> None:
         if state is None:
             state = self.state
         task = self.task
@@ -371,7 +375,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 self.logger.info(f"Hook {hook_name!r} finished running successfully")
 
-    def begin_run(self):
+    def begin_run(self) -> None:
         try:
             self._resolve_parameters()
             self._set_custom_task_run_name()
@@ -390,6 +394,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
         new_state = Running()
 
+        assert self.task_run is not None, "Task run is not set"
         self.task_run.start_time = new_state.timestamp
 
         flow_run_context = FlowRunContext.get()
@@ -406,7 +411,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         # result reference that no longer exists
         if state.is_completed():
             try:
-                state.result(retry_result_failure=False, _sync=True)
+                state.result(retry_result_failure=False, _sync=True)  # type: ignore[reportCallIssue]
             except Exception:
                 state = self.set_state(new_state, force=True)
 
@@ -422,7 +427,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             time.sleep(interval)
             state = self.set_state(new_state)
 
-    def set_state(self, state: State, force: bool = False) -> State:
+    def set_state(self, state: State[R], force: bool = False) -> State[R]:
         last_state = self.state
         if not self.task_run:
             raise ValueError("Task run is not set")
@@ -537,7 +542,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                "Task run failed with exception: %r - " "Retry %s/%s will start %s",
+                "Task run failed with exception: %r - Retry %s/%s will start %s",
                 exc,
                 self.retries + 1,
                 self.task.retries,
@@ -545,7 +550,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
 
             self.set_state(new_state, force=True)
-            self.retries = self.retries + 1
+            self.retries: int = self.retries + 1
             return True
         elif self.retries >= self.task.retries:
             self.logger.error(
@@ -639,7 +644,9 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             stack.enter_context(ConcurrencyContextV1())
             stack.enter_context(ConcurrencyContext())
 
-            self.logger = task_run_logger(task_run=self.task_run, task=self.task)  # type: ignore
+            self.logger: "logging.Logger" = task_run_logger(
+                task_run=self.task_run, task=self.task
+            )  # type: ignore
 
             yield
 
@@ -648,7 +655,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self,
         task_run_id: Optional[UUID] = None,
         dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
-    ) -> Generator["SyncTaskRunEngine", Any, Any]:
+    ) -> Generator[Self, Any, Any]:
         """
         Enters a client context and creates a task run if needed.
         """
@@ -718,7 +725,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                     self._is_started = False
                     self._client = None
 
-    async def wait_until_ready(self):
+    async def wait_until_ready(self) -> None:
         """Waits until the scheduled time (if its the future), then enters Running."""
         if scheduled_time := self.state.state_details.scheduled_time:
             sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
@@ -825,6 +832,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
 @dataclass
 class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
+    task_run: TaskRun | None = None
     _client: Optional[PrefectClient] = None
 
     @property
@@ -866,7 +874,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
             return False
 
-    async def call_hooks(self, state: Optional[State] = None):
+    async def call_hooks(self, state: Optional[State] = None) -> None:
         if state is None:
             state = self.state
         task = self.task
@@ -901,7 +909,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 self.logger.info(f"Hook {hook_name!r} finished running successfully")
 
-    async def begin_run(self):
+    async def begin_run(self) -> None:
         try:
             self._resolve_parameters()
             self._set_custom_task_run_name()
@@ -1067,7 +1075,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                "Task run failed with exception: %r - " "Retry %s/%s will start %s",
+                "Task run failed with exception: %r - Retry %s/%s will start %s",
                 exc,
                 self.retries + 1,
                 self.task.retries,
@@ -1075,7 +1083,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
 
             await self.set_state(new_state, force=True)
-            self.retries = self.retries + 1
+            self.retries: int = self.retries + 1
             return True
         elif self.retries >= self.task.retries:
             self.logger.error(
@@ -1169,7 +1177,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             )
             stack.enter_context(ConcurrencyContext())
 
-            self.logger = task_run_logger(task_run=self.task_run, task=self.task)  # type: ignore
+            self.logger: "logging.Logger" = task_run_logger(
+                task_run=self.task_run, task=self.task
+            )  # type: ignore
 
             yield
 
@@ -1178,7 +1188,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self,
         task_run_id: Optional[UUID] = None,
         dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
-    ) -> AsyncGenerator["AsyncTaskRunEngine", Any]:
+    ) -> AsyncGenerator[Self, Any]:
         """
         Enters a client context and creates a task run if needed.
         """
@@ -1246,7 +1256,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                     self._is_started = False
                     self._client = None
 
-    async def wait_until_ready(self):
+    async def wait_until_ready(self) -> None:
         """Waits until the scheduled time (if its the future), then enters Running."""
         if scheduled_time := self.state.state_details.scheduled_time:
             sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
@@ -1341,7 +1351,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         if transaction.is_committed():
             result = transaction.read()
         else:
-            if self.task_run.tags:
+            if self.task_run and self.task_run.tags:
                 # Acquire a concurrency slot for each tag, but only if a limit
                 # matching the tag already exists.
                 async with aconcurrency(list(self.task_run.tags), self.task_run.id):
@@ -1546,7 +1556,7 @@ def run_task(
     Returns:
         The result of the task run
     """
-    kwargs = dict(
+    kwargs: dict[str, Any] = dict(
         task=task,
         task_run_id=task_run_id,
         task_run=task_run,