prefect-client 2.19.4__py3-none-any.whl → 3.0.0rc2__py3-none-any.whl

prefect/infrastructure/provisioners/cloud_run.py

@@ -5,7 +5,7 @@ import tempfile
 from copy import deepcopy
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 from uuid import UUID
 
 from anyio import run_process
@@ -17,7 +17,7 @@ from rich.prompt import Confirm
 from rich.syntax import Syntax
 
 from prefect.cli._prompts import prompt, prompt_select_from_table
-from prefect.client.orchestration import PrefectClient, ServerType
+from prefect.client.orchestration import ServerType
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectAlreadyExists
@@ -27,6 +27,9 @@ from prefect.settings import (
     update_current_profile,
 )
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 
 class CloudRunPushProvisioner:
     def __init__(self):
@@ -215,7 +218,7 @@ class CloudRunPushProvisioner:
             ) from e
 
     async def _create_gcp_credentials_block(
-        self, block_document_name: str, key: dict, client: PrefectClient
+        self, block_document_name: str, key: dict, client: "PrefectClient"
     ) -> UUID:
         credentials_block_type = await client.read_block_type_by_slug("gcp-credentials")
 
@@ -242,7 +245,9 @@ class CloudRunPushProvisioner:
             )
             return block_doc.id
 
