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

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,

prefect/task_runners.py

@@ -40,6 +40,8 @@ from prefect.utilities.callables import (
 from prefect.utilities.collections import isiterable
 
 if TYPE_CHECKING:
+    import logging
+
     from prefect.tasks import Task
 
 P = ParamSpec("P")
@@ -61,11 +63,11 @@ class TaskRunner(abc.ABC, Generic[F]):
     """
 
     def __init__(self):
-        self.logger = get_logger(f"task_runner.{self.name}")
+        self.logger: "logging.Logger" = get_logger(f"task_runner.{self.name}")
         self._started = False
 
     @property
-    def name(self):
+    def name(self) -> str:
         """The name of this task runner"""
         return type(self).__name__.lower().replace("taskrunner", "")
 
@@ -74,32 +76,42 @@ class TaskRunner(abc.ABC, Generic[F]):
         """Return a new instance of this task runner with the same configuration."""
         ...
 
+    @overload
     @abc.abstractmethod
     def submit(
         self,
-        task: "Task[P, R]",
+        task: "Task[P, Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
         dependencies: dict[str, set[TaskRunInput]] | None = None,
     ) -> F:
-        """
-        Submit a task to the task run engine.
+        ...
 
-        Args:
-            task: The task to submit.
-            parameters: The parameters to use when running the task.
-            wait_for: A list of futures that the task depends on.
+    @overload
+    @abc.abstractmethod
+    def submit(
+        self,
+        task: "Task[Any, R]",
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
+    ) -> F:
+        ...
 
-        Returns:
-            A future object that can be used to wait for the task to complete and
-            retrieve the result.
-        """
+    @abc.abstractmethod
+    def submit(
+        self,
+        task: "Task[P, R]",
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
+    ) -> F:
         ...
 
     def map(
         self,
         task: "Task[P, R]",
-        parameters: dict[str, Any],
+        parameters: dict[str, Any | unmapped[Any] | allow_failure[Any]],
         wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     ) -> PrefectFutureList[F]:
         """
@@ -205,7 +217,7 @@ class TaskRunner(abc.ABC, Generic[F]):
 
         return PrefectFutureList(futures)
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         if self._started:
             raise RuntimeError("This task runner is already started")
 
@@ -213,12 +225,12 @@ class TaskRunner(abc.ABC, Generic[F]):
         self._started = True
         return self
 
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.logger.debug("Stopping task runner")
         self._started = False
 
 
-class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
+class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[R]]):
     def __init__(self, max_workers: Optional[int] = None):
         super().__init__()
         self._executor: Optional[ThreadPoolExecutor] = None
@@ -229,7 +241,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
         )
         self._cancel_events: Dict[uuid.UUID, threading.Event] = {}
 
-    def duplicate(self) -> "ThreadPoolTaskRunner":
+    def duplicate(self) -> "ThreadPoolTaskRunner[R]":
         return type(self)(max_workers=self._max_workers)
 
     @overload
@@ -254,7 +266,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
 
     def submit(
         self,
-        task: "Task[P, R]",
+        task: "Task[P, R | Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
         dependencies: dict[str, set[TaskRunInput]] | None = None,
@@ -345,7 +357,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
     ) -> PrefectFutureList[PrefectConcurrentFuture[R]]:
         return super().map(task, parameters, wait_for)
 
-    def cancel_all(self):
+    def cancel_all(self) -> None:
         for event in self._cancel_events.values():
             event.set()
             self.logger.debug("Set cancel event")
@@ -354,12 +366,12 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
             self._executor.shutdown(cancel_futures=True)
             self._executor = None
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         super().__enter__()
         self._executor = ThreadPoolExecutor(max_workers=self._max_workers)
         return self
 
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.cancel_all()
         if self._executor is not None:
             self._executor.shutdown(cancel_futures=True)
@@ -380,7 +392,7 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture[R]]):
     def __init__(self):
         super().__init__()
 
-    def duplicate(self) -> "PrefectTaskRunner":
+    def duplicate(self) -> "PrefectTaskRunner[R]":
         return type(self)()
 
     @overload

prefect/task_runs.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import asyncio
 import atexit
 import threading
 import uuid
-from typing import Callable, Dict, Optional
+from typing import TYPE_CHECKING, Callable, Dict, Optional
 
 import anyio
 from cachetools import TTLCache
@@ -15,6 +17,9 @@ from prefect.events.clients import get_events_subscriber
 from prefect.events.filters import EventFilter, EventNameFilter
 from prefect.logging.loggers import get_logger
 
