prefect-client 3.0.0rc9__py3-none-any.whl → 3.0.0rc10__py3-none-any.whl

prefect/results.py

@@ -28,7 +28,6 @@ from prefect.client.utilities import inject_client
 from prefect.exceptions import MissingResult, ObjectAlreadyExists
 from prefect.filesystems import (
     LocalFileSystem,
-    ReadableFileSystem,
     WritableFileSystem,
 )
 from prefect.logging import get_logger
@@ -111,22 +110,32 @@ async def _get_or_create_default_storage(block_document_slug: str) -> ResultStor
 
 
 @sync_compatible
-async def get_or_create_default_result_storage() -> ResultStorage:
+async def get_default_result_storage() -> ResultStorage:
     """
     Generate a default file system for result storage.
     """
-    return await _get_or_create_default_storage(
-        PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value()
-    )
+    default_block = PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value()
+
+    if default_block is not None:
+        return await Block.load(default_block)
+
+    # otherwise, use the local file system
+    basepath = PREFECT_LOCAL_STORAGE_PATH.value()
+    return LocalFileSystem(basepath=basepath)
 
 
 async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
     """
     Generate a default file system for background task parameter/result storage.
     """
-    return await _get_or_create_default_storage(
-        PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK.value()
-    )
+    default_block = PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK.value()
+
+    if default_block is not None:
+        return await Block.load(default_block)
+
+    # otherwise, use the local file system
+    basepath = PREFECT_LOCAL_STORAGE_PATH.value()
+    return LocalFileSystem(basepath=basepath)
 
 
 def get_default_result_serializer() -> ResultSerializer:
@@ -177,9 +186,7 @@ class ResultFactory(BaseModel):
                 kwargs.pop(key)
 
         # Apply defaults
-        kwargs.setdefault(
-            "result_storage", await get_or_create_default_result_storage()
-        )
+        kwargs.setdefault("result_storage", await get_default_result_storage())
         kwargs.setdefault("result_serializer", get_default_result_serializer())
         kwargs.setdefault("persist_result", get_default_persist_setting())
         kwargs.setdefault("cache_result_in_memory", True)
@@ -230,9 +237,7 @@ class ResultFactory(BaseModel):
         """
         Create a new result factory for a task.
         """
-        return await cls._from_task(
-            task, get_or_create_default_result_storage, client=client
-        )
+        return await cls._from_task(task, get_default_result_storage, client=client)
 
     @classmethod
     @inject_client
@@ -268,7 +273,14 @@ class ResultFactory(BaseModel):
             if ctx and ctx.result_factory
             else get_default_result_serializer()
         )
-        persist_result = task.persist_result
+        if task.persist_result is None:
+            persist_result = (
+                ctx.result_factory.persist_result
+                if ctx and ctx.result_factory
+                else get_default_persist_setting()
+            )
+        else:
+            persist_result = task.persist_result
 
         cache_result_in_memory = task.cache_result_in_memory
 
@@ -330,16 +342,7 @@ class ResultFactory(BaseModel):
                 # Avoid saving the block if it already has an identifier assigned
                 storage_block_id = storage_block._block_document_id
             else:
-                if persist_result:
-                    # TODO: Overwrite is true to avoid issues where the save collides with
-                    # a previously saved document with a matching hash
-                    storage_block_id = await storage_block._save(
-                        is_anonymous=True, overwrite=True, client=client
-                    )
-                else:
-                    # a None-type UUID on unpersisted storage should not matter
-                    # since the ID is generated on the server
-                    storage_block_id = None
+                storage_block_id = None
         elif isinstance(result_storage, str):
             storage_block = await Block.load(result_storage, client=client)
             storage_block_id = storage_block._block_document_id
@@ -412,9 +415,6 @@ class ResultFactory(BaseModel):
 
     @sync_compatible
     async def store_parameters(self, identifier: UUID, parameters: Dict[str, Any]):
-        assert (
-            self.storage_block_id is not None
-        ), "Unexpected storage block ID. Was it persisted?"
         data = self.serializer.dumps(parameters)
         blob = PersistedResultBlob(serializer=self.serializer, data=data)
         await self.storage_block.write_path(
@@ -423,9 +423,6 @@ class ResultFactory(BaseModel):
 
     @sync_compatible
     async def read_parameters(self, identifier: UUID) -> Dict[str, Any]:
-        assert (
-            self.storage_block_id is not None
-        ), "Unexpected storage block ID. Was it persisted?"
         blob = PersistedResultBlob.model_validate_json(
             await self.storage_block.read_path(f"parameters/{identifier}")
         )
@@ -435,10 +432,7 @@ class ResultFactory(BaseModel):
 @register_base_type
 class BaseResult(BaseModel, abc.ABC, Generic[R]):
     model_config = ConfigDict(extra="forbid")
-
     type: 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__"
@@ -504,11 +498,7 @@ class UnpersistedResult(BaseResult):
         obj: R,
         cache_object: bool = True,
     ) -> "UnpersistedResult[R]":
