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

prefect/_internal/concurrency/services.py

@@ -1,15 +1,14 @@
 import abc
 import asyncio
-import atexit
 import concurrent.futures
 import contextlib
 import logging
 import queue
-import sys
 import threading
-from typing import Awaitable, Dict, Generic, List, Optional, Type, TypeVar, Union
+from collections.abc import AsyncGenerator, Awaitable, Coroutine, Generator, Hashable
+from typing import TYPE_CHECKING, Any, Generic, NoReturn, Optional, Union, cast
 
-from typing_extensions import Self
+from typing_extensions import Self, TypeVar, TypeVarTuple, Unpack
 
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.api import create_call, from_sync
@@ -18,17 +17,19 @@ from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect._internal.concurrency.threads import WorkerThread, get_global_loop
 
 T = TypeVar("T")
+Ts = TypeVarTuple("Ts")
+R = TypeVar("R", infer_variance=True)
 
 
-class QueueService(abc.ABC, Generic[T]):
-    _instances: Dict[int, Self] = {}
+class _QueueServiceBase(abc.ABC, Generic[T]):
+    _instances: dict[int, Self] = {}
     _instance_lock = threading.Lock()
 
-    def __init__(self, *args) -> None:
-        self._queue: queue.Queue = queue.Queue()
+    def __init__(self, *args: Hashable) -> None:
+        self._queue: queue.Queue[Optional[T]] = queue.Queue()
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._done_event: Optional[asyncio.Event] = None
-        self._task: Optional[asyncio.Task] = None
+        self._task: Optional[asyncio.Task[None]] = None
         self._stopped: bool = False
         self._started: bool = False
         self._key = hash(args)
@@ -41,14 +42,14 @@ class QueueService(abc.ABC, Generic[T]):
         )
         self._logger = logging.getLogger(f"{type(self).__name__}")
 
-    def start(self):
+    def start(self) -> None:
         logger.debug("Starting service %r", self)
         loop_thread = get_global_loop()
 
-        if not asyncio.get_running_loop() == loop_thread._loop:
+        if not asyncio.get_running_loop() == getattr(loop_thread, "_loop"):
             raise RuntimeError("Services must run on the global loop thread.")
 
-        self._loop = loop_thread._loop
+        self._loop = asyncio.get_running_loop()
         self._done_event = asyncio.Event()
         self._task = self._loop.create_task(self._run())
         self._queue_get_thread.start()
@@ -58,23 +59,18 @@ class QueueService(abc.ABC, Generic[T]):
         loop_thread.add_shutdown_call(create_call(self.drain))
 
         # Stop at interpreter exit by default
-        if sys.version_info < (3, 9):
-            atexit.register(self._at_exit)
-        else:
-            # See related issue at https://bugs.python.org/issue42647
-            # Handling items may require spawning a thread and in 3.9  new threads
-            # cannot be spawned after the interpreter finalizes threads which happens
-            #  _before_ the normal `atexit` hook is called resulting in failure to
-            # process items. This is particularly relevant for services which use an
-            # httpx client.
-            from threading import _register_atexit
-
-            _register_atexit(self._at_exit)
-
-    def _at_exit(self):
+        # Handling items may require spawning a thread and in 3.9 new threads
+        # cannot be spawned after the interpreter finalizes threads which
+        #  happens _before_ the normal `atexit` hook is called resulting in
+        # failure to process items. This is particularly relevant for services
+        # which use an httpx client. See related issue at
+        # https://github.com/python/cpython/issues/86813
+        threading._register_atexit(self._at_exit)  # pyright: ignore[reportUnknownVariableType, reportAttributeAccessIssue]
+
+    def _at_exit(self) -> None:
         self.drain(at_exit=True)
 