-    async def _create_provision_table(self, work_pool_name: str, client: PrefectClient):
+    async def _create_provision_table(
+        self, work_pool_name: str, client: "PrefectClient"
+    ):
         return Panel(
             dedent(
                 f"""\
@@ -268,7 +273,7 @@ class CloudRunPushProvisioner:
         )
 
     async def _customize_resource_names(
-        self, work_pool_name: str, client: PrefectClient
+        self, work_pool_name: str, client: "PrefectClient"
     ) -> bool:
         self._service_account_name = prompt(
             "Please enter a name for the service account",
@@ -294,7 +299,7 @@ class CloudRunPushProvisioner:
         self,
         work_pool_name: str,
         base_job_template: dict,
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ) -> Dict[str, Any]:
         assert client, "Client injection failed"
         await self._verify_gcloud_ready()

prefect/infrastructure/provisioners/container_instance.py

@@ -18,7 +18,7 @@ import subprocess
 import time
 from copy import deepcopy
 from textwrap import dedent
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 from uuid import UUID
 
 from anyio import run_process
@@ -29,7 +29,6 @@ from rich.prompt import Confirm
 from rich.syntax import Syntax
 
 from prefect.cli._prompts import prompt, prompt_select_from_table
-from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectAlreadyExists, ObjectNotFound
@@ -38,6 +37,9 @@ from prefect.settings import (
     update_current_profile,
 )
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 
 class AzureCLI:
     """
@@ -684,7 +686,7 @@ class ContainerInstancePushProvisioner:
         client_id: str,
         tenant_id: str,
         client_secret: str,
-        client: PrefectClient,
+        client: "PrefectClient",
     ) -> UUID:
         """
         Creates a credentials block for Azure Container Instance.
@@ -755,7 +757,7 @@ class ContainerInstancePushProvisioner:
             raise e
 
     async def _aci_credentials_block_exists(
-        self, block_name: str, client: PrefectClient
+        self, block_name: str, client: "PrefectClient"
     ) -> bool:
         """
         Checks if an ACI credentials block with the given name already exists.
@@ -782,7 +784,9 @@ class ContainerInstancePushProvisioner:
         else:
             return False
 
-    async def _create_provision_table(self, work_pool_name: str, client: PrefectClient):
+    async def _create_provision_table(
+        self, work_pool_name: str, client: "PrefectClient"
+    ):
         return Panel(
             dedent(
                 f"""\
@@ -809,7 +813,7 @@ class ContainerInstancePushProvisioner:
         )
 
     async def _customize_resource_names(
-        self, work_pool_name: str, client: PrefectClient
+        self, work_pool_name: str, client: "PrefectClient"
     ) -> bool:
         self._resource_group_name = prompt(
             "Please enter a name for the resource group",
@@ -851,7 +855,7 @@ class ContainerInstancePushProvisioner:
         self,
         work_pool_name: str,
         base_job_template: Dict[str, Any],
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ) -> Dict[str, Any]:
         """
         Orchestrates the provisioning of Azure resources and setup for the push work pool.

prefect/infrastructure/provisioners/ecs.py

@@ -9,7 +9,7 @@ import sys
 from copy import deepcopy
 from functools import partial
 from textwrap import dedent
-from typing import Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
 
 import anyio
 from anyio import run_process
@@ -20,7 +20,6 @@ from rich.prompt import Confirm
 from rich.syntax import Syntax
 
 from prefect.cli._prompts import prompt
-from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectNotFound
@@ -31,6 +30,9 @@ from prefect.settings import (
 from prefect.utilities.collections import get_from_dict
 from prefect.utilities.importtools import lazy_import
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
 boto3 = lazy_import("boto3")
 
 current_console = contextvars.ContextVar("console", default=Console())
@@ -250,7 +252,7 @@ class CredentialsBlockResource:
 
     @inject_client
     async def requires_provisioning(
-        self, client: Optional[PrefectClient] = None
+        self, client: Optional["PrefectClient"] = None
     ) -> bool:
         if self._requires_provisioning is None:
             try:
@@ -280,7 +282,7 @@ class CredentialsBlockResource:
         self,
         base_job_template: Dict[str, Any],
         advance: Callable[[], None],
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ):
         """
         Provisions an AWS credentials block.

prefect/infrastructure/provisioners/modal.py

@@ -2,20 +2,23 @@ import importlib
 import shlex
 import sys
 from copy import deepcopy
-from typing import Any, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
 
 from anyio import run_process
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from rich.prompt import Confirm
 
-from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.schemas.objects import BlockDocument
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectNotFound
 from prefect.utilities.importtools import lazy_import
 
+if TYPE_CHECKING:
+    from prefect.client.orchestration import PrefectClient
+
+
 modal = lazy_import("modal")
 
 
@@ -24,7 +27,7 @@ class ModalPushProvisioner:
     A infrastructure provisioner for Modal push work pools.
     """
 
-    def __init__(self, client: Optional[PrefectClient] = None):
+    def __init__(self, client: Optional["PrefectClient"] = None):
         self._console = Console()
 
     @property
@@ -91,7 +94,7 @@ class ModalPushProvisioner:
         block_document_name: str,
         modal_token_id: str,
         modal_token_secret: str,
-        client: PrefectClient,
+        client: "PrefectClient",
     ) -> BlockDocument:
         """
         Creates a ModalCredentials block containing the provided token ID and secret.
@@ -142,7 +145,7 @@ class ModalPushProvisioner:
         self,
         work_pool_name: str,
         base_job_template: Dict[str, Any],
-        client: Optional[PrefectClient] = None,
+        client: Optional["PrefectClient"] = None,
     ) -> Dict[str, Any]:
         """
         Provisions resources necessary for a Modal push work pool.

prefect/input/actions.py

@@ -4,7 +4,6 @@ from uuid import UUID
 import orjson
 import pydantic
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
 from prefect.client.utilities import client_injector
 from prefect.context import FlowRunContext
 from prefect.exceptions import PrefectHTTPStatusError
@@ -14,8 +13,7 @@ if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
 
 
-if HAS_PYDANTIC_V2:
-    from prefect._internal.pydantic.v2_schema import is_v2_model
+from prefect._internal.pydantic.v2_schema import is_v2_model
 
 
 def ensure_flow_run_id(flow_run_id: Optional[UUID] = None) -> UUID:
@@ -41,7 +39,7 @@ async def create_flow_run_input_from_model(
         if context is not None and context.flow_run is not None:
             sender = f"prefect.flow-run.{context.flow_run.id}"
 
-    if HAS_PYDANTIC_V2 and is_v2_model(model_instance):
+    if is_v2_model(model_instance):
         json_safe = orjson.loads(model_instance.model_dump_json())
     else:
         json_safe = orjson.loads(model_instance.json())

prefect/input/run_input.py

@@ -18,7 +18,8 @@ Sender flow:
 ```python
 import random
 from uuid import UUID
-from prefect import flow, get_run_logger
+from prefect import flow
+from prefect.logging import get_run_logger
 from prefect.input import RunInput
 
 class NumberData(RunInput):
@@ -43,7 +44,8 @@ Receiver flow:
 ```python
 import random
 from uuid import UUID
-from prefect import flow, get_run_logger
+from prefect import flow
+from prefect.logging import get_run_logger
 from prefect.input import RunInput
 
 class NumberData(RunInput):
@@ -77,8 +79,8 @@ from uuid import UUID, uuid4
 
 import anyio
 import pydantic
+from pydantic import ConfigDict
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
 from prefect.input.actions import (
     create_flow_run_input,
     create_flow_run_input_from_model,
@@ -92,8 +94,7 @@ if TYPE_CHECKING:
     from prefect.client.schemas.objects import FlowRunInput
     from prefect.states import State
 
-if HAS_PYDANTIC_V2:
-    from prefect._internal.pydantic.v2_schema import create_v2_schema
+from prefect._internal.pydantic.v2_schema import create_v2_schema, is_v2_model
 
 R = TypeVar("R", bound="RunInput")
 T = TypeVar("T", bound="object")
@@ -138,13 +139,12 @@ def keyset_from_base_key(base_key: str) -> Keyset:
 
 class RunInputMetadata(pydantic.BaseModel):
     key: str
-    sender: Optional[str]
+    sender: Optional[str] = None
     receiver: UUID
 
 
 class RunInput(pydantic.BaseModel):
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
     _description: Optional[str] = pydantic.PrivateAttr(default=None)
     _metadata: RunInputMetadata = pydantic.PrivateAttr()
@@ -168,7 +168,7 @@ class RunInput(pydantic.BaseModel):
             - flow_run_id (UUID, optional): the flow run ID to save the input for
         """
 
-        if HAS_PYDANTIC_V2:
+        if is_v2_model(cls):
             schema = create_v2_schema(cls.__name__, model_base=cls)
         else:
             schema = cls.schema(by_alias=True)

prefect/logging/formatters.py

@@ -48,7 +48,6 @@ class JsonFormatter(logging.Formatter):
 
         self.serializer = JSONSerializer(
             jsonlib="orjson",
-            object_encoder="pydantic.json.pydantic_encoder",
             dumps_kwargs={"option": orjson.OPT_INDENT_2} if fmt == "pretty" else {},
         )
 
@@ -98,7 +97,6 @@ class PrefectFormatter(logging.Formatter):
             init_kwargs["defaults"] = defaults
             style_kwargs["defaults"] = defaults
 
-        # validate added in 3.8
         init_kwargs["validate"] = validate
 
         super().__init__(format, datefmt, style, **init_kwargs)

prefect/logging/handlers.py

@@ -119,21 +119,13 @@ class APILogHandler(logging.Handler):
             return APILogWorker.drain_all(timeout=5)
 
     @classmethod
-    def aflush(cls):
+    async def aflush(cls):
         """
         Tell the `APILogWorker` to send any currently enqueued logs and block until
         completion.
-
-        If called in a synchronous context, will only block up to 5s before returning.
         """
 
-        if not get_running_loop():
-            raise RuntimeError(
-                "`aflush` cannot be used from a synchronous context; use `flush`"
-                " instead."
-            )
-
-        return APILogWorker.drain_all()
+        return await APILogWorker.drain_all()
 
     def emit(self, record: logging.LogRecord):
         """
@@ -229,7 +221,7 @@ class APILogHandler(logging.Handler):
                 getattr(record, "created", None) or time.time()
             ),
             message=self.format(record),
-        ).dict(json_compatible=True)
+        ).model_dump(mode="json")
 
         log_size = log["__payload_size__"] = self._get_payload_size(log)
         if log_size > PREFECT_LOGGING_TO_API_MAX_LOG_SIZE.value():

prefect/logging/loggers.py

@@ -5,7 +5,7 @@ import warnings
 from builtins import print
 from contextlib import contextmanager
 from functools import lru_cache
-from logging import LogRecord
+from logging import LoggerAdapter, LogRecord
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from typing_extensions import Self
@@ -161,7 +161,7 @@ def flow_run_logger(
     flow_run: Union["FlowRun", "ClientFlowRun"],
     flow: Optional["Flow"] = None,
     **kwargs: str,
-):
+) -> LoggerAdapter:
     """
     Create a flow run logger with the run's metadata attached.
 

prefect/manifests.py

@@ -4,7 +4,8 @@ Manifests are portable descriptions of one or more workflows within a given dire
 They are the foundational building blocks for defining Flow Deployments.
 """
 
