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

prefect/context.py

@@ -1,4 +1,4 @@
-r"""
+"""
 Async and thread safe models for passing runtime context data.
 
 These contexts should never be directly mutated by the user.
@@ -9,50 +9,42 @@ 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 contextmanager
+import weakref
+from contextlib import ExitStack, asynccontextmanager, contextmanager
 from contextvars import ContextVar, Token
-from functools import update_wrapper
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
-    ContextManager,
+    AsyncGenerator,
     Dict,
     Generator,
-    List,
+    Mapping,
     Optional,
     Set,
-    Tuple,
     Type,
     TypeVar,
     Union,
 )
 
-import anyio.abc
 import pendulum
-
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
-
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import BaseModel, Field, PrivateAttr
-else:
-    from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+from pydantic_extra_types.pendulum_dt import DateTime
+from typing_extensions import Self
 
 import prefect.logging
 import prefect.logging.configuration
 import prefect.settings
-from prefect._internal.schemas.fields import DateTimeTZ
-from prefect.client.orchestration import PrefectClient, SyncPrefectClient
+from prefect._internal.compatibility.migration import getattr_migration
+from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.events.worker import EventsWorker
 from prefect.exceptions import MissingContextError
-from prefect.futures import PrefectFuture
-from prefect.results import ResultFactory
+from prefect.results import ResultStore
 from prefect.settings import PREFECT_HOME, Profile, Settings
 from prefect.states import State
-from prefect.task_runners import BaseTaskRunner
-from prefect.utilities.importtools import load_script_as_module
+from prefect.task_runners import TaskRunner
+from prefect.utilities.services import start_client_metrics_server
 
 T = TypeVar("T")
 
@@ -63,7 +55,60 @@ if TYPE_CHECKING:
 # Define the global settings context variable
 # This will be populated downstream but must be null here to facilitate loading the
 # default settings.
-GLOBAL_SETTINGS_CONTEXT = None
+GLOBAL_SETTINGS_CONTEXT = None  # type: ignore
+
+
+def serialize_context() -> Dict[str, Any]:
+    """
+    Serialize the current context for use in a remote execution environment.
+    """
+
+    flow_run_context = EngineContext.get()
+    task_run_context = TaskRunContext.get()
+    tags_context = TagsContext.get()
+    settings_context = SettingsContext.get()
+
+    return {
+        "flow_run_context": flow_run_context.serialize() if flow_run_context else {},
+        "task_run_context": task_run_context.serialize() if task_run_context else {},
+        "tags_context": tags_context.serialize() if tags_context else {},
+        "settings_context": settings_context.serialize() if settings_context else {},
+    }
+
+
+@contextmanager
+def hydrated_context(
+    serialized_context: Optional[Dict[str, Any]] = None,
+    client: Union[PrefectClient, SyncPrefectClient, None] = None,
+):
+    with ExitStack() as stack:
+        if serialized_context:
+            # Set up settings context
+            if settings_context := serialized_context.get("settings_context"):
+                stack.enter_context(SettingsContext(**settings_context))
+            # Set up parent flow run context
+            client = client or get_client(sync_client=True)
+            if flow_run_context := serialized_context.get("flow_run_context"):
+                flow = flow_run_context["flow"]
+                task_runner = stack.enter_context(flow.task_runner.duplicate())
+                flow_run_context = FlowRunContext(
+                    **flow_run_context,
+                    client=client,
+                    task_runner=task_runner,
+                    detached=True,
+                )
+                stack.enter_context(flow_run_context)
+            # Set up parent task run context
+            if parent_task_run_context := serialized_context.get("task_run_context"):
+                task_run_context = TaskRunContext(
+                    **parent_task_run_context,
+                    client=client,
+                )
+                stack.enter_context(task_run_context)
+            # Set up tags context
+            if tags_context := serialized_context.get("tags_context"):
+                stack.enter_context(tags(*tags_context["current_tags"]))
+        yield
 
 
 class ContextModel(BaseModel):
@@ -74,14 +119,13 @@ class ContextModel(BaseModel):
 
     # The context variable for storing data must be defined by the child class
     __var__: ContextVar