-    def _stop(self, at_exit: bool = False):
+    def _stop(self, at_exit: bool = False) -> None:
         """
         Stop running this instance.
 
@@ -100,27 +96,11 @@ class QueueService(abc.ABC, Generic[T]):
             # Signal completion to the loop
             self._queue.put_nowait(None)
 
-    def send(self, item: T):
-        """
-        Send an item to this instance of the service.
-        """
-        with self._lock:
-            if self._stopped:
-                raise RuntimeError("Cannot put items in a stopped service instance.")
-
-            logger.debug("Service %r enqueuing item %r", self, item)
-            self._queue.put_nowait(self._prepare_item(item))
-
-    def _prepare_item(self, item: T) -> T:
-        """
-        Prepare an item for submission to the service. This is called before
-        the item is sent to the service.
-
-        The default implementation returns the item unchanged.
-        """
-        return item
+    @abc.abstractmethod
+    def send(self, item: Any) -> Any:
+        raise NotImplementedError
 
-    async def _run(self):
+    async def _run(self) -> None:
         try:
             async with self._lifespan():
                 await self._main_loop()
@@ -142,14 +122,15 @@ class QueueService(abc.ABC, Generic[T]):
             self._queue_get_thread.shutdown()
 
             self._stopped = True
+            assert self._done_event is not None
             self._done_event.set()
 
-    async def _main_loop(self):
+    async def _main_loop(self) -> None:
         last_log_time = 0
         log_interval = 4  # log every 4 seconds
 
         while True:
-            item: T = await self._queue_get_thread.submit(
+            item: Optional[T] = await self._queue_get_thread.submit(
                 create_call(self._queue.get)
             ).aresult()
 
@@ -183,19 +164,17 @@ class QueueService(abc.ABC, Generic[T]):
                 self._queue.task_done()
 
     @abc.abstractmethod
-    async def _handle(self, item: T):
-        """
-        Process an item sent to the service.
-        """
+    async def _handle(self, item: Any) -> Any:
+        raise NotImplementedError
 
     @contextlib.asynccontextmanager
-    async def _lifespan(self):
+    async def _lifespan(self) -> AsyncGenerator[None, Any]:
         """
         Perform any setup and teardown for the service.
         """
         yield
 
-    def _drain(self, at_exit: bool = False) -> concurrent.futures.Future:
+    def _drain(self, at_exit: bool = False) -> concurrent.futures.Future[bool]:
         """
         Internal implementation for `drain`. Returns a future for sync/async interfaces.
         """
@@ -204,15 +183,17 @@ class QueueService(abc.ABC, Generic[T]):
 
         self._stop(at_exit=at_exit)
 
+        assert self._done_event is not None
         if self._done_event.is_set():
-            future = concurrent.futures.Future()
-            future.set_result(None)
+            future: concurrent.futures.Future[bool] = concurrent.futures.Future()
+            future.set_result(False)
             return future
 
-        future = asyncio.run_coroutine_threadsafe(self._done_event.wait(), self._loop)
-        return future
+        assert self._loop is not None
+        task = cast(Coroutine[Any, Any, bool], self._done_event.wait())
+        return asyncio.run_coroutine_threadsafe(task, self._loop)
 
-    def drain(self, at_exit: bool = False) -> None:
+    def drain(self, at_exit: bool = False) -> Union[bool, Awaitable[bool]]:
         """
         Stop this instance of the service and wait for remaining work to be completed.
 
@@ -226,15 +207,24 @@ class QueueService(abc.ABC, Generic[T]):
 
     @classmethod
     def drain_all(
-        cls, timeout: Optional[float] = None, at_exit=True
-    ) -> Union[Awaitable, None]:
+        cls, timeout: Optional[float] = None, at_exit: bool = True
+    ) -> Union[
+        tuple[
+            set[concurrent.futures.Future[bool]], set[concurrent.futures.Future[bool]]
+        ],
+        Coroutine[
+            Any,
+            Any,
+            Optional[tuple[set[asyncio.Future[bool]], set[asyncio.Future[bool]]]],
+        ],
+    ]:
         """
         Stop all instances of the service and wait for all remaining work to be
         completed.
 
         Returns an awaitable if called from an async context.
         """
-        futures = []
+        futures: list[concurrent.futures.Future[bool]] = []
         with cls._instance_lock:
             instances = tuple(cls._instances.values())
 
@@ -242,26 +232,24 @@ class QueueService(abc.ABC, Generic[T]):
                 futures.append(instance._drain(at_exit=at_exit))
 
         if get_running_loop() is not None:
-            return (
-                asyncio.wait(
+            if futures:
+                return asyncio.wait(
                     [asyncio.wrap_future(fut) for fut in futures], timeout=timeout
                 )
-                if futures
-                # `wait` errors if it receives an empty list but we need to return a
-                # coroutine still
-                else asyncio.sleep(0)
-            )
+            # `wait` errors if it receives an empty list but we need to return a
+            # coroutine still
+            return asyncio.sleep(0)
         else:
             return concurrent.futures.wait(futures, timeout=timeout)
 
-    def wait_until_empty(self):
+    def wait_until_empty(self) -> None:
         """
         Wait until the queue is empty and all items have been processed.
         """
         self._queue.join()
 
     @classmethod
-    def instance(cls: Type[Self], *args) -> Self:
+    def instance(cls, *args: Hashable) -> Self:
         """
         Get an instance of the service.
 
@@ -278,7 +266,7 @@ class QueueService(abc.ABC, Generic[T]):
         self._instances.pop(self._key, None)
 
     @classmethod
-    def _new_instance(cls, *args):
+    def _new_instance(cls, *args: Hashable) -> Self:
         """
         Create and start a new instance of the service.
         """
@@ -295,6 +283,87 @@ class QueueService(abc.ABC, Generic[T]):
         return instance
 
 
+class QueueService(_QueueServiceBase[T]):
+    def send(self, item: T) -> None:
+        """
+        Send an item to this instance of the service.
+        """
+        with self._lock:
+            if self._stopped:
+                raise RuntimeError("Cannot put items in a stopped service instance.")
+
+            logger.debug("Service %r enqueuing item %r", self, item)
+            self._queue.put_nowait(self._prepare_item(item))
+
+    def _prepare_item(self, item: T) -> T:
+        """
+        Prepare an item for submission to the service. This is called before
+        the item is sent to the service.
+
+        The default implementation returns the item unchanged.
+        """
+        return item
+
+    @abc.abstractmethod
+    async def _handle(self, item: T) -> None:
+        """
+        Process an item sent to the service.
+        """
+
+
+class FutureQueueService(
+    _QueueServiceBase[tuple[Unpack[Ts], concurrent.futures.Future[R]]]
+):
+    """Queued service that provides a future that is signalled with the acquired result for each item
+
+    If there was a failure acquiring, the future result is set to the exception.
+
+    Type Parameters:
+        Ts: the tuple of types that make up sent arguments
+        R: the type returned for each item once acquired
+
+    """
+
+    async def _handle(
+        self, item: tuple[Unpack[Ts], concurrent.futures.Future[R]]
+    ) -> None:
+        send_item, future = item[:-1], item[-1]
+        try:
+            response = await self.acquire(*send_item)
+        except Exception as exc:
+            # If the request to the increment endpoint fails in a non-standard
+            # way, we need to set the future's result so it'll be re-raised in
+            # the context of the caller.
+            future.set_exception(exc)
+            raise exc
+        else:
+            future.set_result(response)
+
+    @abc.abstractmethod
+    async def acquire(self, *args: Unpack[Ts]) -> R:
+        raise NotImplementedError
+
+    def send(self, item: tuple[Unpack[Ts]]) -> concurrent.futures.Future[R]:
+        with self._lock:
+            if self._stopped:
+                raise RuntimeError("Cannot put items in a stopped service instance.")
+
+            logger.debug("Service %r enqueuing item %r", self, item)
+            future: concurrent.futures.Future[R] = concurrent.futures.Future()
+            self._queue.put_nowait((*self._prepare_item(item), future))
+
+        return future
+
+    def _prepare_item(self, item: tuple[Unpack[Ts]]) -> tuple[Unpack[Ts]]:
+        """
+        Prepare an item for submission to the service. This is called before
+        the item is sent to the service.
+
+        The default implementation returns the item unchanged.
+        """
+        return item
+
+
 class BatchedQueueService(QueueService[T]):
     """
     A queue service that handles a batch of items instead of a single item at a time.
@@ -310,7 +379,7 @@ class BatchedQueueService(QueueService[T]):
         done = False
 
         while not done:
-            batch = []
+            batch: list[T] = []
             batch_size = 0
 
             # Pull items from the queue until we reach the batch size
@@ -359,13 +428,15 @@ class BatchedQueueService(QueueService[T]):
                 )
 
     @abc.abstractmethod
-    async def _handle_batch(self, items: List[T]):
+    async def _handle_batch(self, items: list[T]) -> None:
         """
         Process a batch of items sent to the service.
         """
 
