prefect-client 3.0.0rc10__py3-none-any.whl → 3.0.0rc11__py3-none-any.whl

prefect/_internal/concurrency/api.py

@@ -151,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 []:

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

prefect/client/cloud.py

@@ -9,7 +9,7 @@ from starlette import status
 import prefect.context
 import prefect.settings
 from prefect.client.base import PrefectHttpxAsyncClient
-from prefect.client.schemas import Workspace
+from prefect.client.schemas.objects import Workspace
 from prefect.exceptions import ObjectNotFound, PrefectException
 from prefect.settings import (
     PREFECT_API_KEY,

prefect/client/schemas/objects.py

@@ -791,7 +791,7 @@ class TaskRun(ObjectBaseModel):
 
     state: Optional[State] = Field(
         default=None,
-        description="The state of the flow run.",
+        description="The state of the task run.",
         examples=["State(type=StateType.COMPLETED)"],
     )
 

prefect/concurrency/asyncio.py

@@ -1,6 +1,6 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import List, Literal, Optional, Union, cast
+from typing import AsyncGenerator, List, Literal, Optional, Union, cast
 
 import httpx
 import pendulum
@@ -34,7 +34,7 @@ async def concurrency(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
-):
+) -> AsyncGenerator[None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
 
@@ -77,7 +77,7 @@ async def concurrency(
         _emit_concurrency_release_events(limits, occupy, emitted_events)
 
 
-async def rate_limit(names: Union[str, List[str]], occupy: int = 1):
+async def rate_limit(names: Union[str, List[str]], occupy: int = 1) -> None:
     """Block execution until an `occupy` number of slots of the concurrency
     limits given in `names` are acquired. Requires that all given concurrency
     limits have a slot decay.

prefect/concurrency/events.py

@@ -54,6 +54,6 @@ def _emit_concurrency_release_events(
     limits: List[MinimalConcurrencyLimitResponse],
     occupy: int,
     events: Dict[UUID, Optional[Event]],
-):
+) -> None:
     for limit in limits:
         _emit_concurrency_event("released", limit, limits, occupy, events[limit.id])

prefect/concurrency/services.py

@@ -3,6 +3,7 @@ import concurrent.futures
 from contextlib import asynccontextmanager
 from typing import (
     TYPE_CHECKING,
+    AsyncGenerator,
     FrozenSet,
     Optional,
     Tuple,
@@ -27,14 +28,14 @@ class ConcurrencySlotAcquisitionService(QueueService):
         self.concurrency_limit_names = sorted(list(concurrency_limit_names))
 
     @asynccontextmanager
-    async def _lifespan(self):
+    async def _lifespan(self) -> AsyncGenerator[None, None]:
         async with get_client() as client:
             self._client = client
             yield
 
     async def _handle(
         self, item: Tuple[int, str, Optional[float], concurrent.futures.Future]
-    ):
+    ) -> None:
         occupy, mode, timeout_seconds, future = item
         try:
             response = await self.acquire_slots(occupy, mode, timeout_seconds)

prefect/concurrency/sync.py

@@ -1,5 +1,15 @@
 from contextlib import contextmanager
-from typing import List, Optional, Union, cast
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Generator,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import pendulum
 
@@ -22,13 +32,15 @@ from .events import (
     _emit_concurrency_release_events,
 )
 
+T = TypeVar("T")
+
 
 @contextmanager
 def concurrency(
     names: Union[str, List[str]],
     occupy: int = 1,
     timeout_seconds: Optional[float] = None,
-):
+) -> Generator[None, None, None]:
     """A context manager that acquires and releases concurrency slots from the
     given concurrency limits.
 
@@ -75,7 +87,7 @@ def concurrency(
         _emit_concurrency_release_events(limits, occupy, emitted_events)
 
 
-def rate_limit(names: Union[str, List[str]], occupy: int = 1):
+def rate_limit(names: Union[str, List[str]], occupy: int = 1) -> None:
     """Block execution until an `occupy` number of slots of the concurrency
     limits given in `names` are acquired. Requires that all given concurrency
     limits have a slot decay.
@@ -91,11 +103,13 @@ def rate_limit(names: Union[str, List[str]], occupy: int = 1):
     _emit_concurrency_acquisition_events(limits, occupy)
 
 
-def _call_async_function_from_sync(fn, *args, **kwargs):
+def _call_async_function_from_sync(
+    fn: Callable[..., Awaitable[T]], *args: Any, **kwargs: Any
+) -> T:
     loop = get_running_loop()
     call = create_call(fn, *args, **kwargs)
 
     if loop is not None:
         return from_sync.call_soon_in_loop_thread(call).result()
     else:
-        return call()
+        return call()  # type: ignore [return-value]

prefect/context.py

@@ -9,6 +9,7 @@ For more user-accessible information about the current run, see [`prefect.runtim
 import os
 import sys
 import warnings
+import weakref
 from contextlib import ExitStack, contextmanager
 from contextvars import ContextVar, Token
 from pathlib import Path
@@ -17,6 +18,7 @@ from typing import (
     Any,
     Dict,
     Generator,
+    Mapping,
     Optional,
     Set,
     Type,
@@ -291,8 +293,12 @@ class EngineContext(RunContext):
     # Counter for flow pauses
     observed_flow_pauses: Dict[str, int] = Field(default_factory=dict)
 
-    # Tracking for result from task runs in this flow run
-    task_run_results: Dict[int, State] = Field(default_factory=dict)
+    # Tracking for result from task runs in this flow run for dependency tracking
+    # Holds the ID of the object returned by the task run and task run state
+    # This is a weakref dictionary to avoid undermining garbage collection
+    task_run_results: Mapping[int, State] = Field(
+        default_factory=weakref.WeakValueDictionary
+    )
 
     # Events worker to emit events to Prefect Cloud
     events: Optional[EventsWorker] = None

prefect/deployments/__init__.py

@@ -1,20 +1,33 @@
+from typing import TYPE_CHECKING
 from prefect._internal.compatibility.migration import getattr_migration
-import prefect.deployments.base
-import prefect.deployments.steps
-from prefect.deployments.base import (
-    initialize_project,
-)
 
-from prefect.deployments.runner import (
-    RunnerDeployment,
-    deploy,
-    DockerImage,
-    EntrypointType,
-)
 
+if TYPE_CHECKING:
+    from .flow_runs import run_deployment
+    from .base import initialize_project
+    from .runner import deploy
 
-from prefect.deployments.flow_runs import (
-    run_deployment,
-)
+_public_api: dict[str, tuple[str, str]] = {
+    "initialize_project": (__spec__.parent, ".base"),
+    "run_deployment": (__spec__.parent, ".flow_runs"),
+    "deploy": (__spec__.parent, ".runner"),
+}
 
-__getattr__ = getattr_migration(__name__)
+# Declare API for type-checkers
+__all__ = ["initialize_project", "deploy", "run_deployment"]
+
+
+def __getattr__(attr_name: str) -> object:
+    dynamic_attr = _public_api.get(attr_name)
+    if dynamic_attr is None:
+        return getattr_migration(__name__)(attr_name)
+
+    package, module_name = dynamic_attr
+
+    from importlib import import_module
+
+    if module_name == "__module__":
+        return import_module(f".{attr_name}", package=package)
+    else:
+        module = import_module(module_name, package=package)
+        return getattr(module, attr_name)

prefect/deployments/steps/pull.py

@@ -6,6 +6,7 @@ import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
+from prefect._internal.retries import retry_async_fn
 from prefect.logging.loggers import get_logger
 from prefect.runner.storage import BlockStorageAdapter, GitRepository, RemoteStorage
 from prefect.utilities.asyncutils import sync_compatible
@@ -31,6 +32,12 @@ def set_working_directory(directory: str) -> dict:
     return dict(directory=directory)
 
 
+@retry_async_fn(
+    max_attempts=3,
+    base_delay=1,
+    max_delay=10,
+    retry_on_exceptions=(RuntimeError,),
+)
 @sync_compatible
 async def git_clone(
     repository: str,

prefect/flow_engine.py

@@ -7,7 +7,6 @@ from dataclasses import dataclass, field
 from typing import (
     Any,
     AsyncGenerator,
-    Callable,
     Coroutine,
     Dict,
     Generator,
@@ -415,7 +414,7 @@ class FlowRunEngine(Generic[P, R]):
 
         return flow_run
 
-    def call_hooks(self, state: Optional[State] = None) -> Iterable[Callable]:
+    def call_hooks(self, state: Optional[State] = None):
         if state is None:
             state = self.state
         flow = self.flow
@@ -613,11 +612,7 @@ class FlowRunEngine(Generic[P, R]):
 
             if self.state.is_running():
                 self.call_hooks()
-            try:
-                yield
-            finally:
-                if self.state.is_final() or self.state.is_cancelling():
-                    self.call_hooks()
+            yield
 
     @contextmanager
     def run_context(self):
@@ -638,6 +633,9 @@ class FlowRunEngine(Generic[P, R]):
             except Exception as exc:
                 self.logger.exception("Encountered exception during execution: %r", exc)
                 self.handle_exception(exc)
+            finally:
+                if self.state.is_final() or self.state.is_cancelling():
+                    self.call_hooks()
 
     def call_flow_fn(self) -> Union[R, Coroutine[Any, Any, R]]:
         """