-    _token: Token = PrivateAttr(None)
-
-    class Config:
-        # allow_mutation = False
-        arbitrary_types_allowed = True
-        extra = "forbid"
+    _token: Optional[Token] = PrivateAttr(None)
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+    )
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         if self._token is not None:
             raise RuntimeError(
                 "Context already entered. Context enter calls cannot be nested."
@@ -98,10 +142,13 @@ class ContextModel(BaseModel):
         self._token = None
 
     @classmethod
-    def get(cls: Type[T]) -> Optional[T]:
+    def get(cls: Type[Self]) -> Optional[Self]:
+        """Get the current context instance"""
         return cls.__var__.get(None)
 
-    def copy(self, **kwargs):
+    def model_copy(
+        self: Self, *, update: Optional[Dict[str, Any]] = None, deep: bool = False
+    ):
         """
         Duplicate the context model, optionally choosing which fields to include, exclude, or change.
 
@@ -115,90 +162,130 @@ class ContextModel(BaseModel):
         Returns:
             A new model instance.
         """
+        new = super().model_copy(update=update, deep=deep)
         # Remove the token on copy to avoid re-entrance errors
-        new = super().copy(**kwargs)
         new._token = None
         return new
 
+    def serialize(self) -> Dict[str, Any]:
+        """
+        Serialize the context model to a dictionary that can be pickled with cloudpickle.
+        """
+        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.
+class SyncClientContext(ContextModel):
+    """
+    A context for managing the sync Prefect client instances.
+
+    Clients were formerly tracked on the TaskRunContext and FlowRunContext, but
+    having two separate places and the addition of both sync and async clients
+    made it difficult to manage. This context is intended to be the single
+    source for sync clients.
+
+    The client creates a sync client, which can either be read directly from
+    the context object OR loaded with get_client, inject_client, or other
+    Prefect utilities.
+
+    with SyncClientContext.get_or_create() as ctx:
+        c1 = get_client(sync_client=True)
+        c2 = get_client(sync_client=True)
+        assert c1 is c2
+        assert c1 is ctx.client
     """
 
-    start_time: DateTimeTZ = Field(default_factory=lambda: pendulum.now("UTC"))
-
-    _instance_registry: Dict[Type[T], List[T]] = PrivateAttr(
-        default_factory=lambda: defaultdict(list)
-    )
+    __var__ = ContextVar("sync-client-context")
+    client: SyncPrefectClient
+    _httpx_settings: Optional[dict[str, Any]] = PrivateAttr(None)
+    _context_stack: int = PrivateAttr(0)
 
-    # 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)
+    def __init__(self, httpx_settings: Optional[dict[str, Any]] = None):
+        super().__init__(
+            client=get_client(sync_client=True, httpx_settings=httpx_settings),
         )
+        self._httpx_settings = httpx_settings
+        self._context_stack = 0
+
+    def __enter__(self):
+        self._context_stack += 1
+        if self._context_stack == 1:
+            self.client.__enter__()
+            return super().__enter__()
+        else:
+            return self
+
+    def __exit__(self, *exc_info):
+        self._context_stack -= 1
+        if self._context_stack == 0:
+            self.client.__exit__(*exc_info)
+            return super().__exit__(*exc_info)
 
     @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__
+    @contextmanager
+    def get_or_create(cls) -> Generator["SyncClientContext", None, None]:
+        ctx = SyncClientContext.get()
+        if ctx:
+            yield ctx
+        else:
+            with SyncClientContext() as ctx:
+                yield ctx
 
-        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)
+class AsyncClientContext(ContextModel):
+    """
+    A context for managing the async Prefect client instances.
+
+    Clients were formerly tracked on the TaskRunContext and FlowRunContext, but
+    having two separate places and the addition of both sync and async clients
+    made it difficult to manage. This context is intended to be the single
+    source for async clients.
+
+    The client creates an async client, which can either be read directly from
+    the context object OR loaded with get_client, inject_client, or other
+    Prefect utilities.
+
+    with AsyncClientContext.get_or_create() as ctx:
+        c1 = get_client(sync_client=False)
+        c2 = get_client(sync_client=False)
+        assert c1 is c2
+        assert c1 is ctx.client
+    """
 
-        type_.__init__ = __register_init__
-        return type_
+    __var__ = ContextVar("async-client-context")
+    client: PrefectClient
+    _httpx_settings: Optional[dict[str, Any]] = PrivateAttr(None)
+    _context_stack: int = PrivateAttr(0)
+
+    def __init__(self, httpx_settings: Optional[dict[str, Any]] = None):
+        super().__init__(
+            client=get_client(sync_client=False, httpx_settings=httpx_settings),
+        )
+        self._httpx_settings = httpx_settings
+        self._context_stack = 0
+
+    async def __aenter__(self):
+        self._context_stack += 1
+        if self._context_stack == 1:
+            await self.client.__aenter__()
+            return super().__enter__()
+        else:
+            return self
+
+    async def __aexit__(self, *exc_info):
+        self._context_stack -= 1
+        if self._context_stack == 0:
+            await self.client.__aexit__(*exc_info)
+            return super().__exit__(*exc_info)
+
+    @classmethod
+    @asynccontextmanager
+    async def get_or_create(cls) -> AsyncGenerator[Self, None]:
+        ctx = cls.get()
+        if ctx:
+            yield ctx
+        else:
+            async with cls() as ctx:
+                yield ctx
 
 
 class RunContext(ContextModel):
