prefect-client 3.0.0rc4__py3-none-any.whl → 3.0.0rc6__py3-none-any.whl

prefect/__init__.py

@@ -49,8 +49,6 @@ import prefect.blocks.system
 # Initialize the process-wide profile and registry at import time
 import prefect.context
 
-prefect.context.initialize_object_registry()
-
 # Perform any forward-ref updates needed for Pydantic models
 import prefect.client.schemas
 

prefect/{records/cache_policies.py → cache_policies.py}

@@ -8,6 +8,10 @@ from prefect.utilities.hashing import hash_objects
 
 @dataclass
 class CachePolicy:
+    """
+    Base class for all cache policies.
+    """
+
     @classmethod
     def from_cache_key_fn(
         cls, cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
@@ -59,6 +63,11 @@ class CachePolicy:
 
 @dataclass
 class CacheKeyFnPolicy(CachePolicy):
+    """
+    This policy accepts a custom function with signature f(task_run_context, task_parameters, flow_parameters) -> str
+    and uses it to compute a task run cache key.
+    """
+
     # making it optional for tests
     cache_key_fn: Optional[
         Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
@@ -77,6 +86,13 @@ class CacheKeyFnPolicy(CachePolicy):
 
 @dataclass
 class CompoundCachePolicy(CachePolicy):
+    """
+    This policy is constructed from two or more other cache policies and works by computing the keys
+    for each policy individually, and then hashing a sorted tuple of all computed keys.
+
+    Any keys that return `None` will be ignored.
+    """
+
     policies: Optional[list] = None
 
     def compute_key(
@@ -88,20 +104,25 @@ class CompoundCachePolicy(CachePolicy):
     ) -> Optional[str]:
         keys = []
         for policy in self.policies or []:
-            keys.append(
-                policy.compute_key(
-                    task_ctx=task_ctx,
-                    inputs=inputs,
-                    flow_parameters=flow_parameters,
-                    **kwargs,
-                )
+            policy_key = policy.compute_key(
+                task_ctx=task_ctx,
+                inputs=inputs,
+                flow_parameters=flow_parameters,
+                **kwargs,
             )
+            if policy_key is not None:
+                keys.append(policy_key)
+        if not keys:
+            return None
         return hash_objects(*keys)
 
 
 @dataclass
-class Default(CachePolicy):
-    "Execution run ID only"
+class _None(CachePolicy):
+    """
+    Policy that always returns `None` for the computed cache key.
+    This policy prevents persistence.
+    """
 
     def compute_key(
         self,
@@ -110,12 +131,14 @@ class Default(CachePolicy):
         flow_parameters: Dict[str, Any],
         **kwargs,
     ) -> Optional[str]:
-        return str(task_ctx.task_run.id)
+        return None
 
 
 @dataclass
-class _None(CachePolicy):
-    "ignore key policies altogether, always run - prevents persistence"
+class TaskSource(CachePolicy):
+    """
+    Policy for computing a cache key based on the source code of the task.
+    """
 
     def compute_key(
         self,
@@ -124,11 +147,22 @@ class _None(CachePolicy):
         flow_parameters: Dict[str, Any],
         **kwargs,
     ) -> Optional[str]:
-        return None
+        if not task_ctx:
+            return None
+        try:
+            lines = inspect.getsource(task_ctx.task)
+        except TypeError:
+            lines = inspect.getsource(task_ctx.task.fn.__class__)
+
+        return hash_objects(lines)
 
 
 @dataclass
-class TaskDef(CachePolicy):
+class FlowParameters(CachePolicy):
+    """
+    Policy that computes the cache key based on a hash of the flow parameters.
+    """
+
     def compute_key(
         self,
         task_ctx: TaskRunContext,
@@ -136,21 +170,37 @@ class TaskDef(CachePolicy):
         flow_parameters: Dict[str, Any],
         **kwargs,
     ) -> Optional[str]:
-        lines = inspect.getsource(task_ctx.task)
-        return hash_objects(lines)
+        if not flow_parameters:
+            return None
+        return hash_objects(flow_parameters)
 
 
 @dataclass
-class FlowParameters(CachePolicy):
-    pass
+class RunId(CachePolicy):
+    """
+    Returns either the prevailing flow run ID, or if not found, the prevailing task
+    run ID.
+    """
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        if not task_ctx:
+            return None
+        run_id = task_ctx.task_run.flow_run_id
+        if run_id is None:
+            run_id = task_ctx.task_run.id
+        return str(run_id)
 
 
 @dataclass
 class Inputs(CachePolicy):
     """
-    Exposes flag for whether to include flow parameters as well.
-
-    And exclude/include config.
+    Policy that computes a cache key based on a hash of the runtime inputs provided to the task..
     """
 
     exclude: Optional[list] = None
@@ -166,6 +216,9 @@ class Inputs(CachePolicy):
         inputs = inputs or {}
         exclude = self.exclude or []
 
+        if not inputs:
+            return None
+
         for key, val in inputs.items():
             if key not in exclude:
                 hashed_inputs[key] = val
@@ -173,7 +226,9 @@ class Inputs(CachePolicy):
         return hash_objects(hashed_inputs)
 
 
-DEFAULT = Default()
 INPUTS = Inputs()
 NONE = _None()
-TASKDEF = TaskDef()
+TASK_SOURCE = TaskSource()
+FLOW_PARAMETERS = FlowParameters()
+RUN_ID = RunId()
+DEFAULT = INPUTS + TASK_SOURCE + RUN_ID

prefect/client/schemas/schedules.py

@@ -3,13 +3,14 @@ Schedule schemas
 """
 
 import datetime
-from typing import Annotated, Optional, Union
+from typing import Annotated, Any, Optional, Union
 
 import dateutil
 import dateutil.rrule
 import pendulum
 from pydantic import AfterValidator, ConfigDict, Field, field_validator, model_validator
 from pydantic_extra_types.pendulum_dt import DateTime
+from typing_extensions import TypeAlias, TypeGuard
 
 from prefect._internal.schemas.bases import PrefectBaseModel
 from prefect._internal.schemas.validators import (
@@ -279,7 +280,13 @@ class NoSchedule(PrefectBaseModel):
     model_config = ConfigDict(extra="forbid")
 
 
-SCHEDULE_TYPES = Union[IntervalSchedule, CronSchedule, RRuleSchedule, NoSchedule]
+SCHEDULE_TYPES: TypeAlias = Union[
+    IntervalSchedule, CronSchedule, RRuleSchedule, NoSchedule
+]
+
+
+def is_schedule_type(obj: Any) -> TypeGuard[SCHEDULE_TYPES]:
+    return isinstance(obj, (IntervalSchedule, CronSchedule, RRuleSchedule, NoSchedule))
 
 
 def construct_schedule(

prefect/client/types/flexible_schedule_list.py

@@ -0,0 +1,11 @@
+from typing import TYPE_CHECKING, Any, Sequence, Union
+
+from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    from prefect.client.schemas.actions import DeploymentScheduleCreate
+    from prefect.client.schemas.schedules import SCHEDULE_TYPES
+
+FlexibleScheduleList: TypeAlias = Sequence[
+    Union[DeploymentScheduleCreate, dict[str, Any], "SCHEDULE_TYPES"]
+]

prefect/concurrency/asyncio.py

@@ -13,7 +13,6 @@ except ImportError:
 
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
-from prefect.utilities.timeout import timeout_async
 
 from .events import (
     _emit_concurrency_acquisition_events,
@@ -26,6 +25,10 @@ class ConcurrencySlotAcquisitionError(Exception):
     """Raised when an unhandlable occurs while acquiring concurrency slots."""
 
 
+class AcquireConcurrencySlotTimeoutError(TimeoutError):
+    """Raised when acquiring a concurrency slot times out."""
+
+
 @asynccontextmanager
 async def concurrency(
     names: Union[str, List[str]],
@@ -58,8 +61,9 @@ async def concurrency(
     ```
     """
     names = names if isinstance(names, list) else [names]
-    with timeout_async(seconds=timeout_seconds):
-        limits = await _acquire_concurrency_slots(names, occupy)
+    limits = await _acquire_concurrency_slots(
+        names, occupy, timeout_seconds=timeout_seconds
+    )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
 
@@ -91,12 +95,18 @@ async def _acquire_concurrency_slots(
     names: List[str],
     slots: int,
     mode: Union[Literal["concurrency"], Literal["rate_limit"]] = "concurrency",
+    timeout_seconds: Optional[float] = None,
 ) -> List[MinimalConcurrencyLimitResponse]:
     service = ConcurrencySlotAcquisitionService.instance(frozenset(names))
-    future = service.send((slots, mode))
+    future = service.send((slots, mode, timeout_seconds))
     response_or_exception = await asyncio.wrap_future(future)
 
     if isinstance(response_or_exception, Exception):
+        if isinstance(response_or_exception, TimeoutError):
+            raise AcquireConcurrencySlotTimeoutError(
+                f"Attempt to acquire concurrency slots timed out after {timeout_seconds} second(s)"
+            ) from response_or_exception
+
         raise ConcurrencySlotAcquisitionError(
             f"Unable to acquire concurrency slots on {names!r}"
         ) from response_or_exception

prefect/concurrency/services.py

@@ -4,6 +4,7 @@ from contextlib import asynccontextmanager
 from typing import (
     TYPE_CHECKING,
     FrozenSet,
+    Optional,
     Tuple,
 )
 
@@ -13,6 +14,7 @@ from starlette import status
 from prefect._internal.concurrency import logger
 from prefect._internal.concurrency.services import QueueService
 from prefect.client.orchestration import get_client
+from prefect.utilities.timeout import timeout_async
 
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
@@ -30,10 +32,12 @@ class ConcurrencySlotAcquisitionService(QueueService):
             self._client = client
             yield
 
-    async def _handle(self, item: Tuple[int, str, concurrent.futures.Future]):
-        occupy, mode, future = item
+    async def _handle(
+        self, item: Tuple[int, str, Optional[float], concurrent.futures.Future]
+    ):
+        occupy, mode, timeout_seconds, future = item
         try:
-            response = await self.acquire_slots(occupy, mode)
+            response = await self.acquire_slots(occupy, mode, timeout_seconds)
         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 that the caller can
@@ -43,25 +47,28 @@ class ConcurrencySlotAcquisitionService(QueueService):
         else:
             future.set_result(response)
 
-    async def acquire_slots(self, slots: int, mode: str) -> httpx.Response:
-        while True:
-            try:
-                response = await self._client.increment_concurrency_slots(
-                    names=self.concurrency_limit_names, slots=slots, mode=mode
-                )
-            except Exception as exc:
-                if (
-                    isinstance(exc, httpx.HTTPStatusError)
-                    and exc.response.status_code == status.HTTP_423_LOCKED
-                ):
-                    retry_after = float(exc.response.headers["Retry-After"])
-                    await asyncio.sleep(retry_after)
+    async def acquire_slots(
+        self, slots: int, mode: str, timeout_seconds: Optional[float] = None
+    ) -> httpx.Response:
+        with timeout_async(seconds=timeout_seconds):
+            while True:
+                try:
+                    response = await self._client.increment_concurrency_slots(
+                        names=self.concurrency_limit_names, slots=slots, mode=mode
+                    )
+                except Exception as exc:
+                    if (
+                        isinstance(exc, httpx.HTTPStatusError)
+                        and exc.response.status_code == status.HTTP_423_LOCKED
+                    ):
+                        retry_after = float(exc.response.headers["Retry-After"])
+                        await asyncio.sleep(retry_after)
+                    else:
+                        raise exc
                 else:
-                    raise exc
-            else:
-                return response
+                    return response
 
-    def send(self, item: Tuple[int, str]):
+    def send(self, item: Tuple[int, str, Optional[float]]) -> concurrent.futures.Future:
         with self._lock:
             if self._stopped:
                 raise RuntimeError("Cannot put items in a stopped service instance.")
@@ -69,7 +76,7 @@ class ConcurrencySlotAcquisitionService(QueueService):
             logger.debug("Service %r enqueuing item %r", self, item)
             future: concurrent.futures.Future = concurrent.futures.Future()
 
-            occupy, mode = item
-            self._queue.put_nowait((occupy, mode, future))
+            occupy, mode, timeout_seconds = item
+            self._queue.put_nowait((occupy, mode, timeout_seconds, future))
 
         return future

prefect/concurrency/sync.py

@@ -12,7 +12,6 @@ except ImportError:
 from prefect._internal.concurrency.api import create_call, from_sync
 from prefect._internal.concurrency.event_loop import get_running_loop
 from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
-from prefect.utilities.timeout import timeout
 
 from .asyncio import (
     _acquire_concurrency_slots,
@@ -57,10 +56,9 @@ def concurrency(
     """
     names = names if isinstance(names, list) else [names]
 
-    with timeout(seconds=timeout_seconds):
-        limits: List[MinimalConcurrencyLimitResponse] = _call_async_function_from_sync(
-            _acquire_concurrency_slots, names, occupy
-        )
+    limits: List[MinimalConcurrencyLimitResponse] = _call_async_function_from_sync(
+        _acquire_concurrency_slots, names, occupy, timeout_seconds=timeout_seconds
+    )
     acquisition_time = pendulum.now("UTC")
     emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
 

prefect/context.py

@@ -9,21 +9,16 @@ For more user-accessible information about the current run, see [`prefect.runtim
 import os
 import sys
 import warnings
-from collections import defaultdict
 from contextlib import ExitStack, contextmanager
 from contextvars import ContextVar, Token
-from functools import update_wrapper
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
-    ContextManager,
     Dict,
     Generator,
-    List,
     Optional,
     Set,
-    Tuple,
     Type,
     TypeVar,
     Union,
@@ -46,7 +41,6 @@ from prefect.settings import PREFECT_HOME, Profile, Settings
 from prefect.states import State
 from prefect.task_runners import TaskRunner
 from prefect.utilities.asyncutils import run_coro_as_sync
-from prefect.utilities.importtools import load_script_as_module
 
 T = TypeVar("T")
 
@@ -180,86 +174,6 @@ class ContextModel(BaseModel):
         return self.model_dump(exclude_unset=True)
 
 
-class PrefectObjectRegistry(ContextModel):
-    """
-    A context that acts as a registry for all Prefect objects that are
-    registered during load and execution.
-
-    Attributes:
-        start_time: The time the object registry was created.
-        block_code_execution: If set, flow calls will be ignored.
-        capture_failures: If set, failures during __init__ will be silenced and tracked.
-    """
-
-    start_time: DateTime = Field(default_factory=lambda: pendulum.now("UTC"))
-
-    _instance_registry: Dict[Type[T], List[T]] = PrivateAttr(
-        default_factory=lambda: defaultdict(list)
-    )
-
-    # Failures will be a tuple of (exception, instance, args, kwargs)
-    _instance_init_failures: Dict[
-        Type[T], List[Tuple[Exception, T, Tuple, Dict]]
-    ] = PrivateAttr(default_factory=lambda: defaultdict(list))
-
-    block_code_execution: bool = False
-    capture_failures: bool = False
-
-    __var__ = ContextVar("object_registry")
-
-    def get_instances(self, type_: Type[T]) -> List[T]:
-        instances = []
-        for registered_type, type_instances in self._instance_registry.items():
-            if type_ in registered_type.mro():
-                instances.extend(type_instances)
-        return instances
-
-    def get_instance_failures(
-        self, type_: Type[T]
-    ) -> List[Tuple[Exception, T, Tuple, Dict]]:
-        failures = []
-        for type__ in type_.mro():
-            failures.extend(self._instance_init_failures[type__])
-        return failures
-
-    def register_instance(self, object):
-        # TODO: Consider using a 'Set' to avoid duplicate entries
-        self._instance_registry[type(object)].append(object)
-
-    def register_init_failure(
-        self, exc: Exception, object: Any, init_args: Tuple, init_kwargs: Dict
-    ):
-        self._instance_init_failures[type(object)].append(
-            (exc, object, init_args, init_kwargs)
-        )
-
-    @classmethod
-    def register_instances(cls, type_: Type[T]) -> Type[T]:
-        """
-        Decorator for a class that adds registration to the `PrefectObjectRegistry`
-        on initialization of instances.
-        """
-        original_init = type_.__init__
-
-        def __register_init__(__self__: T, *args: Any, **kwargs: Any) -> None:
-            registry = cls.get()
-            try:
-                original_init(__self__, *args, **kwargs)
-            except Exception as exc:
-                if not registry or not registry.capture_failures:
-                    raise
-                else:
-                    registry.register_init_failure(exc, __self__, args, kwargs)
-            else:
-                if registry:
-                    registry.register_instance(__self__)
-
-        update_wrapper(__register_init__, original_init)
-
-        type_.__init__ = __register_init__
-        return type_
-
-
 class ClientContext(ContextModel):
     """
     A context for managing the Prefect client instances.
@@ -594,23 +508,6 @@ def tags(*new_tags: str) -> Generator[Set[str], None, None]:
         yield new_tags
 
 
-def registry_from_script(
-    path: str,
-    block_code_execution: bool = True,
-    capture_failures: bool = True,
-) -> PrefectObjectRegistry:
-    """
-    Return a fresh registry with instances populated from execution of a script.
-    """
-    with PrefectObjectRegistry(
-        block_code_execution=block_code_execution,
-        capture_failures=capture_failures,
-    ) as registry:
-        load_script_as_module(path)
-
-    return registry
-
-
 @contextmanager
 def use_profile(
     profile: Union[Profile, str],
@@ -711,14 +608,3 @@ def root_settings_context():
 
 
 GLOBAL_SETTINGS_CONTEXT: SettingsContext = root_settings_context()
-GLOBAL_OBJECT_REGISTRY: Optional[ContextManager[PrefectObjectRegistry]] = None
-
-
-def initialize_object_registry():
-    global GLOBAL_OBJECT_REGISTRY
-
-    if GLOBAL_OBJECT_REGISTRY:
-        return
-
-    GLOBAL_OBJECT_REGISTRY = PrefectObjectRegistry()
-    GLOBAL_OBJECT_REGISTRY.__enter__()

prefect/deployments/__init__.py

@@ -8,7 +8,7 @@ from prefect.deployments.base import (
 from prefect.deployments.runner import (
     RunnerDeployment,
     deploy,
-    DeploymentImage,
+    DockerImage,
     EntrypointType,
 )