+if TYPE_CHECKING:
+    import logging
+
 
 class TaskRunWaiter:
     """
@@ -68,19 +73,19 @@ class TaskRunWaiter:
     _instance_lock = threading.Lock()
 
     def __init__(self):
-        self.logger = get_logger("TaskRunWaiter")
-        self._consumer_task: Optional[asyncio.Task] = None
+        self.logger: "logging.Logger" = get_logger("TaskRunWaiter")
+        self._consumer_task: "asyncio.Task[None] | None" = None
         self._observed_completed_task_runs: TTLCache[uuid.UUID, bool] = TTLCache(
             maxsize=10000, ttl=600
         )
         self._completion_events: Dict[uuid.UUID, asyncio.Event] = {}
-        self._completion_callbacks: Dict[uuid.UUID, Callable] = {}
+        self._completion_callbacks: Dict[uuid.UUID, Callable[[], None]] = {}
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._observed_completed_task_runs_lock = threading.Lock()
         self._completion_events_lock = threading.Lock()
         self._started = False
 
-    def start(self):
+    def start(self) -> None:
         """
         Start the TaskRunWaiter service.
         """
@@ -89,10 +94,12 @@ class TaskRunWaiter:
         self.logger.debug("Starting TaskRunWaiter")
         loop_thread = get_global_loop()
 
-        if not asyncio.get_running_loop() == loop_thread._loop:
+        if not asyncio.get_running_loop() == loop_thread.loop:
             raise RuntimeError("TaskRunWaiter must run on the global loop thread.")
 
-        self._loop = loop_thread._loop
+        self._loop = loop_thread.loop
+        if TYPE_CHECKING:
+            assert self._loop is not None
 
         consumer_started = asyncio.Event()
         self._consumer_task = self._loop.create_task(
@@ -141,7 +148,7 @@ class TaskRunWaiter:
                 except Exception as exc:
                     self.logger.error(f"Error processing event: {exc}")
 
-    def stop(self):
+    def stop(self) -> None:
         """
         Stop the TaskRunWaiter service.
         """
@@ -155,7 +162,7 @@ class TaskRunWaiter:
     @classmethod
     async def wait_for_task_run(
         cls, task_run_id: uuid.UUID, timeout: Optional[float] = None
-    ):
+    ) -> None:
         """
         Wait for a task run to finish.
 
@@ -199,7 +206,9 @@ class TaskRunWaiter:
                 instance._completion_events.pop(task_run_id, None)
 
     @classmethod
-    def add_done_callback(cls, task_run_id: uuid.UUID, callback):
+    def add_done_callback(
+        cls, task_run_id: uuid.UUID, callback: Callable[[], None]
+    ) -> None:
         """
         Add a callback to be called when a task run finishes.
 
@@ -219,7 +228,7 @@ class TaskRunWaiter:
             instance._completion_callbacks[task_run_id] = callback
 
     @classmethod
-    def instance(cls):
+    def instance(cls) -> Self:
         """
         Get the singleton instance of TaskRunWaiter.
         """

prefect/task_worker.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import inspect
 import os
@@ -16,7 +18,7 @@ import pendulum
 import uvicorn
 from exceptiongroup import BaseExceptionGroup  # novermin
 from fastapi import FastAPI
-from typing_extensions import ParamSpec, TypeVar
+from typing_extensions import ParamSpec, Self, TypeVar
 from websockets.exceptions import InvalidStatusCode
 
 from prefect import Task
@@ -42,7 +44,10 @@ from prefect.utilities.processutils import (
 from prefect.utilities.services import start_client_metrics_server
 from prefect.utilities.urls import url_for
 
-logger = get_logger("task_worker")
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger("task_worker")
 
 P = ParamSpec("P")
 R = TypeVar("R", infer_variance=True)
@@ -85,7 +90,7 @@ class TaskWorker:
     def __init__(
         self,
         *tasks: Task[P, R],
-        limit: Optional[int] = 10,
+        limit: int | None = 10,
     ):
         self.tasks: list["Task[..., Any]"] = []
         for t in tasks:
@@ -100,7 +105,7 @@ class TaskWorker:
             else:
                 self.tasks.append(t.with_options(persist_result=True))
 
-        self.task_keys = set(t.task_key for t in tasks if isinstance(t, Task))  # pyright: ignore[reportUnnecessaryIsInstance]
+        self.task_keys: set[str] = set(t.task_key for t in tasks if isinstance(t, Task))  # pyright: ignore[reportUnnecessaryIsInstance]
 
         self._started_at: Optional[pendulum.DateTime] = None
         self.stopping: bool = False
@@ -154,7 +159,7 @@ class TaskWorker:
     def available_tasks(self) -> Optional[int]:
         return int(self._limiter.available_tokens) if self._limiter else None
 
-    def handle_sigterm(self, signum: int, frame: object):
+    def handle_sigterm(self, signum: int, frame: object) -> None:
         """
         Shuts down the task worker when a SIGTERM is received.
         """
@@ -355,14 +360,14 @@ class TaskWorker:
             )
             await asyncio.wrap_future(future)
 
-    async def execute_task_run(self, task_run: TaskRun):
+    async def execute_task_run(self, task_run: TaskRun) -> None:
         """Execute a task run in the task worker."""
         async with self if not self.started else asyncnullcontext():
             token_acquired = await self._acquire_token(task_run.id)
             if token_acquired:
                 await self._safe_submit_scheduled_task_run(task_run)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Self:
         logger.debug("Starting task worker...")
 
         if self._client._closed:  # pyright: ignore[reportPrivateUsage]