-        description = f"Unpersisted result of type `{type(obj).__name__}`"
-        result = cls(
-            artifact_type="result",
-            artifact_description=description,
-        )
+        result = cls()
         # Only store the object in local memory, it will not be sent to the API
         if cache_object:
             result._cache_object(obj)
@@ -528,8 +518,8 @@ class PersistedResult(BaseResult):
     type: str = "reference"
 
     serializer_type: str
-    storage_block_id: uuid.UUID
     storage_key: str
+    storage_block_id: Optional[uuid.UUID] = None
     expiration: Optional[DateTime] = None
 
     _should_cache_object: bool = PrivateAttr(default=True)
@@ -547,6 +537,17 @@ class PersistedResult(BaseResult):
         self._storage_block = storage_block
         self._serializer = serializer
 
+    @inject_client
+    async def _get_storage_block(self, client: "PrefectClient") -> WritableFileSystem:
+        if self._storage_block is not None:
+            return self._storage_block
+        elif self.storage_block_id is not None:
+            block_document = await client.read_block_document(self.storage_block_id)
+            self._storage_block = Block._from_block_document(block_document)
+        else:
+            self._storage_block = await get_default_result_storage()
+        return self._storage_block
+
     @sync_compatible
     @inject_client
     async def get(self, client: "PrefectClient") -> R:
@@ -567,12 +568,8 @@ class PersistedResult(BaseResult):
 
     @inject_client
     async def _read_blob(self, client: "PrefectClient") -> "PersistedResultBlob":
-        assert (
-            self.storage_block_id is not None
-        ), "Unexpected storage block ID. Was it persisted?"
-        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)
+        block = await self._get_storage_block(client=client)
+        content = await block.read_path(self.storage_key)
         blob = PersistedResultBlob.model_validate_json(content)
         return blob
 
@@ -607,10 +604,7 @@ class PersistedResult(BaseResult):
         obj = obj if obj is not NotSet else self._cache
 
         # next, the storage block
-        storage_block = self._storage_block
-        if storage_block is None:
-            block_document = await client.read_block_document(self.storage_block_id)
-            storage_block = Block._from_block_document(block_document)
+        storage_block = await self._get_storage_block(client=client)
 
         # finally, the serializer
         serializer = self._serializer
@@ -673,9 +667,9 @@ class PersistedResult(BaseResult):
         cls: "Type[PersistedResult]",
         obj: R,
         storage_block: WritableFileSystem,
-        storage_block_id: uuid.UUID,
         storage_key_fn: Callable[[], str],
         serializer: Serializer,
+        storage_block_id: Optional[uuid.UUID] = None,
         cache_object: bool = True,
         expiration: Optional[DateTime] = None,
         defer_persistence: bool = False,
