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

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()
                                       """
@@ -602,6 +650,7 @@ class TaskRunEngine(Generic[P, R]):
             key=self.compute_transaction_key(),
             store=ResultFactoryStore(result_factory=result_factory),
             overwrite=overwrite,
+            logger=self.logger,
         ) as txn:
             yield txn
 
@@ -611,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)
@@ -637,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:
@@ -645,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(

prefect/task_runners.py

@@ -1,6 +1,7 @@
 import abc
 import asyncio
 import sys
+import threading
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from contextvars import copy_context
@@ -41,8 +42,8 @@ if TYPE_CHECKING:
 
 P = ParamSpec("P")
 T = TypeVar("T")
-F = TypeVar("F", bound=PrefectFuture)
 R = TypeVar("R")
+F = TypeVar("F", bound=PrefectFuture, default=PrefectConcurrentFuture)
 
 
 class TaskRunner(abc.ABC, Generic[F]):
@@ -220,6 +221,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
         super().__init__()
         self._executor: Optional[ThreadPoolExecutor] = None
         self._max_workers = sys.maxsize if max_workers is None else max_workers
+        self._cancel_events: Dict[uuid.UUID, threading.Event] = {}
 
     def duplicate(self) -> "ThreadPoolTaskRunner":
         return type(self)(max_workers=self._max_workers)
@@ -270,6 +272,8 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
         from prefect.task_engine import run_task_async, run_task_sync
 
         task_run_id = uuid.uuid4()
+        cancel_event = threading.Event()
+        self._cancel_events[task_run_id] = cancel_event
         context = copy_context()
 
         flow_run_ctx = FlowRunContext.get()
@@ -280,31 +284,29 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
         else:
             self.logger.info(f"Submitting task {task.name} to thread pool executor...")
 
+        submit_kwargs = dict(
+            task=task,
+            task_run_id=task_run_id,
+            parameters=parameters,
+            wait_for=wait_for,
+            return_type="state",
+            dependencies=dependencies,
+            context=dict(cancel_event=cancel_event),
+        )
+
         if task.isasync:
             # TODO: Explore possibly using a long-lived thread with an event loop
             # for better performance
             future = self._executor.submit(
                 context.run,
                 asyncio.run,
-                run_task_async(
-                    task=task,
-                    task_run_id=task_run_id,
-                    parameters=parameters,
-                    wait_for=wait_for,
-                    return_type="state",
-                    dependencies=dependencies,
-                ),
+                run_task_async(**submit_kwargs),
             )
         else:
             future = self._executor.submit(
                 context.run,
                 run_task_sync,
-                task=task,
-                task_run_id=task_run_id,
-                parameters=parameters,
-                wait_for=wait_for,
-                return_type="state",
-                dependencies=dependencies,
+                **submit_kwargs,
             )
         prefect_future = PrefectConcurrentFuture(
             task_run_id=task_run_id, wrapped_future=future
@@ -337,14 +339,24 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture]):
     ):
         return super().map(task, parameters, wait_for)
 
+    def cancel_all(self):
+        for event in self._cancel_events.values():
+            event.set()
+            self.logger.debug("Set cancel event")
+
+        if self._executor is not None:
+            self._executor.shutdown(cancel_futures=True)
+            self._executor = None
+
     def __enter__(self):
         super().__enter__()
         self._executor = ThreadPoolExecutor(max_workers=self._max_workers)
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
+        self.cancel_all()
         if self._executor is not None:
-            self._executor.shutdown()
+            self._executor.shutdown(cancel_futures=True)
             self._executor = None
         super().__exit__(exc_type, exc_value, traceback)
 

prefect/task_worker.py

@@ -7,7 +7,7 @@ import sys
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import AsyncExitStack
 from contextvars import copy_context
-from typing import List, Optional
+from typing import Optional
 from uuid import UUID
 
 import anyio
@@ -20,6 +20,7 @@ from websockets.exceptions import InvalidStatusCode
 
 from prefect import Task
 from prefect._internal.concurrency.api import create_call, from_sync
+from prefect.cache_policies import DEFAULT, NONE
 from prefect.client.orchestration import get_client
 from prefect.client.schemas.objects import TaskRun
 from prefect.client.subscriptions import Subscription
@@ -32,9 +33,11 @@ from prefect.settings import (
 )
 from prefect.states import Pending
 from prefect.task_engine import run_task_async, run_task_sync
+from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import asyncnullcontext, sync_compatible
 from prefect.utilities.engine import emit_task_run_state_change_event, propose_state
 from prefect.utilities.processutils import _register_signal
+from prefect.utilities.urls import url_for
 
 logger = get_logger("task_worker")
 
@@ -76,7 +79,16 @@ class TaskWorker:
         *tasks: Task,
         limit: Optional[int] = 10,
     ):
-        self.tasks: List[Task] = list(tasks)
+        self.tasks = []
+        for t in tasks:
+            if isinstance(t, Task):
+                if t.cache_policy in [None, NONE, NotSet]:
+                    self.tasks.append(
+                        t.with_options(persist_result=True, cache_policy=DEFAULT)
+                    )
+                else:
+                    self.tasks.append(t.with_options(persist_result=True))
+
         self.task_keys = set(t.task_key for t in tasks if isinstance(t, Task))
 
         self._started_at: Optional[pendulum.DateTime] = None
@@ -277,10 +289,6 @@ class TaskWorker:
                     await self._client._client.delete(f"/task_runs/{task_run.id}")
                 return
 
-        logger.debug(
-            f"Submitting run {task_run.name!r} of task {task.name!r} to engine"
-        )
-
         try:
             new_state = Pending()
             new_state.state_details.deferred = True
@@ -315,6 +323,11 @@ class TaskWorker:
             validated_state=state,
         )
 
+        if task_run_url := url_for(task_run):
+            logger.info(
+                f"Submitting task run {task_run.name!r} to engine. View run in the UI at {task_run_url!r}"
+            )
+
         if task.isasync:
             await run_task_async(
                 task=task,

prefect/tasks.py

@@ -33,6 +33,9 @@ from uuid import UUID, uuid4
 
 from typing_extensions import Literal, ParamSpec
 
+from prefect._internal.compatibility.deprecated import (
+    deprecated_async_method,
+)
 from prefect.cache_policies import DEFAULT, NONE, CachePolicy
 from prefect.client.orchestration import get_client
 from prefect.client.schemas import TaskRun
@@ -47,7 +50,6 @@ from prefect.futures import PrefectDistributedFuture, PrefectFuture, PrefectFutu
 from prefect.logging.loggers import get_logger
 from prefect.results import ResultFactory, ResultSerializer, ResultStorage
 from prefect.settings import (
-    PREFECT_RESULTS_PERSIST_BY_DEFAULT,
     PREFECT_TASK_DEFAULT_RETRIES,
     PREFECT_TASK_DEFAULT_RETRY_DELAY_SECONDS,
 )
@@ -64,6 +66,7 @@ from prefect.utilities.callables import (
 )
 from prefect.utilities.hashing import hash_objects
 from prefect.utilities.importtools import to_qualified_name
+from prefect.utilities.urls import url_for
 
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
@@ -384,9 +387,20 @@ class Task(Generic[P, R]):
         self.cache_expiration = cache_expiration
         self.refresh_cache = refresh_cache
 
+        # result persistence settings
         if persist_result is None:
-            persist_result = PREFECT_RESULTS_PERSIST_BY_DEFAULT.value()
-        if not persist_result:
+            if any(
+                [
+                    cache_policy and cache_policy != NONE and cache_policy != NotSet,
+                    cache_key_fn is not None,
+                    result_storage_key is not None,
+                    result_storage is not None,
+                    result_serializer is not None,
+                ]
+            ):
+                persist_result = True
+
+        if persist_result is False:
             self.cache_policy = None if cache_policy is None else NONE
             if cache_policy and cache_policy is not NotSet and cache_policy != NONE:
                 logger.warning(
@@ -425,6 +439,14 @@ class Task(Generic[P, R]):
 
         self.retry_jitter_factor = retry_jitter_factor
         self.persist_result = persist_result
+
+        if result_storage and not isinstance(result_storage, str):
+            if getattr(result_storage, "_block_document_id", None) is None:
+                raise TypeError(
+                    "Result storage configuration must be persisted server-side."
+                    " Please call `.save()` on your block before passing it in."
+                )
+
         self.result_storage = result_storage
         self.result_serializer = result_serializer
         self.result_storage_key = result_storage_key
@@ -477,7 +499,7 @@ class Task(Generic[P, R]):
         cache_key_fn: Optional[
             Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
         ] = None,
-        task_run_name: Optional[Union[Callable[[], str], str]] = None,
+        task_run_name: Optional[Union[Callable[[], str], str, Type[NotSet]]] = NotSet,
         cache_expiration: Optional[datetime.timedelta] = None,
         retries: Union[int, Type[NotSet]] = NotSet,
         retry_delay_seconds: Union[
@@ -591,7 +613,9 @@ class Task(Generic[P, R]):
             else self.cache_policy,
             cache_key_fn=cache_key_fn or self.cache_key_fn,
             cache_expiration=cache_expiration or self.cache_expiration,
-            task_run_name=task_run_name,
+            task_run_name=task_run_name
+            if task_run_name is not NotSet
+            else self.task_run_name,
             retries=retries if retries is not NotSet else self.retries,
             retry_delay_seconds=(
                 retry_delay_seconds
@@ -869,6 +893,7 @@ class Task(Generic[P, R]):
     ) -> State[T]:
         ...
 
+    @deprecated_async_method
     def submit(
         self,
         *args: Any,
@@ -1033,6 +1058,7 @@ class Task(Generic[P, R]):
     ) -> PrefectFutureList[State[T]]:
         ...
 
+    @deprecated_async_method
     def map(
         self,
         *args: Any,
@@ -1275,14 +1301,20 @@ class Task(Generic[P, R]):
         # Convert the call args/kwargs to a parameter dict
         parameters = get_call_parameters(self.fn, args, kwargs)
 
-        task_run = run_coro_as_sync(
+        task_run: TaskRun = run_coro_as_sync(
             self.create_run(
                 parameters=parameters,
                 deferred=True,
                 wait_for=wait_for,
                 extra_task_inputs=dependencies,
             )
-        )
+        )  # type: ignore
+
+        if task_run_url := url_for(task_run):
+            logger.info(
+                f"Created task run {task_run.name!r}. View it in the UI at {task_run_url!r}"
+            )
+
         return PrefectDistributedFuture(task_run_id=task_run.id)
 
     def delay(self, *args: P.args, **kwargs: P.kwargs) -> PrefectDistributedFuture:

prefect/transactions.py

@@ -1,3 +1,4 @@
+import logging
 from contextlib import contextmanager
 from contextvars import ContextVar, Token
 from typing import (
@@ -7,21 +8,25 @@ from typing import (
     List,
     Optional,
     Type,
+    Union,
 )
 
 from pydantic import Field
 from typing_extensions import Self
 
 from prefect.context import ContextModel, FlowRunContext, TaskRunContext
+from prefect.exceptions import MissingContextError
+from prefect.logging.loggers import PrefectLogAdapter, get_logger, get_run_logger
 from prefect.records import RecordStore
 from prefect.records.result_store import ResultFactoryStore
 from prefect.results import (
     BaseResult,
     ResultFactory,
-    get_or_create_default_result_storage,
+    get_default_result_storage,
 )
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import AutoEnum
+from prefect.utilities.engine import _get_hook_name
 
 
 class IsolationLevel(AutoEnum):
@@ -58,6 +63,7 @@ class Transaction(ContextModel):
         default_factory=list
     )
     overwrite: bool = False
+    logger: Union[logging.Logger, logging.LoggerAdapter, None] = None
     _staged_value: Any = None
     __var__: ContextVar = ContextVar("transaction")
 
@@ -174,10 +180,13 @@ class Transaction(ContextModel):
             return False
 
         try:
+            hook_name = None
+
             for child in self.children:
                 child.commit()
 
             for hook in self.on_commit_hooks:
+                hook_name = _get_hook_name(hook)
                 hook(self)
 
             if self.store and self.key:
@@ -185,6 +194,19 @@ class Transaction(ContextModel):
             self.state = TransactionState.COMMITTED
             return True
         except Exception:
+            if self.logger:
+                if hook_name:
+                    msg = (
+                        f"An error was encountered while running commit hook {hook_name!r}",
+                    )
+                else:
+                    msg = (
+                        f"An error was encountered while committing transaction {self.key!r}",
+                    )
+                self.logger.exception(
+                    msg,
+                    exc_info=True,
+                )
             self.rollback()
             return False
 
@@ -212,6 +234,7 @@ class Transaction(ContextModel):
 
         try:
             for hook in reversed(self.on_rollback_hooks):
+                hook_name = _get_hook_name(hook)
                 hook(self)
 
             self.state = TransactionState.ROLLED_BACK
@@ -221,6 +244,11 @@ class Transaction(ContextModel):
 
             return True
         except Exception:
+            if self.logger:
+                self.logger.exception(
+                    f"An error was encountered while running rollback hook {hook_name!r}",
+                    exc_info=True,
+                )
             return False
 
     @classmethod
@@ -238,6 +266,7 @@ def transaction(
     store: Optional[RecordStore] = None,
     commit_mode: Optional[CommitMode] = None,
     overwrite: bool = False,
+    logger: Optional[PrefectLogAdapter] = None,
 ) -> Generator[Transaction, None, None]:
     """
     A context manager for opening and managing a transaction.
@@ -268,7 +297,7 @@ def transaction(
                 }
             )
         else:
-            default_storage = get_or_create_default_result_storage(_sync=True)
+            default_storage = get_default_result_storage(_sync=True)
             if existing_factory:
                 new_factory = existing_factory.model_copy(
                     update={
@@ -288,7 +317,16 @@ def transaction(
             result_factory=new_factory,
         )
 
+    try:
+        logger = logger or get_run_logger()
+    except MissingContextError:
+        logger = get_logger("transactions")
+
     with Transaction(
-        key=key, store=store, commit_mode=commit_mode, overwrite=overwrite
+        key=key,
+        store=store,
+        commit_mode=commit_mode,
+        overwrite=overwrite,
+        logger=logger,
     ) as txn:
         yield txn