-from prefect.pydantic import BaseModel, Field
+from pydantic import BaseModel, Field
+
 from prefect.utilities.callables import ParameterSchema
 
 

prefect/records/__init__.py

@@ -0,0 +1 @@
+from .store import RecordStore

prefect/records/cache_policies.py

@@ -0,0 +1,179 @@
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional
+
+from prefect.context import TaskRunContext
+from prefect.utilities.hashing import hash_objects
+
+
+@dataclass
+class CachePolicy:
+    @classmethod
+    def from_cache_key_fn(
+        cls, cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
+    ) -> "CacheKeyFnPolicy":
+        """
+        Given a function generates a key policy.
+        """
+        return CacheKeyFnPolicy(cache_key_fn=cache_key_fn)
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        raise NotImplementedError
+
+    def __sub__(self, other: str) -> "CompoundCachePolicy":
+        if not isinstance(other, str):
+            raise TypeError("Can only subtract strings from key policies.")
+        if isinstance(self, Inputs):
+            exclude = self.exclude or []
+            return Inputs(exclude=exclude + [other])
+        elif isinstance(self, CompoundCachePolicy):
+            new = Inputs(exclude=[other])
+            policies = self.policies or []
+            return CompoundCachePolicy(policies=policies + [new])
+        else:
+            new = Inputs(exclude=[other])
+            return CompoundCachePolicy(policies=[self, new])
+
+    def __add__(self, other: "CachePolicy") -> "CompoundCachePolicy":
+        # adding _None is a no-op
+        if isinstance(other, _None):
+            return self
+        elif isinstance(self, _None):
+            return other
+
+        if isinstance(self, CompoundCachePolicy):
+            policies = self.policies or []
+            return CompoundCachePolicy(policies=policies + [other])
+        elif isinstance(other, CompoundCachePolicy):
+            policies = other.policies or []
+            return CompoundCachePolicy(policies=policies + [self])
+        else:
+            return CompoundCachePolicy(policies=[self, other])
+
+
+@dataclass
+class CacheKeyFnPolicy(CachePolicy):
+    # making it optional for tests
+    cache_key_fn: Optional[
+        Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
+    ] = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        if self.cache_key_fn:
+            return self.cache_key_fn(task_ctx, inputs)
+
+
+@dataclass
+class CompoundCachePolicy(CachePolicy):
+    policies: list = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        keys = []
+        for policy in self.policies:
+            keys.append(
+                policy.compute_key(
+                    task_ctx=task_ctx,
+                    inputs=inputs,
+                    flow_parameters=flow_parameters,
+                    **kwargs,
+                )
+            )
+        return hash_objects(*keys)
+
+
+@dataclass
+class Default(CachePolicy):
+    "Execution run ID only"
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        return str(task_ctx.task_run.id)
+
+
+@dataclass
+class _None(CachePolicy):
+    "ignore key policies altogether, always run - prevents persistence"
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        return None
+
+
+@dataclass
+class TaskDef(CachePolicy):
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        lines = inspect.getsource(task_ctx.task)
+        return hash_objects(lines)
+
+
+@dataclass
+class FlowParameters(CachePolicy):
+    pass
+
+
+@dataclass
+class Inputs(CachePolicy):
+    """
+    Exposes flag for whether to include flow parameters as well.
+
+    And exclude/include config.
+    """
+
+    exclude: list = None
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: Dict[str, Any],
+        flow_parameters: Dict[str, Any],
+        **kwargs,
+    ) -> Optional[str]:
+        hashed_inputs = {}
+        inputs = inputs or {}
+        exclude = self.exclude or []
+
+        for key, val in inputs.items():
+            if key not in exclude:
+                hashed_inputs[key] = val
+
+        return hash_objects(hashed_inputs)
+
+
+DEFAULT = Default()
+INPUTS = Inputs()
+NONE = _None()
+TASKDEF = TaskDef()

