prefect-client 2.19.2__py3-none-any.whl → 3.0.0rc1__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.
@@ -10,7 +10,7 @@ import os
 import sys
 import warnings
 from collections import defaultdict
-from contextlib import contextmanager
+from contextlib import ExitStack, contextmanager
 from contextvars import ContextVar, Token
 from functools import update_wrapper
 from pathlib import Path
@@ -29,21 +29,19 @@ from typing import (
     Union,
 )
 
+import anyio
+import anyio._backends._asyncio
 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 sniffio import AsyncLibraryNotFoundError
+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.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
@@ -51,7 +49,8 @@ from prefect.futures import PrefectFuture
 from prefect.results import ResultFactory
 from prefect.settings import PREFECT_HOME, Profile, Settings
 from prefect.states import State
-from prefect.task_runners import BaseTaskRunner
+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")
@@ -63,7 +62,71 @@ 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
+            # TODO: This task group isn't necessary in the new engine. Remove the background tasks
+            # attribute from FlowRunContext.
+            client = client or get_client(sync_client=True)
+            if flow_run_context := serialized_context.get("flow_run_context"):
+                try:
+                    task_group = anyio.create_task_group()
+                except AsyncLibraryNotFoundError:
+                    task_group = anyio._backends._asyncio.TaskGroup()
+                flow = flow_run_context["flow"]
+                flow_run_context = FlowRunContext(
+                    **flow_run_context,
+                    client=client,
+                    background_tasks=task_group,
+                    result_factory=run_coro_as_sync(ResultFactory.from_flow(flow)),
+                    task_runner=flow.task_runner.duplicate(),
+                    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"):
+                parent_task = parent_task_run_context["task"]
+                task_run_context = TaskRunContext(
+                    **parent_task_run_context,
+                    client=client,
+                    result_factory=run_coro_as_sync(
+                        ResultFactory.from_autonomous_task(parent_task)
+                    ),
+                )
+                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,12 +137,11 @@ 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):
         if self._token is not None:
@@ -98,10 +160,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,11 +180,17 @@ 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):
     """
@@ -132,7 +203,7 @@ class PrefectObjectRegistry(ContextModel):
         capture_failures: If set, failures during __init__ will be silenced and tracked.
     """
 