-    async def _handle(self, item: T):
-        assert False, "`_handle` should never be called for batched queue services"
+    async def _handle(self, item: T) -> NoReturn:
+        raise AssertionError(
+            "`_handle` should never be called for batched queue services"
+        )
 
     def _get_size(self, item: T) -> int:
         """
@@ -376,12 +447,15 @@ class BatchedQueueService(QueueService[T]):
 
 
 @contextlib.contextmanager
-def drain_on_exit(service: QueueService):
+def drain_on_exit(service: QueueService[Any]) -> Generator[None, Any, None]:
     yield
     service.drain_all(at_exit=True)
 
 
 @contextlib.asynccontextmanager
-async def drain_on_exit_async(service: QueueService):
+async def drain_on_exit_async(service: QueueService[Any]) -> AsyncGenerator[None, Any]:
     yield
-    await service.drain_all(at_exit=True)
+    drain_all = service.drain_all(at_exit=True)
+    if TYPE_CHECKING:
+        assert not isinstance(drain_all, tuple)
+    await drain_all

prefect/_internal/concurrency/waiters.py

@@ -10,7 +10,7 @@ import inspect
 import queue
 import threading
 from collections import deque
-from collections.abc import Awaitable
+from collections.abc import AsyncGenerator, Awaitable, Generator
 from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
 from weakref import WeakKeyDictionary
 
@@ -24,7 +24,7 @@ from prefect._internal.concurrency.primitives import Event
 T = TypeVar("T")
 
 
-# Waiters are stored in a stack for each thread
+# Waiters are stored in a queue for each thread
 _WAITERS_BY_THREAD: "WeakKeyDictionary[threading.Thread, deque[Waiter[Any]]]" = (
     WeakKeyDictionary()
 )
@@ -49,8 +49,9 @@ class Waiter(Portal, abc.ABC, Generic[T]):
     """
 
     def __init__(self, call: Call[T]) -> None:
-        if not isinstance(call, Call):  # Guard against common mistake
-            raise TypeError(f"Expected call of type `Call`; got {call!r}.")
+        if not TYPE_CHECKING:
+            if not isinstance(call, Call):  # Guard against common mistake
+                raise TypeError(f"Expected call of type `Call`; got {call!r}.")
 
         self._call = call
         self._owner_thread = threading.current_thread()
@@ -107,7 +108,7 @@ class SyncWaiter(Waiter[T]):
         call.set_runner(self)
         return call
 
-    def _handle_waiting_callbacks(self):
+    def _handle_waiting_callbacks(self) -> None:
         logger.debug("Waiter %r watching for callbacks", self)
         while True:
             callback = self._queue.get()
@@ -121,7 +122,7 @@ class SyncWaiter(Waiter[T]):
             del callback
 
     @contextlib.contextmanager
-    def _handle_done_callbacks(self):
+    def _handle_done_callbacks(self) -> Generator[None, Any, None]:
         try:
             yield
         finally:
@@ -195,10 +196,13 @@ class AsyncWaiter(Waiter[T]):
             call_soon_in_loop(self._loop, self._queue.put_nowait, call)
         self._early_submissions = []
 
-    async def _handle_waiting_callbacks(self):
+    async def _handle_waiting_callbacks(self) -> None:
         logger.debug("Waiter %r watching for callbacks", self)
         tasks: list[Awaitable[None]] = []
 
+        if TYPE_CHECKING:
+            assert self._queue is not None
+
         try:
             while True:
                 callback = await self._queue.get()
@@ -221,7 +225,7 @@ class AsyncWaiter(Waiter[T]):
             self._done_waiting = True
 
     @contextlib.asynccontextmanager
-    async def _handle_done_callbacks(self):
+    async def _handle_done_callbacks(self) -> AsyncGenerator[None, Any]:
         try:
             yield
         finally:
@@ -244,7 +248,7 @@ class AsyncWaiter(Waiter[T]):
         else:
             self._done_callbacks.append(callback)
 
-    def _signal_stop_waiting(self):
+    def _signal_stop_waiting(self) -> None:
         # Only send a `None` to the queue if the waiter is still blocked reading from
         # the queue. Otherwise, it's possible that the event loop is stopped.
         if not self._done_waiting:

prefect/_internal/pydantic/annotations/pendulum.py

@@ -3,39 +3,39 @@ This file contains compat code to handle pendulum.DateTime objects during jsonsc
 generation and validation.
 """
 
-import typing as t
+from typing import Annotated, Any, Union
 
 import pendulum
 from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
 from pydantic.json_schema import JsonSchemaValue
 from pydantic_core import core_schema
-from typing_extensions import Annotated
 
 
 class _PendulumDateTimeAnnotation:
-    _pendulum_type: t.Type[
-        t.Union[pendulum.DateTime, pendulum.Date, pendulum.Time]
+    _pendulum_type: type[
+        Union[pendulum.DateTime, pendulum.Date, pendulum.Time, pendulum.Duration]
     ] = pendulum.DateTime
 
     _pendulum_types_to_schemas = {
         pendulum.DateTime: core_schema.datetime_schema(),
         pendulum.Date: core_schema.date_schema(),
+        pendulum.Time: core_schema.time_schema(),
         pendulum.Duration: core_schema.timedelta_schema(),
     }
 
     @classmethod
     def __get_pydantic_core_schema__(
         cls,
-        _source_type: t.Any,
+        _source_type: Any,
         _handler: GetCoreSchemaHandler,
     ) -> core_schema.CoreSchema:
         def validate_from_str(
             value: str,
-        ) -> t.Union[pendulum.DateTime, pendulum.Date, pendulum.Time]:
+        ) -> Union[pendulum.DateTime, pendulum.Date, pendulum.Time, pendulum.Duration]:
             return pendulum.parse(value)
 
         def to_str(
-            value: t.Union[pendulum.DateTime, pendulum.Date, pendulum.Time],
+            value: Union[pendulum.DateTime, pendulum.Date, pendulum.Time],
         ) -> str:
             return value.isoformat()
 

prefect/_internal/pytz.py

@@ -14,11 +14,12 @@ License: MIT
 """
 
 try:
-    import pytz  # noqa: F401
+    import pytz as pytz
 except ImportError:
-    HAS_PYTZ = False
+    _has_pytz = False
 else:
-    HAS_PYTZ = True
+    _has_pytz = True
+HAS_PYTZ = _has_pytz
 
 
 all_timezones_set = {

prefect/_internal/retries.py

@@ -1,6 +1,7 @@
 import asyncio
+from collections.abc import Coroutine
 from functools import wraps
-from typing import Callable, Optional, Tuple, Type, TypeVar
+from typing import Any, Callable, Optional, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -25,9 +26,11 @@ def retry_async_fn(
     ] = exponential_backoff_with_jitter,
     base_delay: float = 1,
     max_delay: float = 10,
-    retry_on_exceptions: Tuple[Type[Exception], ...] = (Exception,),
+    retry_on_exceptions: tuple[type[Exception], ...] = (Exception,),
     operation_name: Optional[str] = None,
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, Optional[R]]]
+]:
     """A decorator for retrying an async function.
 
     Args:
@@ -43,9 +46,11 @@ def retry_async_fn(
             the function name. If None, uses the function name.
     """
 
-    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+    def decorator(
+        func: Callable[P, Coroutine[Any, Any, R]],
+    ) -> Callable[P, Coroutine[Any, Any, Optional[R]]]:
         @wraps(func)
-        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
             name = operation_name or func.__name__
             for attempt in range(max_attempts):
                 try:

prefect/_internal/schemas/bases.py

@@ -4,18 +4,16 @@ Utilities for creating and working with Prefect REST API schemas.
 
 import datetime
 import os
-from typing import Any, ClassVar, Generator, Optional, Set, TypeVar, cast
+from typing import Any, ClassVar, Optional, TypeVar, cast
 from uuid import UUID, uuid4
 
 import pendulum
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-)
+from pydantic import BaseModel, ConfigDict, Field
+from rich.repr import RichReprResult
 from typing_extensions import Self
 
 from prefect.types import DateTime
+from prefect.utilities.generics import validate_list
 
 T = TypeVar("T")
 
@@ -32,7 +30,7 @@ class PrefectBaseModel(BaseModel):
     subtle unintentional testing errors.
     """
 
-    _reset_fields: ClassVar[Set[str]] = set()
+    _reset_fields: ClassVar[set[str]] = set()
 
     model_config: ClassVar[ConfigDict] = ConfigDict(
         ser_json_timedelta="float",
@@ -59,7 +57,18 @@ class PrefectBaseModel(BaseModel):
         else:
             return copy_dict == other
 
-    def __rich_repr__(self) -> Generator[tuple[str, Any, Any], None, None]:
+    @classmethod
+    def model_validate_list(
+        cls,
+        obj: Any,
+        *,
+        strict: Optional[bool] = None,
+        from_attributes: Optional[bool] = None,
+        context: Optional[Any] = None,
+    ) -> list[Self]:
+        return validate_list(cls, obj)
+
+    def __rich_repr__(self) -> RichReprResult:
         # Display all of the fields in the model if they differ from the default value
         for name, field in self.model_fields.items():
             value = getattr(self, name)
@@ -102,7 +111,7 @@ class IDBaseModel(PrefectBaseModel):
     The ID is reset on copy() and not included in equality comparisons.
     """
 
-    _reset_fields: ClassVar[Set[str]] = {"id"}
+    _reset_fields: ClassVar[set[str]] = {"id"}
     id: UUID = Field(default_factory=uuid4)
 
 
@@ -115,7 +124,7 @@ class ObjectBaseModel(IDBaseModel):
     equality comparisons.
     """
 
-    _reset_fields: ClassVar[Set[str]] = {"id", "created", "updated"}
+    _reset_fields: ClassVar[set[str]] = {"id", "created", "updated"}
     model_config: ClassVar[ConfigDict] = ConfigDict(from_attributes=True)
 
     created: Optional[DateTime] = Field(default=None, repr=False)