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

prefect/task_engine.py

@@ -5,6 +5,7 @@ import time
 from asyncio import CancelledError
 from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
+from functools import wraps
 from textwrap import dedent
 from typing import (
     Any,
@@ -53,6 +54,7 @@ from prefect.records.result_store import ResultFactoryStore
 from prefect.results import BaseResult, ResultFactory, _format_user_supplied_storage_key
 from prefect.settings import (
     PREFECT_DEBUG_MODE,
+    PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION,
     PREFECT_TASKS_REFRESH_CACHE,
 )
 from prefect.states import (
@@ -124,8 +126,7 @@ class TaskRunEngine(Generic[P, R]):
             raise ValueError("Task run is not set")
         return self.task_run.state
 
-    @property
-    def can_retry(self) -> bool:
+    def can_retry(self, exc: Exception) -> bool:
         retry_condition: Optional[
             Callable[[Task[P, Coroutine[Any, Any, R]], TaskRun, State], bool]
         ] = self.task.retry_condition_fn
@@ -136,9 +137,19 @@ class TaskRunEngine(Generic[P, R]):
                 f"Running `retry_condition_fn` check {retry_condition!r} for task"
                 f" {self.task.name!r}"
             )
-            return not retry_condition or retry_condition(
-                self.task, self.task_run, self.state
+            state = Failed(
+                data=exc,
+                message=f"Task run encountered unexpected exception: {repr(exc)}",
             )
+            if inspect.iscoroutinefunction(retry_condition):
+                should_retry = run_coro_as_sync(
+                    retry_condition(self.task, self.task_run, state)
+                )
+            elif inspect.isfunction(retry_condition):
+                should_retry = retry_condition(self.task, self.task_run, state)
+            else:
+                should_retry = not retry_condition
+            return should_retry
         except Exception:
             self.logger.error(
                 (
@@ -269,6 +280,17 @@ class TaskRunEngine(Generic[P, R]):
             return
 
         new_state = Running()
+
+        if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            self.task_run.start_time = new_state.timestamp
+            self.task_run.run_count += 1
+
+            flow_run_context = FlowRunContext.get()
+            if flow_run_context:
+                # Carry forward any task run information from the flow run
+                flow_run = flow_run_context.flow_run
+                self.task_run.flow_run_run_count = flow_run.run_count
+
         state = self.set_state(new_state)
 
         # TODO: this is temporary until the API stops rejecting state transitions
@@ -298,24 +320,37 @@ class TaskRunEngine(Generic[P, R]):
         last_state = self.state
         if not self.task_run:
             raise ValueError("Task run is not set")
-        try:
-            new_state = propose_state_sync(
-                self.client, state, task_run_id=self.task_run.id, force=force
-            )
-        except Pause as exc:
-            # We shouldn't get a pause signal without a state, but if this happens,
-            # just use a Paused state to assume an in-process pause.
-            new_state = exc.state if exc.state else Paused()
-            if new_state.state_details.pause_reschedule:
-                # If we're being asked to pause and reschedule, we should exit the
-                # task and expect to be resumed later.
-                raise
 
-        # currently this is a hack to keep a reference to the state object
-        # that has an in-memory result attached to it; using the API state
+        if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            self.task_run.state = new_state = state
+
+            # Ensure that the state_details are populated with the current run IDs
+            new_state.state_details.task_run_id = self.task_run.id
+            new_state.state_details.flow_run_id = self.task_run.flow_run_id
+
+            # Predictively update the de-normalized task_run.state_* attributes
+            self.task_run.state_id = new_state.id
+            self.task_run.state_type = new_state.type
+            self.task_run.state_name = new_state.name
+        else:
+            try:
+                new_state = propose_state_sync(
+                    self.client, state, task_run_id=self.task_run.id, force=force
+                )
+            except Pause as exc:
+                # We shouldn't get a pause signal without a state, but if this happens,
+                # just use a Paused state to assume an in-process pause.
+                new_state = exc.state if exc.state else Paused()
+                if new_state.state_details.pause_reschedule:
+                    # If we're being asked to pause and reschedule, we should exit the
+                    # task and expect to be resumed later.
+                    raise
+
+            # currently this is a hack to keep a reference to the state object
+            # that has an in-memory result attached to it; using the API state
+            # could result in losing that reference
+            self.task_run.state = new_state
 
-        # could result in losing that reference
-        self.task_run.state = new_state
         # emit a state change event
         self._last_event = emit_task_run_state_change_event(
             task_run=self.task_run,
@@ -323,6 +358,7 @@ class TaskRunEngine(Generic[P, R]):
             validated_state=self.task_run.state,
             follows=self._last_event,
         )
+
         return new_state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
@@ -367,11 +403,19 @@ class TaskRunEngine(Generic[P, R]):
         )
         transaction.stage(
             terminal_state.data,
-            on_rollback_hooks=self.task.on_rollback_hooks,
-            on_commit_hooks=self.task.on_commit_hooks,
+            on_rollback_hooks=[
+                _with_transaction_hook_logging(hook, "rollback", self.logger)
+                for hook in self.task.on_rollback_hooks
+            ],
+            on_commit_hooks=[
+                _with_transaction_hook_logging(hook, "commit", self.logger)
+                for hook in self.task.on_commit_hooks
+            ],
         )
         if transaction.is_committed():
             terminal_state.name = "Cached"
+
+        self.record_terminal_state_timing(terminal_state)
         self.set_state(terminal_state)
         self._return_value = result
         return result
@@ -383,7 +427,7 @@ class TaskRunEngine(Generic[P, R]):
           - If the task has a retry delay, place in AwaitingRetry state with a delayed scheduled time.
         - If the task has no retries left, or the retry condition is not met, return False.
         """
-        if self.retries < self.task.retries and self.can_retry:
+        if self.retries < self.task.retries and self.can_retry(exc):
             if self.task.retry_delay_seconds:
                 delay = (
                     self.task.retry_delay_seconds[
@@ -398,6 +442,8 @@ class TaskRunEngine(Generic[P, R]):
             else:
                 delay = None
                 new_state = Retrying()
+                if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+                    self.task_run.run_count += 1
 
             self.logger.info(
                 "Task run failed with exception: %r - " "Retry %s/%s will start %s",
@@ -432,6 +478,7 @@ class TaskRunEngine(Generic[P, R]):
                     result_factory=getattr(context, "result_factory", None),
                 )
             )
+            self.record_terminal_state_timing(state)
             self.set_state(state)
             self._raised = exc
 
@@ -454,9 +501,20 @@ class TaskRunEngine(Generic[P, R]):
         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.record_terminal_state_timing(state)
         self.set_state(state, force=True)
         self._raised = exc
 
+    def record_terminal_state_timing(self, state: State) -> None:
+        if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            if self.task_run.start_time and not self.task_run.end_time:
+                self.task_run.end_time = state.timestamp
+
+                if self.task_run.state.is_running():
+                    self.task_run.total_run_time += (
+                        state.timestamp - self.task_run.state.timestamp
+                    )
+
     @contextmanager
     def setup_run_context(self, client: Optional[SyncPrefectClient] = None):
         from prefect.utilities.engine import (
@@ -469,7 +527,8 @@ class TaskRunEngine(Generic[P, R]):
         if not self.task_run:
             raise ValueError("Task run is not set")
 
-        self.task_run = client.read_task_run(self.task_run.id)
+        if not PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            self.task_run = client.read_task_run(self.task_run.id)
         with ExitStack() as stack:
             if log_prints := should_log_prints(self.task):
                 stack.enter_context(patch_print())
@@ -483,23 +542,24 @@ class TaskRunEngine(Generic[P, R]):
                     client=client,
                 )
             )
-            # set the logger to the task run logger
+
             self.logger = task_run_logger(task_run=self.task_run, task=self.task)  # type: ignore
 
-            # update the task run name if necessary
-            if not self._task_name_set and self.task.task_run_name:
-                task_run_name = _resolve_custom_task_run_name(
-                    task=self.task, parameters=self.parameters
-                )
-                self.client.set_task_run_name(
-                    task_run_id=self.task_run.id, name=task_run_name
-                )
-                self.logger.extra["task_run_name"] = task_run_name
-                self.logger.debug(
-                    f"Renamed task run {self.task_run.name!r} to {task_run_name!r}"
-                )
-                self.task_run.name = task_run_name
-                self._task_name_set = True
+            if not PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+                # update the task run name if necessary
+                if not self._task_name_set and self.task.task_run_name:
+                    task_run_name = _resolve_custom_task_run_name(
+                        task=self.task, parameters=self.parameters
+                    )
+                    self.client.set_task_run_name(
+                        task_run_id=self.task_run.id, name=task_run_name
+                    )
+                    self.logger.extra["task_run_name"] = task_run_name
+                    self.logger.debug(
+                        f"Renamed task run {self.task_run.name!r} to {task_run_name!r}"
+                    )
+                    self.task_run.name = task_run_name
+                    self._task_name_set = True
             yield
 
     @contextmanager
@@ -511,22 +571,47 @@ class TaskRunEngine(Generic[P, R]):
         """
         Enters a client context and creates a task run if needed.
         """
+
         with hydrated_context(self.context):
             with ClientContext.get_or_create() as client_ctx:
                 self._client = client_ctx.sync_client
                 self._is_started = True
                 try:
                     if not self.task_run:
-                        self.task_run = run_coro_as_sync(
-                            self.task.create_run(
-                                id=task_run_id,
-                                parameters=self.parameters,
-                                flow_run_context=FlowRunContext.get(),
-                                parent_task_run_context=TaskRunContext.get(),
-                                wait_for=self.wait_for,
-                                extra_task_inputs=dependencies,
+                        if PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+                            # TODO - this maybe should be a method on Task?
+                            from prefect.utilities.engine import (
+                                _resolve_custom_task_run_name,
+                            )
+
+                            task_run_name = None
+                            if not self._task_name_set and self.task.task_run_name:
+                                task_run_name = _resolve_custom_task_run_name(
+                                    task=self.task, parameters=self.parameters
+                                )
+
+                            self.task_run = run_coro_as_sync(
+                                self.task.create_local_run(
+                                    id=task_run_id,
+                                    parameters=self.parameters,
+                                    flow_run_context=FlowRunContext.get(),
+                                    parent_task_run_context=TaskRunContext.get(),
+                                    wait_for=self.wait_for,
+                                    extra_task_inputs=dependencies,
+                                    task_run_name=task_run_name,
+                                )
+                            )
+                        else:
+                            self.task_run = run_coro_as_sync(
+                                self.task.create_run(
+                                    id=task_run_id,
+                                    parameters=self.parameters,
+                                    flow_run_context=FlowRunContext.get(),
+                                    parent_task_run_context=TaskRunContext.get(),
+                                    wait_for=self.wait_for,
+                                    extra_task_inputs=dependencies,
+                                )
                             )
-                        )
                     # Emit an event to capture that the task run was in the `PENDING` state.
                     self._last_event = emit_task_run_state_change_event(
                         task_run=self.task_run,
@@ -916,3 +1001,28 @@ def run_task(
         return run_task_async(**kwargs)
     else:
         return run_task_sync(**kwargs)
+
+
+def _with_transaction_hook_logging(
+    hook: Callable[[Transaction], None],
+    hook_type: Literal["rollback", "commit"],
+    logger: logging.Logger,
+) -> Callable[[Transaction], None]:
+    @wraps(hook)
+    def _hook(txn: Transaction) -> None:
+        hook_name = _get_hook_name(hook)
+        logger.info(f"Running {hook_type} hook {hook_name!r}")
+
+        try:
+            hook(txn)
+        except Exception as exc:
+            logger.error(
+                f"An error was encountered while running {hook_type} hook {hook_name!r}",
+            )
+            raise exc
+        else:
+            logger.info(
+                f"{hook_type.capitalize()} hook {hook_name!r} finished running successfully"
+            )
+
+    return _hook

prefect/task_worker.py

@@ -325,7 +325,7 @@ class TaskWorker:
 
         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}"
+                f"Submitting task run {task_run.name!r} to engine. View in the UI: {task_run_url}"
             )
 
         if task.isasync:

prefect/tasks.py

@@ -6,7 +6,6 @@ Module containing the base workflow task class and decorator - for most use case
 
 import datetime
 import inspect
-import os
 from copy import copy
 from functools import partial, update_wrapper
 from typing import (
@@ -33,13 +32,19 @@ from uuid import UUID, uuid4
 
 from typing_extensions import Literal, ParamSpec
 
+import prefect.states
 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
-from prefect.client.schemas.objects import TaskRunInput, TaskRunResult
+from prefect.client.schemas.objects import (
+    StateDetails,
+    TaskRunInput,
+    TaskRunPolicy,
+    TaskRunResult,
+)
 from prefect.context import (
     FlowRunContext,
     TagsContext,
@@ -50,6 +55,7 @@ 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_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION,
     PREFECT_TASK_DEFAULT_RETRIES,
     PREFECT_TASK_DEFAULT_RETRY_DELAY_SECONDS,
 )
@@ -181,6 +187,31 @@ def _infer_parent_task_runs(
     return parents
 
 
+def _generate_task_key(fn: Callable[..., Any]) -> str:
+    """Generate a task key based on the function name and source code.
+
+    We may eventually want some sort of top-level namespace here to
+    disambiguate tasks with the same function name in different modules,
+    in a more human-readable way, while avoiding relative import problems (see #12337).
+
+    As long as the task implementations are unique (even if named the same), we should
+    not have any collisions.
+
+    Args:
+        fn: The function to generate a task key for.
+    """
+    if not hasattr(fn, "__qualname__"):
+        return to_qualified_name(type(fn))
+
+    qualname = fn.__qualname__.split(".")[-1]
+
+    code_hash = (
+        h[:NUM_CHARS_DYNAMIC_KEY] if (h := hash_objects(fn.__code__)) else "unknown"
+    )
+
+    return f"{qualname}-{code_hash}"
+
+
 class Task(Generic[P, R]):
     """
     A Prefect task definition.
@@ -263,7 +294,7 @@ class Task(Generic[P, R]):
         description: Optional[str] = None,
         tags: Optional[Iterable[str]] = None,
         version: Optional[str] = None,
-        cache_policy: Optional[CachePolicy] = NotSet,
+        cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
         cache_key_fn: Optional[
             Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
         ] = None,
@@ -362,17 +393,7 @@ class Task(Generic[P, R]):
 
         self.tags = set(tags if tags else [])
 
-        if not hasattr(self.fn, "__qualname__"):
-            self.task_key = to_qualified_name(type(self.fn))
-        else:
-            try:
-                task_origin_hash = hash_objects(
-                    self.name, os.path.abspath(inspect.getsourcefile(self.fn))
-                )
-            except TypeError:
-                task_origin_hash = "unknown-source-file"
-
-            self.task_key = f"{self.fn.__qualname__}-{task_origin_hash}"
+        self.task_key = _generate_task_key(self.fn)
 
         if cache_policy is not NotSet and cache_key_fn is not None:
             logger.warning(
@@ -786,6 +807,130 @@ class Task(Generic[P, R]):
 
             return task_run
 
+    async def create_local_run(
+        self,
+        client: Optional["PrefectClient"] = None,
+        id: Optional[UUID] = None,
+        parameters: Optional[Dict[str, Any]] = None,
+        flow_run_context: Optional[FlowRunContext] = None,
+        parent_task_run_context: Optional[TaskRunContext] = None,
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
+        extra_task_inputs: Optional[Dict[str, Set[TaskRunInput]]] = None,
+        deferred: bool = False,
+        task_run_name: Optional[str] = None,
+    ) -> TaskRun:
+        if not PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION:
+            raise RuntimeError(
+                "Cannot call `Task.create_local_run` unless "
+                "PREFECT_EXPERIMENTAL_ENABLE_CLIENT_SIDE_TASK_ORCHESTRATION is True"
+            )
+
+        from prefect.utilities.engine import (
+            _dynamic_key_for_task_run,
+            collect_task_run_inputs_sync,
+        )
+
+        if flow_run_context is None:
+            flow_run_context = FlowRunContext.get()
+        if parent_task_run_context is None:
+            parent_task_run_context = TaskRunContext.get()
+        if parameters is None:
+            parameters = {}
+        if client is None:
+            client = get_client()
+
+        async with client:
+            if not flow_run_context:
+                dynamic_key = f"{self.task_key}-{str(uuid4().hex)}"
+                task_run_name = task_run_name or self.name
+            else:
+                dynamic_key = _dynamic_key_for_task_run(
+                    context=flow_run_context, task=self
+                )
+                task_run_name = task_run_name or f"{self.name}-{dynamic_key}"
+
+            if deferred:
+                state = Scheduled()
+                state.state_details.deferred = True
+            else:
+                state = Pending()
+
+            # store parameters for background tasks so that task worker
+            # can retrieve them at runtime
+            if deferred and (parameters or wait_for):
+                parameters_id = uuid4()
+                state.state_details.task_parameters_id = parameters_id
+
+                # TODO: Improve use of result storage for parameter storage / reference
+                self.persist_result = True
+
+                factory = await ResultFactory.from_autonomous_task(self, client=client)
+                context = serialize_context()
+                data: Dict[str, Any] = {"context": context}
+                if parameters:
+                    data["parameters"] = parameters
+                if wait_for:
+                    data["wait_for"] = wait_for
+                await factory.store_parameters(parameters_id, data)
+
+            # collect task inputs
+            task_inputs = {
+                k: collect_task_run_inputs_sync(v) for k, v in parameters.items()
+            }
+
+            # collect all parent dependencies
+            if task_parents := _infer_parent_task_runs(
+                flow_run_context=flow_run_context,
+                task_run_context=parent_task_run_context,
+                parameters=parameters,
+            ):
+                task_inputs["__parents__"] = task_parents
+
+            # check wait for dependencies
+            if wait_for:
+                task_inputs["wait_for"] = collect_task_run_inputs_sync(wait_for)
+
+            # Join extra task inputs
+            for k, extras in (extra_task_inputs or {}).items():
+                task_inputs[k] = task_inputs[k].union(extras)
+
+            flow_run_id = (
+                getattr(flow_run_context.flow_run, "id", None)
+                if flow_run_context and flow_run_context.flow_run
+                else None
+            )
+            task_run_id = id or uuid4()
+            state = prefect.states.Pending(
+                state_details=StateDetails(
+                    task_run_id=task_run_id,
+                    flow_run_id=flow_run_id,
+                )
+            )
+            task_run = TaskRun(
+                id=task_run_id,
+                name=task_run_name,
+                flow_run_id=flow_run_id,
+                task_key=self.task_key,
+                dynamic_key=str(dynamic_key),
+                task_version=self.version,
+                empirical_policy=TaskRunPolicy(
+                    retries=self.retries,
+                    retry_delay=self.retry_delay_seconds,
+                    retry_jitter_factor=self.retry_jitter_factor,
+                ),
+                tags=list(set(self.tags).union(TagsContext.get().current_tags or [])),
+                task_inputs=task_inputs or {},
+                expected_start_time=state.timestamp,
+                state_id=state.id,
+                state_type=state.type,
+                state_name=state.name,
+                state=state,
+                created=state.timestamp,
+                updated=state.timestamp,
+            )
+
+            return task_run
+
     @overload
     def __call__(
         self: "Task[P, NoReturn]",
@@ -1365,7 +1510,7 @@ class Task(Generic[P, R]):
 
         Args:
             task_runner: The task runner to use for serving the task. If not provided,
-                the default ConcurrentTaskRunner will be used.
+                the default task runner will be used.
 
         Examples:
             Serve a task using the default task runner
@@ -1392,7 +1537,7 @@ def task(
     description: Optional[str] = None,
     tags: Optional[Iterable[str]] = None,
     version: Optional[str] = None,
-    cache_policy: CachePolicy = NotSet,
+    cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
     cache_key_fn: Optional[
         Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
     ] = None,
@@ -1430,7 +1575,9 @@ def task(
     tags: Optional[Iterable[str]] = None,
     version: Optional[str] = None,
     cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
-    cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]] = None,
+    cache_key_fn: Union[
+        Callable[["TaskRunContext", Dict[str, Any]], Optional[str]], None
+    ] = None,
     cache_expiration: Optional[datetime.timedelta] = None,
     task_run_name: Optional[Union[Callable[[], str], str]] = None,
     retries: Optional[int] = None,

prefect/transactions.py

@@ -26,7 +26,6 @@ from prefect.results import (
 )
 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):
@@ -180,13 +179,10 @@ 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:
@@ -195,16 +191,8 @@ class Transaction(ContextModel):
             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,
+                    f"An error was encountered while committing transaction {self.key!r}",
                     exc_info=True,
                 )
             self.rollback()
@@ -234,7 +222,6 @@ 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
@@ -246,7 +233,7 @@ class Transaction(ContextModel):
         except Exception:
             if self.logger:
                 self.logger.exception(
-                    f"An error was encountered while running rollback hook {hook_name!r}",
+                    f"An error was encountered while rolling back transaction {self.key!r}",
                     exc_info=True,
                 )
             return False