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

prefect/_internal/concurrency/api.py

@@ -6,7 +6,6 @@ import abc
 import asyncio
 import concurrent.futures
 import contextlib
-import threading
 from typing import (
     Awaitable,
     Callable,
@@ -19,7 +18,6 @@ from typing import (
 
 from typing_extensions import ParamSpec
 
-from prefect._internal.concurrency.calls import get_current_call
 from prefect._internal.concurrency.threads import (
     WorkerThread,
     get_global_loop,
@@ -29,7 +27,6 @@ from prefect._internal.concurrency.waiters import (
     AsyncWaiter,
     Call,
     SyncWaiter,
-    get_waiter_for_thread,
 )
 
 P = ParamSpec("P")
@@ -154,7 +151,7 @@ class from_async(_base):
         __call: Union[Callable[[], T], Call[T]],
         timeout: Optional[float] = None,
         done_callbacks: Optional[Iterable[Call]] = None,
-    ) -> Call[T]:
+    ) -> T:
         call = _cast_to_call(__call)
         waiter = AsyncWaiter(call=call)
         for callback in done_callbacks or []:
@@ -163,22 +160,6 @@ class from_async(_base):
         await waiter.wait()
         return call.result()
 
-    @staticmethod
-    def call_soon_in_waiting_thread(
-        __call: Union[Callable[[], T], Call[T]],
-        thread: threading.Thread,
-        timeout: Optional[float] = None,
-    ) -> Call[T]:
-        call = _cast_to_call(__call)
-        parent_call = get_current_call()
-        waiter = get_waiter_for_thread(thread, parent_call)
-        if waiter is None:
-            raise RuntimeError(f"No waiter found for thread {thread}.")
-
-        call.set_timeout(timeout)
-        waiter.submit(call)
-        return call
-
     @staticmethod
     def call_in_new_thread(
         __call: Union[Callable[[], T], Call[T]], timeout: Optional[float] = None
@@ -231,21 +212,6 @@ class from_sync(_base):
         waiter.wait()
         return call.result()
 
-    @staticmethod
-    def call_soon_in_waiting_thread(
-        __call: Union[Callable[[], T], Call[T]],
-        thread: threading.Thread,
-        timeout: Optional[float] = None,
-    ) -> Call[T]:
-        call = _cast_to_call(__call)
-        waiter = get_waiter_for_thread(thread)
-        if waiter is None:
-            raise RuntimeError(f"No waiter found for thread {thread}.")
-
-        call.set_timeout(timeout)
-        waiter.submit(call)
-        return call
-
     @staticmethod
     def call_in_new_thread(
         __call: Union[Callable[[], T], Call[T]], timeout: Optional[float] = None

prefect/_internal/concurrency/calls.py

@@ -49,12 +49,6 @@ current_call: contextvars.ContextVar["weakref.ref[Call]"] = (  # novm
 _ASYNC_TASK_REFS = set()
 
 
-def get_current_call() -> Optional["Call"]:
-    call_ref = current_call.get(None)
-    if call_ref:
-        return call_ref()
-
-
 @contextlib.contextmanager
 def set_current_call(call: "Call"):
     token = current_call.set(weakref.ref(call))

prefect/_internal/concurrency/cancellation.py

@@ -76,9 +76,6 @@ class ThreadShield:
 
 
 class CancelledError(asyncio.CancelledError):
-    # In Python 3.7, `asyncio.CancelledError` is identical to `concurrent.futures.CancelledError`
-    # but in 3.8+ it is a separate class that inherits from `BaseException` instead
-    # See https://bugs.python.org/issue32528
     # We want our `CancelledError` to be treated as a `BaseException` and defining it
     # here simplifies downstream logic that needs to know "which" cancelled error to
     # handle.

prefect/_internal/concurrency/event_loop.py

@@ -25,26 +25,6 @@ def get_running_loop() -> Optional[asyncio.BaseEventLoop]:
         return None
 
 
-def call_in_loop(
-    __loop: asyncio.AbstractEventLoop,
-    __fn: Callable[P, T],
-    *args: P.args,
-    **kwargs: P.kwargs,
-) -> T:
-    """
-    Run a synchronous call in event loop's thread from another thread.
-
-    This function is blocking and not safe to call from an asynchronous context.
-
-    Returns the result of the call.
-    """
-    if __loop is get_running_loop():
-        return __fn(*args, **kwargs)
-    else:
-        future = call_soon_in_loop(__loop, __fn, *args, **kwargs)
-        return future.result()
-
-
 def call_soon_in_loop(
     __loop: asyncio.AbstractEventLoop,
     __fn: Callable[P, T],

prefect/_internal/concurrency/inspection.py

@@ -7,10 +7,10 @@ import linecache
 import sys
 import threading
 from types import FrameType
-from typing import List
+from typing import List, Optional
 
 """
-The following functions are derived from dask/distributed which is licensed under the 
+The following functions are derived from dask/distributed which is licensed under the
 BSD 3-Clause License.
 
 Copyright (c) 2015, Anaconda, Inc. and contributors
@@ -75,7 +75,7 @@ def repr_frame(frame: FrameType) -> str:
 def call_stack(frame: FrameType) -> List[str]:
     """Create a call text stack from a frame"""
     L = []
-    cur_frame: FrameType | None = frame
+    cur_frame: Optional[FrameType] = frame
     while cur_frame:
         L.append(repr_frame(cur_frame))
         cur_frame = cur_frame.f_back

prefect/_internal/concurrency/primitives.py

@@ -1,6 +1,7 @@
 """
 Thread-safe async synchronization primitives.
 """
+
 import asyncio
 import collections
 import threading

prefect/_internal/concurrency/services.py

@@ -39,6 +39,7 @@ class QueueService(abc.ABC, Generic[T]):
             daemon=True,
             name=f"{type(self).__name__}Thread",
         )
+        self._logger = logging.getLogger(f"{type(self).__name__}")
 
     def start(self):
         logger.debug("Starting service %r", self)
@@ -144,13 +145,27 @@ class QueueService(abc.ABC, Generic[T]):
             self._done_event.set()
 
     async def _main_loop(self):
+        last_log_time = 0
+        log_interval = 4  # log every 4 seconds
+
         while True:
             item: T = await self._queue_get_thread.submit(
                 create_call(self._queue.get)
             ).aresult()
 
+            if self._stopped:
+                current_time = asyncio.get_event_loop().time()
+                queue_size = self._queue.qsize()
+
+                if current_time - last_log_time >= log_interval and queue_size > 0:
+                    self._logger.warning(
+                        f"Still processing items: {queue_size} items remaining..."
+                    )
+                    last_log_time = current_time
+
             if item is None:
                 logger.debug("Exiting service %r", self)
+                self._queue.task_done()
                 break
 
             try:
@@ -164,6 +179,8 @@ class QueueService(abc.ABC, Generic[T]):
                     item,
                     exc_info=log_traceback,
                 )
+            finally:
+                self._queue.task_done()
 
     @abc.abstractmethod
     async def _handle(self, item: T):
@@ -235,6 +252,12 @@ class QueueService(abc.ABC, Generic[T]):
         else:
             return concurrent.futures.wait(futures, timeout=timeout)
 
+    def wait_until_empty(self):
+        """
+        Wait until the queue is empty and all items have been processed.
+        """
+        self._queue.join()
+
     @classmethod
     def instance(cls: Type[Self], *args) -> Self:
         """

prefect/_internal/concurrency/threads.py

@@ -1,6 +1,7 @@
 """
 Utilities for managing worker threads.
 """
+
 import asyncio
 import atexit
 import concurrent.futures
@@ -233,7 +234,10 @@ class EventLoopThread(Portal):
         self.shutdown()
 
 
+# the GLOBAL LOOP is used for background services, like logs
 GLOBAL_LOOP: Optional[EventLoopThread] = None
+# the RUN SYNC LOOP is used exclusively for running async functions in a sync context via asyncutils.run_sync
+RUN_SYNC_LOOP: Optional[EventLoopThread] = None
 
 
 def get_global_loop() -> EventLoopThread:
@@ -267,6 +271,37 @@ def in_global_loop() -> bool:
     return get_global_loop()._loop == get_running_loop()
 
 
+def get_run_sync_loop() -> EventLoopThread:
+    """
+    Get the run_sync loop thread.
+
+    Creates a new one if there is not one available.
+    """
+    global RUN_SYNC_LOOP
+
+    # Create a new worker on first call or if the existing worker is dead
+    if (
+        RUN_SYNC_LOOP is None
+        or not RUN_SYNC_LOOP.thread.is_alive()
+        or RUN_SYNC_LOOP._shutdown_event.is_set()
+    ):
+        RUN_SYNC_LOOP = EventLoopThread(daemon=True, name="RunSyncEventLoopThread")
+        RUN_SYNC_LOOP.start()
+
+    return RUN_SYNC_LOOP
+
+
+def in_run_sync_loop() -> bool:
+    """
+    Check if called from the global loop.
+    """
+    if RUN_SYNC_LOOP is None:
+        # Avoid creating a global loop if there isn't one
+        return False
+
+    return get_run_sync_loop()._loop == get_running_loop()
+
+
 def wait_for_global_loop_exit(timeout: Optional[float] = None) -> None:
     """
     Shutdown the global loop and wait for it to exit.

prefect/_internal/concurrency/waiters.py

@@ -29,34 +29,6 @@ _WAITERS_BY_THREAD: "WeakKeyDictionary[threading.Thread, deque[Waiter]]" = (
 )
 
 
-def get_waiter_for_thread(
-    thread: threading.Thread, parent_call: Optional[Call] = None
-) -> Optional["Waiter"]:
-    """
-    Get the current waiter for a thread and an optional parent call.
-
-    To avoid assigning outer callbacks to inner waiters in the case of nested calls,
-    the parent call is used to determine which waiter to return. If a parent call is
-    not provided, we return the most recently created waiter (last in the stack).
-
-    see https://github.com/PrefectHQ/prefect/issues/12036
-
-    Returns `None` if no active waiter is found for the thread.
-    """
-
-    waiters: "Optional[deque[Waiter]]" = _WAITERS_BY_THREAD.get(thread)
-
-    if waiters and (active_waiters := [w for w in waiters if not w.call_is_done()]):
-        if parent_call and (
-            matching_waiter := next(
-                (w for w in active_waiters if w._call == parent_call), None
-            )
-        ):  # if exists an active waiter responsible for the parent call, return it
-            return matching_waiter
-        else:  # otherwise, return the most recently created waiter
-            return active_waiters[-1]
-
-
 def add_waiter_for_thread(waiter: "Waiter", thread: threading.Thread):
     """
     Add a waiter for a thread.

prefect/_internal/integrations.py

@@ -0,0 +1,7 @@
+KNOWN_EXTRAS_FOR_PACKAGES = {
+    "prefect-kubernetes": "prefect[kubernetes]",
+    "prefect-aws": "prefect[aws]",
+    "prefect-gcp": "prefect[gcp]",
+    "prefect-azure": "prefect[azure]",
+    "prefect-docker": "prefect[docker]",
+}

prefect/_internal/pydantic/__init__.py

@@ -1,46 +1 @@
-### A convenience module to allow for easy switching between pydantic versions.
 
-### Note this introduces a marginally worse import time, since
-### the import of any one of these symbols will import all of them.
-
-### This is a tradeoff we're willing to make for now until pydantic v1 is
-### no longer supported.
-
-
-from ._flags import HAS_PYDANTIC_V2
-
-from ._compat import (
-    model_dump,
-    model_json_schema,
-    model_validate,
-    model_dump_json,
-    model_copy,
-    model_validate_json,
-    TypeAdapter,
-    validate_python,
-    BaseModel,
-    Field,
-    FieldInfo,
-    field_validator,
-    model_validator,
-)
-
-from ._types import IncEx
-
-__all__ = [
-    "model_dump",
-    "model_json_schema",
-    "model_validate",
-    "IncEx",
-    "model_dump_json",
-    "model_copy",
-    "model_validate_json",
-    "TypeAdapter",
-    "validate_python",
-    "BaseModel",
-    "HAS_PYDANTIC_V2",
-    "Field",
-    "FieldInfo",
-    "field_validator",
-    "model_validator",
-]

prefect/_internal/pydantic/annotations/pendulum.py

@@ -1,6 +1,6 @@
 """
-This file contains compat code to handle pendulum.DateTime objects during jsonschema 
-generation and validation. 
+This file contains compat code to handle pendulum.DateTime objects during jsonschema
+generation and validation.
 """
 
 import typing as t

prefect/_internal/pydantic/v1_schema.py

@@ -1,37 +1,36 @@
 import inspect
 import typing
+import warnings
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import BaseModel as V1BaseModel
-else:
-    from pydantic import BaseModel as V1BaseModel
+import pydantic
+from pydantic.v1 import BaseModel as V1BaseModel
 
 
 def is_v1_model(v) -> bool:
-    if isinstance(v, V1BaseModel):
-        return True
-    try:
-        if inspect.isclass(v) and issubclass(v, V1BaseModel):
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=pydantic.warnings.PydanticDeprecatedSince20
+        )
+
+        if isinstance(v, V1BaseModel):
             return True
-    except TypeError:
-        pass
+        try:
+            if inspect.isclass(v) and issubclass(v, V1BaseModel):
+                return True
+        except TypeError:
+            pass
 
-    return False
+        return False
 
 
 def is_v1_type(v) -> bool:
-    if HAS_PYDANTIC_V2:
-        if is_v1_model(v):
-            return True
-
-        try:
-            return v.__module__.startswith("pydantic.v1.types")
-        except AttributeError:
-            return False
+    if is_v1_model(v):
+        return True
 
-    return True
+    try:
+        return v.__module__.startswith("pydantic.v1.types")
+    except AttributeError:
+        return False
 
 
 def has_v1_type_as_param(signature: inspect.Signature) -> bool:

prefect/_internal/pydantic/v2_schema.py

@@ -1,5 +1,3 @@
-# Note: This file should only be imported under a HAS_PYDANTIC_V2 flag
-
 import inspect
 import typing
 import typing as t

prefect/_internal/pydantic/v2_validated_func.py

@@ -1,16 +1,15 @@
 """
-This module contains an implementation of pydantic v1's ValidateFunction 
+This module contains an implementation of pydantic v1's ValidateFunction
 modified to validate function arguments and return a pydantic v2 model.
 
-Specifically it allows for us to validate v2 models used as flow/task 
+Specifically it allows for us to validate v2 models used as flow/task
 arguments.
 """
 
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 # importing directly from v2 to be able to create a v2 model
-from pydantic import BaseModel, create_model
-from pydantic.v1 import validator
+from pydantic import BaseModel, ConfigDict, create_model, field_validator
 from pydantic.v1.decorator import ValidatedFunction
 from pydantic.v1.errors import ConfigError
 from pydantic.v1.utils import to_camel
@@ -28,28 +27,25 @@ class V2ValidatedFunction(ValidatedFunction):
         fields: Dict[str, Any],
         takes_args: bool,
         takes_kwargs: bool,
-        config: "ConfigType",
+        config: ConfigDict,
     ) -> None:
         pos_args = len(self.arg_mapping)
 
-        class CustomConfig:
-            pass
-
-        if not TYPE_CHECKING:  # pragma: no branch
-            if isinstance(config, dict):
-                CustomConfig = type("Config", (), config)  # noqa: F811
-            elif config is not None:
-                CustomConfig = config  # noqa: F811
-
-        if hasattr(CustomConfig, "fields") or hasattr(CustomConfig, "alias_generator"):
+        if config.get("fields") or config.get("alias_generator"):
             raise ConfigError(
                 'Setting the "fields" and "alias_generator" property on custom Config'
                 " for @validate_arguments is not yet supported, please remove."
             )
 
+        if "extra" not in config:
+            config["extra"] = "forbid"
+
         # This is the key change -- inheriting the BaseModel class from v2
         class DecoratorBaseModel(BaseModel):
-            @validator(self.v_args_name, check_fields=False, allow_reuse=True)
+            model_config = config
+
+            @field_validator(self.v_args_name, check_fields=False)
+            @classmethod
             def check_args(cls, v: Optional[List[Any]]) -> Optional[List[Any]]:
                 if takes_args or v is None:
                     return v
@@ -59,7 +55,8 @@ class V2ValidatedFunction(ValidatedFunction):
                     f" {pos_args + len(v)} given"
                 )
 
-            @validator(self.v_kwargs_name, check_fields=False, allow_reuse=True)
+            @field_validator(self.v_kwargs_name, check_fields=False)
+            @classmethod
             def check_kwargs(
                 cls, v: Optional[Dict[str, Any]]
             ) -> Optional[Dict[str, Any]]:
@@ -70,7 +67,8 @@ class V2ValidatedFunction(ValidatedFunction):
                 keys = ", ".join(map(repr, v.keys()))
                 raise TypeError(f"unexpected keyword argument{plural}: {keys}")
 
-            @validator(V_POSITIONAL_ONLY_NAME, check_fields=False, allow_reuse=True)
+            @field_validator(V_POSITIONAL_ONLY_NAME, check_fields=False)
+            @classmethod
             def check_positional_only(cls, v: Optional[List[str]]) -> None:
                 if v is None:
                     return
@@ -82,7 +80,8 @@ class V2ValidatedFunction(ValidatedFunction):
                     f" argument{plural}: {keys}"
                 )
 
-            @validator(V_DUPLICATE_KWARGS, check_fields=False, allow_reuse=True)
+            @field_validator(V_DUPLICATE_KWARGS, check_fields=False)
+            @classmethod
             def check_duplicate_kwargs(cls, v: Optional[List[str]]) -> None:
                 if v is None:
                     return
@@ -91,10 +90,6 @@ class V2ValidatedFunction(ValidatedFunction):
                 keys = ", ".join(map(repr, v))
                 raise TypeError(f"multiple values for argument{plural}: {keys}")
 
-            class Config(CustomConfig):
-                # extra = getattr(CustomConfig, "extra", Extra.forbid)
-                extra = getattr(CustomConfig, "extra", "forbid")
-
         self.model = create_model(
             to_camel(self.raw_function.__name__),
             __base__=DecoratorBaseModel,

prefect/_internal/pytz.py

@@ -1,5 +1,5 @@
 """
-This is a compatibility module that allows us to drop our dependency on pytz in the minimal prefect-client. 
+This is a compatibility module that allows us to drop our dependency on pytz in the minimal prefect-client.
 
 All values here are taken from:
 

prefect/_internal/retries.py

@@ -0,0 +1,61 @@
+import asyncio
+from functools import wraps
+from typing import Any, Callable, Tuple, Type
+
+from prefect.logging.loggers import get_logger
+from prefect.utilities.math import clamped_poisson_interval
+
+logger = get_logger("retries")
+
+
+def exponential_backoff_with_jitter(
+    attempt: int, base_delay: float, max_delay: float
+) -> float:
+    average_interval = min(base_delay * (2**attempt), max_delay)
+    return clamped_poisson_interval(average_interval, clamping_factor=0.3)
+
+
+def retry_async_fn(
+    max_attempts: int = 3,
+    backoff_strategy: Callable[
+        [int, float, float], float
+    ] = exponential_backoff_with_jitter,
+    base_delay: float = 1,
+    max_delay: float = 10,
+    retry_on_exceptions: Tuple[Type[Exception], ...] = (Exception,),
+):
+    """A decorator for retrying an async function.
+
+    Args:
+        max_attempts: The maximum number of times to retry the function.
+        backoff_strategy: A function that takes in the number of attempts, the base
+            delay, and the maximum delay, and returns the delay to use for the next
+            attempt. Defaults to an exponential backoff with jitter.
+        base_delay: The base delay to use for the first attempt.
+        max_delay: The maximum delay to use for the last attempt.
+        retry_on_exceptions: A tuple of exception types to retry on. Defaults to
+            retrying on all exceptions.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            for attempt in range(max_attempts):
+                try:
+                    return await func(*args, **kwargs)
+                except retry_on_exceptions as e:
+                    if attempt == max_attempts - 1:
+                        logger.exception(
+                            f"Function {func.__name__!r} failed after {max_attempts} attempts"
+                        )
+                        raise
+                    delay = backoff_strategy(attempt, base_delay, max_delay)
+                    logger.warning(
+                        f"Attempt {attempt + 1} of function {func.__name__!r} failed with {type(e).__name__}. "
+                        f"Retrying in {delay:.2f} seconds..."
+                    )
+                    await asyncio.sleep(delay)
+
+        return wrapper
+
+    return decorator