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

prefect/task_runners.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import abc
 import asyncio
 import sys
@@ -14,7 +16,6 @@ from typing import (
     Iterable,
     List,
     Optional,
-    Set,
     overload,
 )
 
@@ -39,12 +40,14 @@ from prefect.utilities.callables import (
 from prefect.utilities.collections import isiterable
 
 if TYPE_CHECKING:
+    import logging
+
     from prefect.tasks import Task
 
 P = ParamSpec("P")
 T = TypeVar("T")
 R = TypeVar("R")
-F = TypeVar("F", bound=PrefectFuture, default=PrefectConcurrentFuture)
+F = TypeVar("F", bound=PrefectFuture[Any], default=PrefectConcurrentFuture[Any])
 
 
 class TaskRunner(abc.ABC, Generic[F]):
@@ -60,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", "")
 
@@ -73,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",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+        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]:
         """
@@ -138,9 +151,9 @@ class TaskRunner(abc.ABC, Generic[F]):
         # Ensure that any parameters in kwargs are expanded before this check
         parameters = explode_variadic_parameter(task.fn, parameters)
 
-        iterable_parameters = {}
-        static_parameters = {}
-        annotated_parameters = {}
+        iterable_parameters: dict[str, Any] = {}
+        static_parameters: dict[str, Any] = {}
+        annotated_parameters: dict[str, Any] = {}
         for key, val in parameters.items():
             if isinstance(val, (allow_failure, quote)):
                 # Unwrap annotated parameters to determine if they are iterable
@@ -172,9 +185,9 @@ class TaskRunner(abc.ABC, Generic[F]):
 
         map_length = list(lengths)[0]
 
-        futures: List[PrefectFuture] = []
+        futures: List[PrefectFuture[Any]] = []
         for i in range(map_length):
-            call_parameters = {
+            call_parameters: dict[str, Any] = {
                 key: value[i] for key, value in iterable_parameters.items()
             }
             call_parameters.update(
@@ -204,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")
 
@@ -212,12 +225,12 @@ class TaskRunner(abc.ABC, Generic[F]):
         self._started = True
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    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]):
+class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[R]]):
     def __init__(self, max_workers: Optional[int] = None):
         super().__init__()
         self._executor: Optional[ThreadPoolExecutor] = None
@@ -228,16 +241,16 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
         )
         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
     def submit(
         self,
         task: "Task[P, Coroutine[Any, Any, R]]",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
     ) -> PrefectConcurrentFuture[R]:
         ...
 
@@ -245,19 +258,19 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
     def submit(
         self,
         task: "Task[Any, R]",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
     ) -> PrefectConcurrentFuture[R]:
         ...
 
     def submit(
         self,
-        task: "Task",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
-    ):
+        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,
+    ) -> PrefectConcurrentFuture[R]:
         """
         Submit a task to the task run engine running in a separate thread.
 
@@ -289,7 +302,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
         else:
             self.logger.debug(f"Submitting task {task.name} to thread pool executor...")
 
-        submit_kwargs = dict(
+        submit_kwargs: dict[str, Any] = dict(
             task=task,
             task_run_id=task_run_id,
             parameters=parameters,
@@ -322,8 +335,8 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
     def map(
         self,
         task: "Task[P, Coroutine[Any, Any, R]]",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
     ) -> PrefectFutureList[PrefectConcurrentFuture[R]]:
         ...
 
@@ -331,20 +344,20 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
     def map(
         self,
         task: "Task[Any, R]",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
     ) -> PrefectFutureList[PrefectConcurrentFuture[R]]:
         ...
 
     def map(
         self,
-        task: "Task",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-    ):
+        task: "Task[P, R]",
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+    ) -> 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")
@@ -353,12 +366,12 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
             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, exc_value, traceback):
+    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)
@@ -375,20 +388,20 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
 ConcurrentTaskRunner = ThreadPoolTaskRunner
 
 
-class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture]):
+class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture[R]]):
     def __init__(self):
         super().__init__()
 
-    def duplicate(self) -> "PrefectTaskRunner":
+    def duplicate(self) -> "PrefectTaskRunner[R]":
         return type(self)()
 
     @overload
     def submit(
         self,
         task: "Task[P, Coroutine[Any, Any, R]]",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
     ) -> PrefectDistributedFuture[R]:
         ...
 
@@ -396,19 +409,19 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture]):
     def submit(
         self,
         task: "Task[Any, R]",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
     ) -> PrefectDistributedFuture[R]:
         ...
 
     def submit(
         self,
-        task: "Task",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-        dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
-    ):
+        task: "Task[P, R]",
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
+    ) -> PrefectDistributedFuture[R]:
         """
         Submit a task to the task run engine running in a separate thread.
 
@@ -443,8 +456,8 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture]):
     def map(
         self,
         task: "Task[P, Coroutine[Any, Any, R]]",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
     ) -> PrefectFutureList[PrefectDistributedFuture[R]]:
         ...
 
@@ -452,15 +465,15 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture]):
     def map(
         self,
         task: "Task[Any, R]",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
     ) -> PrefectFutureList[PrefectDistributedFuture[R]]:
         ...
 
     def map(
         self,
-        task: "Task",
-        parameters: Dict[str, Any],
-        wait_for: Optional[Iterable[PrefectFuture]] = None,
-    ):
+        task: "Task[P, R]",
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+    ) -> PrefectFutureList[PrefectDistributedFuture[R]]:
         return super().map(task, parameters, wait_for)

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,12 +18,12 @@ 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
 from prefect._internal.concurrency.api import create_call, from_sync
-from prefect.cache_policies import DEFAULT, NONE
+from prefect.cache_policies import DEFAULT, NO_CACHE
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.objects import TaskRun
 from prefect.client.subscriptions import Subscription
@@ -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:
@@ -93,14 +98,14 @@ class TaskWorker:
                 if not isinstance(t, Task):
                     continue
 
-            if t.cache_policy in [None, NONE, NotSet]:
+            if t.cache_policy in [None, NO_CACHE, NotSet]:
                 self.tasks.append(
                     t.with_options(persist_result=True, cache_policy=DEFAULT)
                 )
             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]