@@ -211,10 +298,21 @@ class RunContext(ContextModel):
         client: The Prefect client instance being used for API communication
     """
 
-    start_time: DateTimeTZ = Field(default_factory=lambda: pendulum.now("UTC"))
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        start_client_metrics_server()
+
+    start_time: DateTime = Field(default_factory=lambda: pendulum.now("UTC"))
     input_keyset: Optional[Dict[str, Dict[str, str]]] = None
     client: Union[PrefectClient, SyncPrefectClient]
 
+    def serialize(self):
+        return self.model_dump(
+            include={"start_time", "input_keyset"},
+            exclude_unset=True,
+        )
+
 
 class EngineContext(RunContext):
     """
@@ -229,19 +327,20 @@ class EngineContext(RunContext):
         task_run_states: A list of states for task runs created within this flow run
         task_run_results: A mapping of result ids to task run states for this flow run
         flow_run_states: A list of states for flow runs created within this flow run
-        sync_portal: A blocking portal for sync task/flow runs in an async flow
-        timeout_scope: The cancellation scope for flow level timeouts
     """
 
     flow: Optional["Flow"] = None
     flow_run: Optional[FlowRun] = None
-    autonomous_task_run: Optional[TaskRun] = None
-    task_runner: BaseTaskRunner
+    task_runner: TaskRunner
     log_prints: bool = False
     parameters: Optional[Dict[str, Any]] = None
 
+    # Flag signaling if the flow run context has been serialized and sent
+    # to remote infrastructure.
+    detached: bool = False
+
     # Result handling
-    result_factory: ResultFactory
+    result_store: ResultStore
 
     # Counter for task calls allowing unique
     task_run_dynamic_keys: Dict[str, int] = Field(default_factory=dict)
@@ -249,24 +348,31 @@ class EngineContext(RunContext):
     # Counter for flow pauses
     observed_flow_pauses: Dict[str, int] = Field(default_factory=dict)
 
-    # Tracking for objects created by this flow run
-    task_run_futures: List[PrefectFuture] = Field(default_factory=list)
-    task_run_states: List[State] = Field(default_factory=list)
-    task_run_results: Dict[int, State] = Field(default_factory=dict)
-    flow_run_states: List[State] = Field(default_factory=list)
-
-    # The synchronous portal is only created for async flows for creating engine calls
-    # from synchronous task and subflow calls
-    sync_portal: Optional[anyio.abc.BlockingPortal] = None
-    timeout_scope: Optional[anyio.abc.CancelScope] = None
-
-    # Task group that can be used for background tasks during the flow run
-    background_tasks: anyio.abc.TaskGroup
+    # 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 worker to emit events
     events: Optional[EventsWorker] = None
 
-    __var__ = ContextVar("flow_run")
+    __var__: ContextVar = ContextVar("flow_run")
+
+    def serialize(self):
+        return self.model_dump(
+            include={
+                "flow_run",
+                "flow",
+                "parameters",
+                "log_prints",
+                "start_time",
+                "input_keyset",
+                "result_store",
+            },
+            exclude_unset=True,
+        )
 
 
 FlowRunContext = EngineContext  # for backwards compatibility
@@ -288,10 +394,24 @@ class TaskRunContext(RunContext):
     parameters: Dict[str, Any]
 
     # Result handling
-    result_factory: ResultFactory
+    result_store: ResultStore
 
     __var__ = ContextVar("task_run")
 
+    def serialize(self):
+        return self.model_dump(
+            include={
+                "task_run",
+                "task",
+                "parameters",
+                "log_prints",
+                "start_time",
+                "input_keyset",
+                "result_store",
+            },
+            exclude_unset=True,
+        )
+
 
 class TagsContext(ContextModel):
     """
@@ -452,23 +572,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],
@@ -555,7 +658,7 @@ def root_settings_context():
             ),
             file=sys.stderr,
         )