prefect/records/result_store.py

@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from typing import Any
+
+from prefect.exceptions import ObjectNotFound
+from prefect.results import BaseResult, PersistedResult, ResultFactory
+from prefect.utilities.asyncutils import run_coro_as_sync
+
+from .store import RecordStore
+
+
+@dataclass
+class ResultFactoryStore(RecordStore):
+    result_factory: ResultFactory
+    cache: PersistedResult = None
+
+    def exists(self, key: str) -> bool:
+        try:
+            result = self.read(key)
+            result.get(_sync=True)
+            self.cache = result
+            return True
+        except (ObjectNotFound, ValueError):
+            return False
+
+    def read(self, key: str) -> BaseResult:
+        if self.cache:
+            return self.cache
+        try:
+            result = PersistedResult(
+                serializer_type=self.result_factory.serializer.type,
+                storage_block_id=self.result_factory.storage_block_id,
+                storage_key=key,
+            )
+            return result
+        except Exception:
+            # this is a bit of a bandaid for functionality
+            raise ValueError("Result could not be read")
+
+    def write(self, key: str, value: Any) -> BaseResult:
+        if isinstance(value, BaseResult):
+            return value
+        return run_coro_as_sync(self.result_factory.create_result(obj=value, key=key))

prefect/records/store.py

@@ -0,0 +1,9 @@
+class RecordStore:
+    def read(self, key: str):
+        raise NotImplementedError
+
+    def write(self, key: str, value: dict):
+        raise NotImplementedError
+
+    def exists(self, key: str) -> bool:
+        return False