@@ -686,31 +680,21 @@ class PersistedResult(BaseResult):
         The object will be serialized and written to the storage block under a unique
         key. It will then be cached on the returned result.
         """
-        assert (
-            storage_block_id is not None
-        ), "Unexpected storage block ID. Was it saved?"
-
         key = storage_key_fn()
         if not isinstance(key, str):
             raise TypeError(
                 f"Expected type 'str' for result storage key; got value {key!r}"
             )
-        description = f"Result of type `{type(obj).__name__}`"
         uri = cls._infer_path(storage_block, key)
-        if uri:
-            if isinstance(storage_block, LocalFileSystem):
-                description += f" persisted to: `{uri}`"
-            else:
-                description += f" persisted to [{uri}]({uri})."
-        else:
-            description += f" persisted with storage block `{storage_block_id}`."
+
+        # in this case we store an absolute path
+        if storage_block_id is None and uri is not None:
+            key = str(uri)
 
         result = cls(
             serializer_type=serializer.type,
             storage_block_id=storage_block_id,
             storage_key=key,
-            artifact_type="result",
-            artifact_description=description,
             expiration=expiration,
         )
 
@@ -787,5 +771,4 @@ class UnknownResult(BaseResult):
                 "Only None is supported."
             )
 
-        description = "Unknown result persisted to Prefect."
-        return cls(value=obj, artifact_type="result", artifact_description=description)
+        return cls(value=obj)

prefect/serializers.py

@@ -13,7 +13,7 @@ bytes to an object respectively.
 
 import abc
 import base64
-from typing import Any, Dict, Generic, Optional, Type, TypeVar
+from typing import Any, Dict, Generic, Optional, Type
 
 from pydantic import (
     BaseModel,
@@ -23,7 +23,7 @@ from pydantic import (
     ValidationError,
     field_validator,
 )
-from typing_extensions import Literal, Self
+from typing_extensions import Literal, Self, TypeVar
 
 from prefect._internal.schemas.validators import (
     cast_type_names_to_serializers,
@@ -36,7 +36,7 @@ from prefect.utilities.dispatch import get_dispatch_key, lookup_type, register_b
 from prefect.utilities.importtools import from_qualified_name, to_qualified_name
 from prefect.utilities.pydantic import custom_pydantic_encoder
 
-D = TypeVar("D")
+D = TypeVar("D", default=Any)
 
 
 def prefect_json_object_encoder(obj: Any) -> Any:

prefect/settings.py

@@ -42,7 +42,7 @@ dependent on the value of other settings or perform other dynamic effects.
 
 import logging
 import os
-import socket
+import re
 import string
 import warnings
 from contextlib import contextmanager
@@ -85,7 +85,6 @@ from prefect._internal.schemas.validators import validate_settings
 from prefect.exceptions import MissingProfileError
 from prefect.utilities.names import OBFUSCATED_PREFIX, obfuscate
 from prefect.utilities.pydantic import add_cloudpickle_reduction
-from prefect.utilities.slugify import slugify
 
 T = TypeVar("T")
 
@@ -404,18 +403,6 @@ def warn_on_misconfigured_api_url(values):
     return values
 
 
-def default_result_storage_block_name(
-    settings: Optional["Settings"] = None, value: Optional[str] = None
-):
-    """
-    `value_callback` for `PREFECT_DEFAULT_RESULT_STORAGE_BLOCK` that sets the default
-    value to the hostname of the machine.
-    """
-    if value is None:
-        return f"local-file-system/{slugify(socket.gethostname())}-storage"
-    return value
-
-
 def default_database_connection_url(settings, value):
     templater = template_with_settings(PREFECT_HOME, PREFECT_API_DATABASE_PASSWORD)
 
@@ -474,10 +461,8 @@ def default_cloud_ui_url(settings, value):
     # Otherwise, infer a value from the API URL
     ui_url = api_url = PREFECT_CLOUD_API_URL.value_from(settings)
 
-    if api_url.startswith("https://api.prefect.cloud"):
-        ui_url = ui_url.replace(
-            "https://api.prefect.cloud", "https://app.prefect.cloud", 1
-        )
+    if re.match(r"^https://api[\.\w]*.prefect.[^\.]+/", api_url):
+        ui_url = ui_url.replace("https://api", "https://app", 1)
 
     if ui_url.endswith("/api"):
         ui_url = ui_url[:-4]
@@ -1323,15 +1308,6 @@ PREFECT_API_MAX_FLOW_RUN_GRAPH_ARTIFACTS = Setting(int, default=10000)
 The maximum number of artifacts to show on a flow run graph on the v2 API
 """
 
-PREFECT_EXPERIMENTAL_ENABLE_WORKERS = Setting(bool, default=True)
-"""
-Whether or not to enable experimental Prefect workers.
-"""
-
-PREFECT_EXPERIMENTAL_WARN_WORKERS = Setting(bool, default=False)
-"""
-Whether or not to warn when experimental Prefect workers are used.
-"""
 
 PREFECT_EXPERIMENTAL_ENABLE_ENHANCED_CANCELLATION = Setting(bool, default=True)
 """
@@ -1423,10 +1399,7 @@ PREFECT_API_SERVICES_TASK_SCHEDULING_ENABLED = Setting(bool, default=True)
 Whether or not to start the task scheduling service in the server application.
 """
 
-PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK = Setting(
-    str,
-    default="local-file-system/prefect-task-scheduling",
-)
+PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK = Setting(Optional[str], default=None)
 """The `block-type/block-document` slug of a block to use as the default storage
 for autonomous tasks."""
 
@@ -1479,7 +1452,8 @@ PREFECT_EXPERIMENTAL_ENABLE_SCHEDULE_CONCURRENCY = Setting(bool, default=False)
 # Defaults -----------------------------------------------------------------------------
 
 PREFECT_DEFAULT_RESULT_STORAGE_BLOCK = Setting(
-    Optional[str], default=None, value_callback=default_result_storage_block_name
+    Optional[str],
+    default=None,
 )
 """The `block-type/block-document` slug of a block to use as the default result storage."""
 