-    start_time: DateTimeTZ = Field(default_factory=lambda: pendulum.now("UTC"))
+    start_time: DateTime = Field(default_factory=lambda: pendulum.now("UTC"))
 
     _instance_registry: Dict[Type[T], List[T]] = PrivateAttr(
         default_factory=lambda: defaultdict(list)
@@ -201,6 +272,67 @@ class PrefectObjectRegistry(ContextModel):
         return type_
 
 
+class ClientContext(ContextModel):
+    """
+    A context for managing the 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 clients.
+
+    The client creates both sync and async clients, which can either be read
+    directly from the context object OR loaded with get_client, inject_client,
+    or other Prefect utilities.
+
+    with ClientContext.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.sync_client
+    """
+
+    __var__ = ContextVar("clients")
+    sync_client: SyncPrefectClient
+    async_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__(
+            sync_client=get_client(sync_client=True, httpx_settings=httpx_settings),
+            async_client=get_client(sync_client=False, 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.sync_client.__enter__()
+            run_coro_as_sync(self.async_client.__aenter__())
+            return super().__enter__()
+        else:
+            return self
+
+    def __exit__(self, *exc_info):
+        self._context_stack -= 1
+        if self._context_stack == 0:
+            self.sync_client.__exit__(*exc_info)
+            run_coro_as_sync(self.async_client.__aexit__(*exc_info))
+            return super().__exit__(*exc_info)
+
+    @classmethod
+    @contextmanager
+    def get_or_create(cls) -> Generator["ClientContext", None, None]:
+        ctx = ClientContext.get()
+        if ctx:
+            yield ctx
+        else:
+            with ClientContext() as ctx:
+                yield ctx
+
+
 class RunContext(ContextModel):
     """
     The base context for a flow or task run. Data in this context will always be
@@ -211,10 +343,16 @@ class RunContext(ContextModel):
         client: The Prefect client instance being used for API communication
     """
 
-    start_time: DateTimeTZ = Field(default_factory=lambda: pendulum.now("UTC"))
+    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):
     """
@@ -236,10 +374,14 @@ class EngineContext(RunContext):
     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
 
@@ -266,7 +408,20 @@ class EngineContext(RunContext):
     # Events worker to emit events to Prefect Cloud
     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",
+            },
+            exclude_unset=True,
+        )
 
 
 FlowRunContext = EngineContext  # for backwards compatibility
@@ -292,6 +447,19 @@ class TaskRunContext(RunContext):
 
     __var__ = ContextVar("task_run")
 
+    def serialize(self):
+        return self.model_dump(
+            include={
+                "task_run",
+                "task",
+                "parameters",
+                "log_prints",
+                "start_time",
+                "input_keyset",
+            },
+            exclude_unset=True,
+        )
+
 
 class TagsContext(ContextModel):
     """
@@ -569,7 +737,7 @@ def root_settings_context():
 
 
 GLOBAL_SETTINGS_CONTEXT: SettingsContext = root_settings_context()
-GLOBAL_OBJECT_REGISTRY: ContextManager[PrefectObjectRegistry] = None
+GLOBAL_OBJECT_REGISTRY: Optional[ContextManager[PrefectObjectRegistry]] = None
 
 
 def initialize_object_registry():

prefect/deployments/__init__.py

@@ -4,15 +4,14 @@ 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 prefect.deployments.flow_runs import (
+    run_deployment,
+)

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
@@ -260,7 +260,7 @@ def _format_deployment_for_saving_to_prefect_file(
         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()
+            deployment["schedule"] = deployment["schedule"].model_dump()
 
         if "is_schedule_active" in deployment:
             deployment["schedule"]["active"] = deployment.pop("is_schedule_active")
@@ -268,16 +268,18 @@ def _format_deployment_for_saving_to_prefect_file(
     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 +297,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()
 

prefect/deployments/flow_runs.py

@@ -0,0 +1,185 @@
+from datetime import datetime
+from typing import TYPE_CHECKING, Iterable, Optional, Union
+from uuid import UUID
+
+import anyio
+import pendulum
+
+from prefect._internal.compatibility.deprecated import deprecated_parameter
+from prefect.client.schemas import FlowRun
+from prefect.client.utilities import inject_client
+from prefect.context import FlowRunContext, TaskRunContext
+from prefect.logging import get_logger
+from prefect.states import Pending, Scheduled
+from prefect.tasks import Task
+from prefect.utilities.asyncutils import sync_compatible
+from prefect.utilities.slugify import slugify
+
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+    from prefect.client.schemas.objects import FlowRun
+
+
+logger = get_logger(__name__)
+
+
+@sync_compatible
+@deprecated_parameter(
+    "infra_overrides",
+    start_date="Apr 2024",
+    help="Use `job_variables` instead.",
+)
+@inject_client
+async def run_deployment(
+    name: Union[str, UUID],
+    client: Optional["PrefectClient"] = None,
+    parameters: Optional[dict] = None,
+    scheduled_time: Optional[datetime] = None,
+    flow_run_name: Optional[str] = None,
+    timeout: Optional[float] = None,
+    poll_interval: Optional[float] = 5,
+    tags: Optional[Iterable[str]] = None,
+    idempotency_key: Optional[str] = None,
+    work_queue_name: Optional[str] = None,
+    as_subflow: Optional[bool] = True,
+    infra_overrides: Optional[dict] = None,
+    job_variables: Optional[dict] = None,
+) -> "FlowRun":
+    """
+    Create a flow run for a deployment and return it after completion or a timeout.
+
+    By default, this function blocks until the flow run finishes executing.
+    Specify a timeout (in seconds) to wait for the flow run to execute before
+    returning flow run metadata. To return immediately, without waiting for the
+    flow run to execute, set `timeout=0`.
+
+    Note that if you specify a timeout, this function will return the flow run
+    metadata whether or not the flow run finished executing.
+
+    If called within a flow or task, the flow run this function creates will
+    be linked to the current flow run as a subflow. Disable this behavior by
+    passing `as_subflow=False`.
+
+    Args:
+        name: The deployment id or deployment name in the form:
+            `"flow name/deployment name"`
+        parameters: Parameter overrides for this flow run. Merged with the deployment
+            defaults.
+        scheduled_time: The time to schedule the flow run for, defaults to scheduling
+            the flow run to start now.
+        flow_run_name: A name for the created flow run
+        timeout: The amount of time to wait (in seconds) for the flow run to
+            complete before returning. Setting `timeout` to 0 will return the flow
+            run metadata immediately. Setting `timeout` to None will allow this
+            function to poll indefinitely. Defaults to None.
+        poll_interval: The number of seconds between polls
+        tags: A list of tags to associate with this flow run; tags can be used in
+            automations and for organizational purposes.
+        idempotency_key: A unique value to recognize retries of the same run, and
+            prevent creating multiple flow runs.
+        work_queue_name: The name of a work queue to use for this run. Defaults to
+            the default work queue for the deployment.
+        as_subflow: Whether to link the flow run as a subflow of the current
+            flow or task run.
+        job_variables: A dictionary of dot delimited infrastructure overrides that
+            will be applied at runtime; for example `env.CONFIG_KEY=config_value` or
+            `namespace='prefect'`
+    """
+    if timeout is not None and timeout < 0:
+        raise ValueError("`timeout` cannot be negative")
+
+    if scheduled_time is None:
+        scheduled_time = pendulum.now("UTC")
+
+    parameters = parameters or {}
+
+    deployment_id = None
+
+    if isinstance(name, UUID):
+        deployment_id = name
+    else:
+        try:
+            deployment_id = UUID(name)
+        except ValueError:
+            pass
+
+    if deployment_id:
+        deployment = await client.read_deployment(deployment_id=deployment_id)
+    else:
+        deployment = await client.read_deployment_by_name(name)
+
+    flow_run_ctx = FlowRunContext.get()
+    task_run_ctx = TaskRunContext.get()
+    if as_subflow and (flow_run_ctx or task_run_ctx):
+        # TODO: this logic can likely be simplified by using `Task.create_run`
+        from prefect.utilities.engine import (
+            _dynamic_key_for_task_run,
+            collect_task_run_inputs,
+        )
+
+        # This was called from a flow. Link the flow run as a subflow.
+        task_inputs = {
+            k: await collect_task_run_inputs(v) for k, v in parameters.items()
+        }
+
+        if deployment_id:
+            flow = await client.read_flow(deployment.flow_id)
+            deployment_name = f"{flow.name}/{deployment.name}"
+        else:
+            deployment_name = name
+
+        # Generate a task in the parent flow run to represent the result of the subflow
+        dummy_task = Task(
+            name=deployment_name,
+            fn=lambda: None,
+            version=deployment.version,
+        )
+        # Override the default task key to include the deployment name
+        dummy_task.task_key = f"{__name__}.run_deployment.{slugify(deployment_name)}"
+        flow_run_id = (
+            flow_run_ctx.flow_run.id
+            if flow_run_ctx
+            else task_run_ctx.task_run.flow_run_id
+        )
+        dynamic_key = (
+            _dynamic_key_for_task_run(flow_run_ctx, dummy_task)
+            if flow_run_ctx
+            else task_run_ctx.task_run.dynamic_key
+        )
+        parent_task_run = await client.create_task_run(
+            task=dummy_task,
+            flow_run_id=flow_run_id,
+            dynamic_key=dynamic_key,
+            task_inputs=task_inputs,
+            state=Pending(),
+        )
+        parent_task_run_id = parent_task_run.id
+    else:
+        parent_task_run_id = None
+
+    flow_run = await client.create_flow_run_from_deployment(
+        deployment.id,
+        parameters=parameters,
+        state=Scheduled(scheduled_time=scheduled_time),
+        name=flow_run_name,
+        tags=tags,
+        idempotency_key=idempotency_key,
+        parent_task_run_id=parent_task_run_id,
+        work_queue_name=work_queue_name,
+        job_variables=job_variables,
+    )
+
+    flow_run_id = flow_run.id
+
+    if timeout == 0:
+        return flow_run
+
+    with anyio.move_on_after(timeout):
+        while True:
+            flow_run = await client.read_flow_run(flow_run_id)
+            flow_state = flow_run.state
+            if flow_state and flow_state.is_final():
+                return flow_run
+            await anyio.sleep(poll_interval)
+
+    return flow_run