-        active_name = "default"
+        active_name = "ephemeral"
 
     with use_profile(
         profiles[active_name],
@@ -569,14 +672,8 @@ def root_settings_context():
 
 
 GLOBAL_SETTINGS_CONTEXT: SettingsContext = root_settings_context()
-GLOBAL_OBJECT_REGISTRY: ContextManager[PrefectObjectRegistry] = None
-
-
-def initialize_object_registry():
-    global GLOBAL_OBJECT_REGISTRY
 
-    if GLOBAL_OBJECT_REGISTRY:
-        return
 
-    GLOBAL_OBJECT_REGISTRY = PrefectObjectRegistry()
-    GLOBAL_OBJECT_REGISTRY.__enter__()
+# 2024-07-02: This surfaces an actionable error message for removed objects
+# in Prefect 3.0 upgrade.
+__getattr__ = getattr_migration(__name__)

prefect/deployments/__init__.py

@@ -1,18 +1,33 @@
-import prefect.deployments.base
-import prefect.deployments.steps
-from prefect.deployments.base import (
-    initialize_project,
-)
-
-from prefect.deployments.deployments import (
-    run_deployment,
-    load_flow_from_flow_run,
-    load_deployments_from_yaml,
-    Deployment,
-)
-from prefect.deployments.runner import (
-    RunnerDeployment,
-    deploy,
-    DeploymentImage,
-    EntrypointType,
-)
+from typing import TYPE_CHECKING
+from prefect._internal.compatibility.migration import getattr_migration
+
+
+if TYPE_CHECKING:
+    from .flow_runs import run_deployment
+    from .base import initialize_project
+    from .runner import deploy
+
+_public_api: dict[str, tuple[str, str]] = {
+    "initialize_project": (__spec__.parent, ".base"),
+    "run_deployment": (__spec__.parent, ".flow_runs"),
+    "deploy": (__spec__.parent, ".runner"),
+}
+
+# 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/base.py

@@ -19,7 +19,7 @@ import anyio
 import yaml
 from ruamel.yaml import YAML
 
-from prefect.client.schemas.objects import MinimalDeploymentSchedule
+from prefect.client.schemas.actions import DeploymentScheduleCreate
 from prefect.client.schemas.schedules import IntervalSchedule
 from prefect.logging import get_logger
 from prefect.settings import PREFECT_DEBUG_MODE
@@ -29,7 +29,7 @@ from prefect.utilities.templating import apply_values
 
 
 def create_default_prefect_yaml(
-    path: str, name: str = None, contents: Optional[Dict[str, Any]] = None
+    path: str, name: Optional[str] = None, contents: Optional[Dict[str, Any]] = None
 ) -> bool:
     """
     Creates default `prefect.yaml` file in the provided path if one does not already exist;
@@ -176,7 +176,9 @@ def _get_git_branch() -> Optional[str]:
 
 
 def initialize_project(
-    name: str = None, recipe: str = None, inputs: Optional[Dict[str, Any]] = None
+    name: Optional[str] = None,
+    recipe: Optional[str] = None,
+    inputs: Optional[Dict[str, Any]] = None,
 ) -> List[str]:
     """
     Initializes a basic project structure with base files.  If no name is provided, the name
@@ -256,28 +258,21 @@ def _format_deployment_for_saving_to_prefect_file(
     # Only want entrypoint to avoid errors
     deployment.pop("flow_name", None)
 
-    if deployment.get("schedule"):
-        if isinstance(deployment["schedule"], IntervalSchedule):
-            deployment["schedule"] = _interval_schedule_to_dict(deployment["schedule"])
-        else:  # all valid SCHEDULE_TYPES are subclasses of BaseModel
-            deployment["schedule"] = deployment["schedule"].dict()
-
-        if "is_schedule_active" in deployment:
-            deployment["schedule"]["active"] = deployment.pop("is_schedule_active")
-
     if deployment.get("schedules"):
         schedules = []
         for deployment_schedule in cast(
-            List[MinimalDeploymentSchedule], deployment["schedules"]
+            List[DeploymentScheduleCreate], deployment["schedules"]
         ):
             if isinstance(deployment_schedule.schedule, IntervalSchedule):
                 schedule_config = _interval_schedule_to_dict(
                     deployment_schedule.schedule
                 )
             else:  # all valid SCHEDULE_TYPES are subclasses of BaseModel
-                schedule_config = deployment_schedule.schedule.dict()
+                schedule_config = deployment_schedule.schedule.model_dump()
 
             schedule_config["active"] = deployment_schedule.active
+            schedule_config["max_active_runs"] = deployment_schedule.max_active_runs
+            schedule_config["catchup"] = deployment_schedule.catchup
             schedules.append(schedule_config)
 
         deployment["schedules"] = schedules
@@ -295,7 +290,7 @@ def _interval_schedule_to_dict(schedule: IntervalSchedule) -> Dict:
     Returns:
         - Dict: the schedule as a dictionary
     """
-    schedule_config = schedule.dict()
+    schedule_config = schedule.model_dump()
     schedule_config["interval"] = schedule_config["interval"].total_seconds()
     schedule_config["anchor_date"] = schedule_config["anchor_date"].isoformat()