prefect/task_engine.py

@@ -1,6 +1,8 @@
 import inspect
 import logging
+import threading
 import time
+from asyncio import CancelledError
 from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
 from textwrap import dedent
@@ -17,6 +19,7 @@ from typing import (
     Optional,
     Sequence,
     Set,
+    Type,
     TypeVar,
     Union,
 )
@@ -36,17 +39,18 @@ from prefect.context import (
     TaskRunContext,
     hydrated_context,
 )
-from prefect.events.schemas.events import Event
+from prefect.events.schemas.events import Event as PrefectEvent
 from prefect.exceptions import (
     Abort,
     Pause,
     PrefectException,
+    TerminationSignal,
     UpstreamTaskError,
 )
 from prefect.futures import PrefectFuture
 from prefect.logging.loggers import get_logger, patch_print, task_run_logger
 from prefect.records.result_store import ResultFactoryStore
-from prefect.results import ResultFactory, _format_user_supplied_storage_key
+from prefect.results import BaseResult, ResultFactory, _format_user_supplied_storage_key
 from prefect.settings import (
     PREFECT_DEBUG_MODE,
     PREFECT_TASKS_REFRESH_CACHE,
@@ -63,6 +67,7 @@ from prefect.states import (
     return_value_to_state,
 )
 from prefect.transactions import Transaction, transaction
+from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import call_with_parameters, parameters_to_args_kwargs
 from prefect.utilities.collections import visit_collection
@@ -80,6 +85,10 @@ P = ParamSpec("P")
 R = TypeVar("R")
 
 
+class TaskRunTimeoutError(TimeoutError):
+    """Raised when a task run exceeds its timeout."""
+
+
 @dataclass
 class TaskRunEngine(Generic[P, R]):
     task: Union[Task[P, R], Task[P, Coroutine[Any, Any, R]]]
@@ -89,11 +98,15 @@ class TaskRunEngine(Generic[P, R]):
     retries: int = 0
     wait_for: Optional[Iterable[PrefectFuture]] = None
     context: Optional[Dict[str, Any]] = None
+    # holds the return value from the user code
+    _return_value: Union[R, Type[NotSet]] = NotSet
+    # holds the exception raised by the user code, if any
+    _raised: Union[Exception, Type[NotSet]] = NotSet
     _initial_run_context: Optional[TaskRunContext] = None
     _is_started: bool = False
     _client: Optional[SyncPrefectClient] = None
     _task_name_set: bool = False
-    _last_event: Optional[Event] = None
+    _last_event: Optional[PrefectEvent] = None
 
     def __post_init__(self):
         if self.parameters is None:
@@ -136,7 +149,16 @@ class TaskRunEngine(Generic[P, R]):
             )
             return False
 
-    def call_hooks(self, state: State = None) -> Iterable[Callable]:
+    def is_cancelled(self) -> bool:
+        if (
+            self.context
+            and "cancel_event" in self.context
+            and isinstance(self.context["cancel_event"], threading.Event)
+        ):
+            return self.context["cancel_event"].is_set()
+        return False
+
+    def call_hooks(self, state: Optional[State] = None):
         if state is None:
             state = self.state
         task = self.task
@@ -171,7 +193,7 @@ class TaskRunEngine(Generic[P, R]):
             else:
                 self.logger.info(f"Hook {hook_name!r} finished running successfully")
 
-    def compute_transaction_key(self) -> str:
+    def compute_transaction_key(self) -> Optional[str]:
         key = None
         if self.task.cache_policy:
             flow_run_context = FlowRunContext.get()
@@ -304,12 +326,24 @@ class TaskRunEngine(Generic[P, R]):
         return new_state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
-        _result = self.state.result(raise_on_failure=raise_on_failure, fetch=True)
-        # state.result is a `sync_compatible` function that may or may not return an awaitable
-        # depending on whether the parent frame is sync or not
-        if inspect.isawaitable(_result):
-            _result = run_coro_as_sync(_result)
-        return _result
+        if self._return_value is not NotSet:
+            # if the return value is a BaseResult, we need to fetch it
+            if isinstance(self._return_value, BaseResult):
+                _result = self._return_value.get()
+                if inspect.isawaitable(_result):
+                    _result = run_coro_as_sync(_result)
+                return _result
+
+            # otherwise, return the value as is
+            return self._return_value
+
+        if self._raised is not NotSet:
+            # if the task raised an exception, raise it
+            if raise_on_failure:
+                raise self._raised
+
+            # otherwise, return the exception
+            return self._raised
 
     def handle_success(self, result: R, transaction: Transaction) -> R:
         result_factory = getattr(TaskRunContext.get(), "result_factory", None)
