prefect-client 2.19.2__py3-none-any.whl → 3.0.0rc1__py3-none-any.whl

prefect/results.py

@@ -16,13 +16,14 @@ from typing import (
 )
 from uuid import UUID
 
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, ValidationError
+from pydantic_core import PydanticUndefinedType
 from typing_extensions import ParamSpec, Self
 
 import prefect
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
 from prefect.blocks.core import Block
 from prefect.client.utilities import inject_client
-from prefect.exceptions import MissingResult
+from prefect.exceptions import MissingResult, ObjectAlreadyExists
 from prefect.filesystems import (
     LocalFileSystem,
     ReadableFileSystem,
@@ -41,13 +42,6 @@ from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import sync_compatible
 from prefect.utilities.pydantic import get_dispatch_key, lookup_type, register_base_type
 
-if HAS_PYDANTIC_V2:
-    import pydantic.v1 as pydantic
-
-else:
-    import pydantic
-
-
 if TYPE_CHECKING:
     from prefect import Flow, Task
     from prefect.client.orchestration import PrefectClient
@@ -84,14 +78,14 @@ _default_task_scheduling_storages: Dict[Tuple[str, str], WritableFileSystem] = {
 
 async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
     """
-    Generate a default file system for autonomous task parameter/result storage.
+    Generate a default file system for background task parameter/result storage.
     """
     default_storage_name, storage_path = cache_key = (
         PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK.value(),
         PREFECT_LOCAL_STORAGE_PATH.value(),
     )
 
-    async def get_storage():
+    async def get_storage() -> WritableFileSystem:
         try:
             return await Block.load(default_storage_name)
         except ValueError as e:
@@ -110,12 +104,14 @@ async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
 
         try:
             await block.save(name, overwrite=False)
-            return block
         except ValueError as e:
             if "already in use" not in str(e):
                 raise e
+        except ObjectAlreadyExists:
+            # Another client created the block before we reached this line
+            block = await Block.load(default_storage_name)
 
-        return await Block.load(default_storage_name)
+        return block
 
     try:
         return _default_task_scheduling_storages[cache_key]
@@ -171,14 +167,14 @@ def task_features_require_result_persistence(task: "Task") -> bool:
     return False
 
 
-def _format_user_supplied_storage_key(key):
+def _format_user_supplied_storage_key(key: str) -> str:
     # Note here we are pinning to task runs since flow runs do not support storage keys
     # yet; we'll need to split logic in the future or have two separate functions
     runtime_vars = {key: getattr(prefect.runtime, key) for key in dir(prefect.runtime)}
     return key.format(**runtime_vars, parameters=prefect.runtime.task_run.parameters)
 
 
-class ResultFactory(pydantic.BaseModel):
+class ResultFactory(BaseModel):
     """
     A utility to generate `Result` types.
     """
@@ -186,7 +182,7 @@ class ResultFactory(pydantic.BaseModel):
     persist_result: bool
     cache_result_in_memory: bool
     serializer: Serializer
-    storage_block_id: Optional[uuid.UUID]
+    storage_block_id: Optional[uuid.UUID] = None
     storage_block: WritableFileSystem
     storage_key_fn: Callable[[], str]
 
@@ -430,7 +426,7 @@ class ResultFactory(pydantic.BaseModel):
             )
 
     @sync_compatible
-    async def create_result(self, obj: R) -> Union[R, "BaseResult[R]"]:
+    async def create_result(self, obj: R, key: str = None) -> Union[R, "BaseResult[R]"]:
         """
         Create a result type for the given object.
 
@@ -450,11 +446,20 @@ class ResultFactory(pydantic.BaseModel):
         if type(obj) in LITERAL_TYPES:
             return await LiteralResult.create(obj)
 
+        if key:
+
+            def key_fn():
+                return key
+
+            storage_key_fn = key_fn
+        else:
+            storage_key_fn = self.storage_key_fn
+
         return await PersistedResult.create(
             obj,
             storage_block=self.storage_block,
             storage_block_id=self.storage_block_id,
-            storage_key_fn=self.storage_key_fn,
+            storage_key_fn=storage_key_fn,
             serializer=self.serializer,
             cache_object=should_cache_object,
         )
@@ -475,17 +480,19 @@ class ResultFactory(pydantic.BaseModel):
         assert (
             self.storage_block_id is not None
         ), "Unexpected storage block ID. Was it persisted?"
-        blob = PersistedResultBlob.parse_raw(
+        blob = PersistedResultBlob.model_validate_json(
             await self.storage_block.read_path(f"parameters/{identifier}")
         )
         return self.serializer.loads(blob.data)
 
 
 @register_base_type
-class BaseResult(pydantic.BaseModel, abc.ABC, Generic[R]):
+class BaseResult(BaseModel, abc.ABC, Generic[R]):
+    model_config = ConfigDict(extra="forbid")
+
     type: str
-    artifact_type: Optional[str]
-    artifact_description: Optional[str]
+    artifact_type: Optional[str] = None
+    artifact_description: Optional[str] = None
 
     def __init__(self, **data: Any) -> None:
         type_string = get_dispatch_key(self) if type(self) != BaseResult else "__base__"
@@ -497,12 +504,12 @@ class BaseResult(pydantic.BaseModel, abc.ABC, Generic[R]):
             try:
                 subcls = lookup_type(cls, dispatch_key=kwargs["type"])
             except KeyError as exc:
-                raise pydantic.ValidationError(errors=[exc], model=cls)
+                raise ValidationError(errors=[exc], model=cls)
             return super().__new__(subcls)
         else:
             return super().__new__(cls)
 
-    _cache: Any = pydantic.PrivateAttr(NotSet)
+    _cache: Any = PrivateAttr(NotSet)
 
     def _cache_object(self, obj: Any) -> None:
         self._cache = obj
@@ -524,12 +531,10 @@ class BaseResult(pydantic.BaseModel, abc.ABC, Generic[R]):
     ) -> "BaseResult[R]":
         ...
 
-    class Config:
-        extra = "forbid"
-
     @classmethod
     def __dispatch_key__(cls, **kwargs):
-        return cls.__fields__.get("type").get_default()
+        default = cls.model_fields.get("type").get_default()
+        return cls.__name__ if isinstance(default, PydanticUndefinedType) else default
 
 
 class UnpersistedResult(BaseResult):
@@ -537,7 +542,7 @@ class UnpersistedResult(BaseResult):
     Result type for results that are not persisted outside of local memory.
     """
 
-    type = "unpersisted"
+    type: str = "unpersisted"
 
     @sync_compatible
     async def get(self) -> R:
@@ -572,8 +577,8 @@ class LiteralResult(BaseResult):
     They are not persisted to external result storage.
     """
 
-    type = "literal"
-    value: Any
+    type: str = "literal"
+    value: Any = None
 
     def has_cached_object(self) -> bool:
         # This result type always has the object cached in memory
@@ -609,13 +614,13 @@ class PersistedResult(BaseResult):
     content was written.
     """
 
-    type = "reference"
+    type: str = "reference"
 
     serializer_type: str
     storage_block_id: uuid.UUID
     storage_key: str
 
-    _should_cache_object: bool = pydantic.PrivateAttr(default=True)
+    _should_cache_object: bool = PrivateAttr(default=True)
 
     @sync_compatible
     @inject_client
@@ -643,7 +648,7 @@ class PersistedResult(BaseResult):
         block_document = await client.read_block_document(self.storage_block_id)
         storage_block: ReadableFileSystem = Block._from_block_document(block_document)
         content = await storage_block.read_path(self.storage_key)
-        blob = PersistedResultBlob.parse_raw(content)
+        blob = PersistedResultBlob.model_validate_json(content)
         return blob
 
     @staticmethod
@@ -686,7 +691,6 @@ class PersistedResult(BaseResult):
             raise TypeError(
                 f"Expected type 'str' for result storage key; got value {key!r}"
             )
-
         await storage_block.write_path(key, content=blob.to_bytes())
 
         description = f"Result of type `{type(obj).__name__}`"
@@ -716,7 +720,7 @@ class PersistedResult(BaseResult):
         return result
 
 
-class PersistedResultBlob(pydantic.BaseModel):
+class PersistedResultBlob(BaseModel):
     """
     The format of the content stored by a persisted result.
 
@@ -725,10 +729,10 @@ class PersistedResultBlob(pydantic.BaseModel):
 
     serializer: Serializer
     data: bytes
-    prefect_version: str = pydantic.Field(default=prefect.__version__)
+    prefect_version: str = Field(default=prefect.__version__)
 
     def to_bytes(self) -> bytes:
-        return self.json().encode()
+        return self.model_dump_json(serialize_as_any=True).encode()
 
 
 class UnknownResult(BaseResult):
@@ -742,7 +746,7 @@ class UnknownResult(BaseResult):
     completed task.
     """
 
-    type = "unknown"
+    type: str = "unknown"
     value: None
 
     def has_cached_object(self) -> bool:

prefect/runner/runner.py

@@ -71,19 +71,15 @@ from prefect.client.schemas.objects import (
     StateType,
 )
 from prefect.client.schemas.schedules import SCHEDULE_TYPES
-from prefect.deployments.deployments import load_flow_from_flow_run
 from prefect.deployments.runner import (
     EntrypointType,
     RunnerDeployment,
 )
 from prefect.deployments.schedules import FlexibleScheduleList
 from prefect.events import DeploymentTriggerTypes, TriggerTypes
-from prefect.exceptions import (
-    Abort,
-)
-from prefect.flows import Flow
+from prefect.exceptions import Abort, ObjectNotFound
+from prefect.flows import Flow, load_flow_from_flow_run
 from prefect.logging.loggers import PrefectLogAdapter, flow_run_logger, get_logger
-from prefect.runner.server import start_webserver
 from prefect.runner.storage import RunnerStorage
 from prefect.settings import (
     PREFECT_API_URL,
@@ -233,7 +229,7 @@ class Runner:
         description: Optional[str] = None,
         tags: Optional[List[str]] = None,
         version: Optional[str] = None,
-        enforce_parameter_schema: bool = False,
+        enforce_parameter_schema: bool = True,
         entrypoint_type: EntrypointType = EntrypointType.FILE_PATH,
     ) -> UUID:
         """
@@ -366,6 +362,8 @@ class Runner:
                 runner.start()
             ```
         """
+        from prefect.runner.server import start_webserver
+
         _register_signal(signal.SIGTERM, self.handle_sigterm)
 
         webserver = webserver if webserver is not None else self.webserver
@@ -1082,7 +1080,7 @@ class Runner:
         state_updates = state_updates or {}
         state_updates.setdefault("name", "Cancelled")
         state_updates.setdefault("type", StateType.CANCELLED)
-        state = flow_run.state.copy(update=state_updates)
+        state = flow_run.state.model_copy(update=state_updates)
 
         await self._client.set_flow_run_state(flow_run.id, state, force=True)
 
@@ -1130,12 +1128,18 @@ class Runner:
         Run the hooks for a flow.
         """
         if state.is_cancelling():
-            flow = await load_flow_from_flow_run(
-                flow_run, client=self._client, storage_base_path=str(self._tmp_dir)
-            )
-            hooks = flow.on_cancellation or []
+            try:
+                flow = await load_flow_from_flow_run(
+                    flow_run, storage_base_path=str(self._tmp_dir)
+                )
+                hooks = flow.on_cancellation_hooks or []
 
-            await _run_hooks(hooks, flow_run, flow, state)
+                await _run_hooks(hooks, flow_run, flow, state)
+            except ObjectNotFound:
+                run_logger = self._get_flow_run_logger(flow_run)
+                run_logger.warning(
+                    f"Runner cannot retrieve flow to execute cancellation hooks for flow run {flow_run.id!r}."
+                )
 
     async def _run_on_crashed_hooks(
         self,
@@ -1147,9 +1151,9 @@ class Runner:
         """
         if state.is_crashed():
             flow = await load_flow_from_flow_run(
-                flow_run, client=self._client, storage_base_path=str(self._tmp_dir)
+                flow_run, storage_base_path=str(self._tmp_dir)
             )
-            hooks = flow.on_crashed or []
+            hooks = flow.on_crashed_hooks or []
 
             await _run_hooks(hooks, flow_run, flow, state)
 

prefect/runner/server.py

@@ -3,11 +3,10 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
 
 import pendulum
 import uvicorn
-from prefect._vendor.fastapi import APIRouter, FastAPI, HTTPException, status
-from prefect._vendor.fastapi.responses import JSONResponse
+from fastapi import APIRouter, FastAPI, HTTPException, status
+from fastapi.responses import JSONResponse
 from typing_extensions import Literal
 
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
 from prefect._internal.schemas.validators import validate_values_conform_to_schema
 from prefect.client.orchestration import get_client
 from prefect.exceptions import MissingFlowError, ScriptError
@@ -30,10 +29,7 @@ if TYPE_CHECKING:
     from prefect.client.schemas.responses import DeploymentResponse
     from prefect.runner import Runner
 
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import BaseModel
-else:
-    from pydantic import BaseModel
+from pydantic import BaseModel
 
 logger = get_logger("webserver")
 
@@ -161,7 +157,7 @@ async def get_subflow_schemas(runner: "Runner") -> Dict[str, Dict]:
             script = deployment.entrypoint.split(":")[0]
             subflows = load_flows_from_script(script)
             for flow in subflows:
-                schemas[flow.name] = flow.parameters.dict()
+                schemas[flow.name] = flow.parameters.model_dump()
 
     return schemas
 
@@ -183,7 +179,7 @@ def _flow_schema_changed(flow: Flow, schemas: Dict[str, Dict]) -> bool:
     flow_name_with_dashes = flow.name.replace("_", "-")
 
     schema = schemas.get(flow.name, None) or schemas.get(flow_name_with_dashes, None)
-    if schema is not None and flow.parameters.dict() != schema:
+    if schema is not None and flow.parameters.model_dump() != schema:
         return True
     return False
 
@@ -236,7 +232,7 @@ def _build_generic_endpoint_for_flows(
 
         return JSONResponse(
             status_code=status.HTTP_201_CREATED,
-            content=flow_run.dict(json_compatible=True),
+            content=flow_run.model_dump(mode="json"),
         )
 
     return _create_flow_run_for_flow_from_fqn

prefect/runner/storage.py

@@ -8,20 +8,15 @@ from uuid import uuid4
 
 import fsspec
 from anyio import run_process
+from pydantic import SecretStr
 
 from prefect._internal.concurrency.api import create_call, from_async
-from prefect._internal.pydantic import HAS_PYDANTIC_V2
 from prefect.blocks.core import Block, BlockNotSavedError
 from prefect.blocks.system import Secret
 from prefect.filesystems import ReadableDeploymentStorage, WritableDeploymentStorage
 from prefect.logging.loggers import get_logger
 from prefect.utilities.collections import visit_collection
 
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import SecretStr
-else:
-    from pydantic import SecretStr
-
 
 @runtime_checkable
 class RunnerStorage(Protocol):
@@ -159,7 +154,7 @@ class GitRepository:
         url_components = urlparse(self._url)
 
         credentials = (
-            self._credentials.dict()
+            self._credentials.model_dump()
             if isinstance(self._credentials, Block)
             else deepcopy(self._credentials)
         )
@@ -392,7 +387,7 @@ class RemoteStorage:
                 if hasattr(obj, "value"):
                     return obj.value
                 else:
-                    return obj.dict()
+                    return obj.model_dump()
             return obj
 
         settings_with_block_values = visit_collection(

prefect/runner/submit.py

@@ -43,7 +43,7 @@ async def _submit_flow_to_runner(
     Returns:
         A `FlowRun` object representing the flow run that was submitted.
     """
-    from prefect.engine import (
+    from prefect.utilities.engine import (
         _dynamic_key_for_task_run,
         collect_task_run_inputs,
         resolve_inputs,
@@ -90,7 +90,7 @@ async def _submit_flow_to_runner(
         )
         response.raise_for_status()
 
-        return FlowRun.parse_obj(response.json())
+        return FlowRun.model_validate(response.json())
 
 
 @overload

prefect/runner/utils.py

@@ -1,8 +1,8 @@
 from copy import deepcopy
 from typing import Any, Dict
 
-from prefect._vendor.fastapi import FastAPI
-from prefect._vendor.fastapi.openapi.utils import get_openapi
+from fastapi import FastAPI
+from fastapi.openapi.utils import get_openapi
 
 from prefect import __version__ as PREFECT_VERSION
 

prefect/serializers.py

@@ -15,6 +15,14 @@ import abc
 import base64
 from typing import Any, Dict, Generic, Optional, Type, TypeVar
 
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    TypeAdapter,
+    ValidationError,
+    field_validator,
+)
 from typing_extensions import Literal, Self
 
 from prefect._internal.schemas.validators import (
@@ -23,32 +31,11 @@ from prefect._internal.schemas.validators import (
     validate_dump_kwargs,
     validate_load_kwargs,
     validate_picklelib,
-    validate_picklelib_version,
 )
-from prefect.pydantic import HAS_PYDANTIC_V2
 from prefect.utilities.dispatch import get_dispatch_key, lookup_type, register_base_type
 from prefect.utilities.importtools import from_qualified_name, to_qualified_name
 from prefect.utilities.pydantic import custom_pydantic_encoder
 
-if HAS_PYDANTIC_V2:
-    from pydantic.v1 import (
-        BaseModel,
-        Field,
-        ValidationError,
-        parse_obj_as,
-        root_validator,
-        validator,
-    )
-else:
-    from pydantic import (
-        BaseModel,
-        Field,
-        ValidationError,
-        parse_obj_as,
-        root_validator,
-        validator,
-    )
-
 D = TypeVar("D")
 
 
@@ -73,7 +60,9 @@ def prefect_json_object_decoder(result: dict):
     with `prefect_json_object_encoder`
     """
     if "__class__" in result:
-        return parse_obj_as(from_qualified_name(result["__class__"]), result["data"])
+        return TypeAdapter(from_qualified_name(result["__class__"])).validate_python(
+            result["data"]
+        )
     elif "__exc_type__" in result:
         return from_qualified_name(result["__exc_type__"])(result["message"])
     else:
@@ -112,12 +101,12 @@ class Serializer(BaseModel, Generic[D], abc.ABC):
     def loads(self, blob: bytes) -> D:
         """Decode the blob of bytes into an object."""
 
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
     @classmethod
-    def __dispatch_key__(cls):
-        return cls.__fields__.get("type").get_default()
+    def __dispatch_key__(cls) -> str:
+        type_str = cls.model_fields["type"].default
+        return type_str if isinstance(type_str, str) else None
 
 
 class PickleSerializer(Serializer):
@@ -133,15 +122,15 @@ class PickleSerializer(Serializer):
     type: Literal["pickle"] = "pickle"
 
     picklelib: str = "cloudpickle"
-    picklelib_version: str = None
+    picklelib_version: Optional[str] = None
 
-    @validator("picklelib")
+    @field_validator("picklelib")
     def check_picklelib(cls, value):
         return validate_picklelib(value)
 
-    @root_validator
-    def check_picklelib_version(cls, values):
-        return validate_picklelib_version(values)
+    # @model_validator(mode="before")
+    # def check_picklelib_version(cls, values):
+    #     return validate_picklelib_version(values)
 
     def dumps(self, obj: Any) -> bytes:
         pickler = from_qualified_name(self.picklelib)
@@ -185,11 +174,11 @@ class JSONSerializer(Serializer):
     dumps_kwargs: Dict[str, Any] = Field(default_factory=dict)
     loads_kwargs: Dict[str, Any] = Field(default_factory=dict)
 
-    @validator("dumps_kwargs")
+    @field_validator("dumps_kwargs")
     def dumps_kwargs_cannot_contain_default(cls, value):
         return validate_dump_kwargs(value)
 
-    @validator("loads_kwargs")
+    @field_validator("loads_kwargs")
     def loads_kwargs_cannot_contain_object_hook(cls, value):
         return validate_load_kwargs(value)
 
@@ -229,11 +218,11 @@ class CompressedSerializer(Serializer):
     serializer: Serializer
     compressionlib: str = "lzma"
 
-    @validator("serializer", pre=True)
+    @field_validator("serializer", mode="before")
     def validate_serializer(cls, value):
         return cast_type_names_to_serializers(value)
 
-    @validator("compressionlib")
+    @field_validator("compressionlib")
     def check_compressionlib(cls, value):
         return validate_compressionlib(value)
 

prefect/server/api/collections_data/views/aggregate-worker-metadata.json

@@ -1,14 +1,5 @@
 {
   "prefect": {
-    "prefect-agent": {
-      "type": "prefect-agent",
-      "documentation_url": "https://docs.prefect.io/latest/concepts/work-pools/#agent-overview",
-      "display_name": "Prefect Agent",
-      "logo_url": "https://cdn.sanity.io/images/3ugk85nk/production/c771bb53894c877e169c8db158c5598558b8f175-24x24.svg",
-      "install_command": "pip install prefect",
-      "default_base_job_configuration": {},
-      "description": "Execute flow runs on heterogeneous infrastructure using infrastructure blocks."
-    },
     "process": {
       "default_base_job_configuration": {
         "job_configuration": {
@@ -806,7 +797,7 @@
             "image": {
               "title": "Image",
               "description": "The image reference of a container image to use for created jobs. If not set, the latest Prefect image will be used.",
-              "example": "docker.io/prefecthq/prefect:2-latest",
+              "example": "docker.io/prefecthq/prefect:3-latest",
               "type": "string"
             },
             "registry_credentials": {
@@ -1040,7 +1031,7 @@
             "image": {
               "title": "Image Name",
               "description": "The image to use for a new Cloud Run Job. If not set, the latest Prefect image will be used. See https://cloud.google.com/run/docs/deploying#images.",
-              "example": "docker.io/prefecthq/prefect:2-latest",
+              "example": "docker.io/prefecthq/prefect:3-latest",
               "type": "string"
             },
             "cpu": {
@@ -1211,7 +1202,7 @@
             "image": {
               "title": "Image Name",
               "description": "The image to use for the Cloud Run job. If not provided the default Prefect image will be used.",
-              "default": "prefecthq/prefect:2-latest",
+              "default": "prefecthq/prefect:3-latest",
               "type": "string"
             },
             "args": {
@@ -1600,7 +1591,7 @@
             "image": {
               "title": "Image",
               "description": "The image reference of a container image to use for created jobs. If not set, the latest Prefect image will be used.",
-              "example": "docker.io/prefecthq/prefect:2-latest",
+              "example": "docker.io/prefecthq/prefect:3-latest",
               "type": "string"
             },
             "service_account_name": {
@@ -1688,4 +1679,4 @@
       "type": "kubernetes"
     }
   }
-}
+}