@@ -339,6 +373,7 @@ class TaskRunEngine(Generic[P, R]):
         if transaction.is_committed():
             terminal_state.name = "Cached"
         self.set_state(terminal_state)
+        self._return_value = result
         return result
 
     def handle_retry(self, exc: Exception) -> bool:
@@ -365,9 +400,11 @@ class TaskRunEngine(Generic[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                f"Task run failed with exception {exc!r} - "
-                f"Retry {self.retries + 1}/{self.task.retries} will start "
-                f"{str(delay) + ' second(s) from now' if delay else 'immediately'}"
+                "Task run failed with exception: %r - " "Retry %s/%s will start %s",
+                exc,
+                self.retries + 1,
+                self.task.retries,
+                str(delay) + " second(s) from now" if delay else "immediately",
             )
 
             self.set_state(new_state, force=True)
@@ -375,7 +412,9 @@ class TaskRunEngine(Generic[P, R]):
             return True
         elif self.retries >= self.task.retries:
             self.logger.error(
-                f"Task run failed with exception {exc!r} - Retries are exhausted"
+                "Task run failed with exception: %r - Retries are exhausted",
+                exc,
+                exc_info=True,
             )
             return False
 
@@ -394,12 +433,14 @@ class TaskRunEngine(Generic[P, R]):
                 )
             )
             self.set_state(state)
+            self._raised = exc
 
     def handle_timeout(self, exc: TimeoutError) -> None:
         if not self.handle_retry(exc):
-            message = (
-                f"Task run exceeded timeout of {self.task.timeout_seconds} seconds"
-            )
+            if isinstance(exc, TaskRunTimeoutError):
+                message = f"Task run exceeded timeout of {self.task.timeout_seconds} second(s)"
+            else:
+                message = f"Task run failed due to timeout: {exc!r}"
             self.logger.error(message)
             state = Failed(
                 data=exc,
@@ -407,12 +448,14 @@ class TaskRunEngine(Generic[P, R]):
                 name="TimedOut",
             )
             self.set_state(state)
+            self._raised = exc
 
     def handle_crash(self, exc: BaseException) -> None:
         state = run_coro_as_sync(exception_to_crashed_state(exc))
         self.logger.error(f"Crash detected! {state.message}")
         self.logger.debug("Crash details:", exc_info=exc)
         self.set_state(state, force=True)
+        self._raised = exc
 
     @contextmanager
     def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
@@ -498,6 +541,11 @@ class TaskRunEngine(Generic[P, R]):
                         )
                         yield self
 
+                except TerminationSignal as exc:
+                    # TerminationSignals are caught and handled as crashes
+                    self.handle_crash(exc)
+                    raise exc
+
                 except Exception:
                     # regular exceptions are caught and re-raised to the user
                     raise
@@ -539,8 +587,8 @@ class TaskRunEngine(Generic[P, R]):
 
                             @flow
                             def example_flow():
-                                say_hello.submit(name="Marvin)
-                                say_hello.wait()
+                                future = say_hello.submit(name="Marvin)
+                                future.wait()
 
                             example_flow()
                                       """
@@ -612,10 +660,16 @@ class TaskRunEngine(Generic[P, R]):
         # reenter the run context to ensure it is up to date for every run
         with self.setup_run_context():
             try:
-                with timeout_context(seconds=self.task.timeout_seconds):
+                with timeout_context(
+                    seconds=self.task.timeout_seconds,
+                    timeout_exc_type=TaskRunTimeoutError,
+                ):
                     self.logger.debug(
                         f"Executing task {self.task.name!r} for task run {self.task_run.name!r}..."
                     )
+                    if self.is_cancelled():
+                        raise CancelledError("Task run cancelled by the task runner")
+
                     yield self
             except TimeoutError as exc:
                 self.handle_timeout(exc)
@@ -638,6 +692,7 @@ class TaskRunEngine(Generic[P, R]):
                 else:
                     result = await call_with_parameters(self.task.fn, parameters)
                 self.handle_success(result, transaction=transaction)
+                return result
 
             return _call_task_fn()
         else:
@@ -646,6 +701,7 @@ class TaskRunEngine(Generic[P, R]):
             else:
                 result = call_with_parameters(self.task.fn, parameters)
             self.handle_success(result, transaction=transaction)
+            return result
 
 
 def run_task_sync(