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

prefect/flow_engine.py

@@ -7,7 +7,6 @@ from dataclasses import dataclass, field
 from typing import (
     Any,
     AsyncGenerator,
-    Callable,
     Coroutine,
     Dict,
     Generator,
@@ -16,6 +15,7 @@ from typing import (
     Literal,
     Optional,
     Tuple,
+    Type,
     TypeVar,
     Union,
     cast,
@@ -30,7 +30,13 @@ from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.filters import FlowRunFilter
 from prefect.client.schemas.sorting import FlowRunSort
 from prefect.context import ClientContext, FlowRunContext, TagsContext
-from prefect.exceptions import Abort, Pause, PrefectException, UpstreamTaskError
+from prefect.exceptions import (
+    Abort,
+    Pause,
+    PrefectException,
+    TerminationSignal,
+    UpstreamTaskError,
+)
 from prefect.flows import Flow, load_flow_from_entrypoint, load_flow_from_flow_run
 from prefect.futures import PrefectFuture, resolve_futures_to_states
 from prefect.logging.loggers import (
@@ -39,7 +45,7 @@ from prefect.logging.loggers import (
     get_run_logger,
     patch_print,
 )
-from prefect.results import ResultFactory
+from prefect.results import BaseResult, ResultFactory
 from prefect.settings import PREFECT_DEBUG_MODE
 from prefect.states import (
     Failed,
@@ -50,6 +56,7 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
+from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import (
     call_with_parameters,
@@ -72,6 +79,10 @@ P = ParamSpec("P")
 R = TypeVar("R")
 
 
+class FlowRunTimeoutError(TimeoutError):
+    """Raised when a flow run exceeds its defined timeout."""
+
+
 def load_flow_and_flow_run(flow_run_id: UUID) -> Tuple[FlowRun, Flow]:
     ## TODO: add error handling to update state and log tracebacks
     entrypoint = os.environ.get("PREFECT__FLOW_ENTRYPOINT")
@@ -95,6 +106,10 @@ class FlowRunEngine(Generic[P, R]):
     flow_run_id: Optional[UUID] = None
     logger: logging.Logger = field(default_factory=lambda: get_logger("engine"))
     wait_for: Optional[Iterable[PrefectFuture]] = 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
     _is_started: bool = False
     _client: Optional[SyncPrefectClient] = None
     short_circuit: bool = False
@@ -208,6 +223,30 @@ class FlowRunEngine(Generic[P, R]):
         return state
 
     def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
+        if self._return_value is not NotSet and not isinstance(
+            self._return_value, State
+        ):
+            if isinstance(self._return_value, BaseResult):
+                _result = self._return_value.get()
+            else:
+                _result = self._return_value
+
+            if inspect.isawaitable(_result):
+                # getting the value for a BaseResult may return an awaitable
+                # depending on whether the parent frame is sync or not
+                _result = run_coro_as_sync(_result)
+            return _result
+
+        if self._raised is not NotSet:
+            if raise_on_failure:
+                raise self._raised
+            return self._raised
+
+        # This is a fall through case which leans on the existing state result mechanics to get the
+        # return value. This is necessary because we currently will return a State object if the
+        # the State was Prefect-created.
+        # TODO: Remove the need to get the result from a State except in cases where the return value
+        # is a State object.
         _result = self.state.result(raise_on_failure=raise_on_failure, fetch=True)  # type: ignore
         # 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
@@ -219,13 +258,15 @@ class FlowRunEngine(Generic[P, R]):
         result_factory = getattr(FlowRunContext.get(), "result_factory", None)
         if result_factory is None:
             raise ValueError("Result factory is not set")
+        resolved_result = resolve_futures_to_states(result)
         terminal_state = run_coro_as_sync(
             return_value_to_state(
-                resolve_futures_to_states(result),
+                resolved_result,
                 result_factory=result_factory,
             )
         )
         self.set_state(terminal_state)
+        self._return_value = resolved_result
         return result
 
     def handle_exception(
@@ -252,10 +293,16 @@ class FlowRunEngine(Generic[P, R]):
                 ),
             )
             state = self.set_state(Running())
+        self._raised = exc
         return state
 
     def handle_timeout(self, exc: TimeoutError) -> None:
-        message = f"Flow run exceeded timeout of {self.flow.timeout_seconds} seconds"
+        if isinstance(exc, FlowRunTimeoutError):
+            message = (
+                f"Flow run exceeded timeout of {self.flow.timeout_seconds} second(s)"
+            )
+        else:
+            message = f"Flow run failed due to timeout: {exc!r}"
         self.logger.error(message)
         state = Failed(
             data=exc,
@@ -263,12 +310,14 @@ class FlowRunEngine(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
 
     def load_subflow_run(
         self,
@@ -315,7 +364,9 @@ class FlowRunEngine(Generic[P, R]):
                 limit=1,
             )
             if flow_runs:
-                return flow_runs[-1]
+                loaded_flow_run = flow_runs[-1]
+                self._return_value = loaded_flow_run.state
+                return loaded_flow_run
 
     def create_flow_run(self, client: SyncPrefectClient) -> FlowRun:
         flow_run_ctx = FlowRunContext.get()
@@ -363,7 +414,7 @@ class FlowRunEngine(Generic[P, R]):
 
         return flow_run
 
-    def call_hooks(self, state: State = None) -> Iterable[Callable]:
+    def call_hooks(self, state: Optional[State] = None):
         if state is None:
             state = self.state
         flow = self.flow
@@ -504,6 +555,11 @@ class FlowRunEngine(Generic[P, R]):
                 )
             try:
                 yield self
+
+            except TerminationSignal as exc:
+                self.cancel_all_tasks()
+                self.handle_crash(exc)
+                raise
             except Exception:
                 # regular exceptions are caught and re-raised to the user
                 raise
@@ -539,6 +595,10 @@ class FlowRunEngine(Generic[P, R]):
             return False  # TODO: handle this differently?
         return getattr(self, "flow_run").state.is_pending()
 
+    def cancel_all_tasks(self):
+        if hasattr(self.flow.task_runner, "cancel_all"):
+            self.flow.task_runner.cancel_all()  # type: ignore
+
     # --------------------------
     #
     # The following methods compose the main task run loop
@@ -552,11 +612,7 @@ class FlowRunEngine(Generic[P, R]):
 
             if self.state.is_running():
                 self.call_hooks()
-            try:
-                yield
-            finally:
-                if self.state.is_final() or self.state.is_cancelling():
-                    self.call_hooks()
+            yield
 
     @contextmanager
     def run_context(self):
@@ -564,7 +620,10 @@ class FlowRunEngine(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.flow.timeout_seconds):
+                with timeout_context(
+                    seconds=self.flow.timeout_seconds,
+                    timeout_exc_type=FlowRunTimeoutError,
+                ):
                     self.logger.debug(
                         f"Executing flow {self.flow.name!r} for flow run {self.flow_run.name!r}..."
                     )
@@ -572,8 +631,11 @@ class FlowRunEngine(Generic[P, R]):
             except TimeoutError as exc:
                 self.handle_timeout(exc)
             except Exception as exc:
-                self.logger.exception(f"Encountered exception during execution: {exc}")
+                self.logger.exception("Encountered exception during execution: %r", exc)
                 self.handle_exception(exc)
+            finally:
+                if self.state.is_final() or self.state.is_cancelling():
+                    self.call_hooks()
 
     def call_flow_fn(self) -> Union[R, Coroutine[Any, Any, R]]:
         """

prefect/flows.py

@@ -5,6 +5,7 @@ Module containing the base workflow class and decorator - for most use cases, us
 # This file requires type-checking with pyright because mypy does not yet support PEP612
 # See https://github.com/python/mypy/issues/8645
 import ast
+import asyncio
 import datetime
 import importlib.util
 import inspect
@@ -28,6 +29,8 @@ from typing import (
     List,
     NoReturn,
     Optional,
+    Set,
+    Tuple,
     Type,
     TypeVar,
     Union,
@@ -44,7 +47,9 @@ from pydantic.v1.errors import ConfigError  # TODO
 from rich.console import Console
 from typing_extensions import Literal, ParamSpec, Self
 
-from prefect._internal.compatibility.deprecated import deprecated_parameter
+from prefect._internal.compatibility.deprecated import (
+    deprecated_parameter,
+)
 from prefect._internal.concurrency.api import create_call, from_async
 from prefect.blocks.core import Block
 from prefect.client.orchestration import get_client
@@ -60,6 +65,7 @@ from prefect.exceptions import (
     MissingFlowError,
     ObjectNotFound,
     ParameterTypeError,
+    ScriptError,
     UnspecifiedFlowError,
 )
 from prefect.filesystems import LocalFileSystem, ReadableDeploymentStorage
@@ -187,7 +193,7 @@ class Flow(Generic[P, R]):
         timeout_seconds: Union[int, float, None] = None,
         validate_parameters: bool = True,
         persist_result: Optional[bool] = None,
-        result_storage: Optional[ResultStorage] = None,
+        result_storage: Optional[Union[ResultStorage, str]] = None,
         result_serializer: Optional[ResultSerializer] = None,
         cache_result_in_memory: bool = True,
         log_prints: Optional[bool] = None,
@@ -335,7 +341,18 @@ class Flow(Generic[P, R]):
                     "Disable validation or change the argument names."
                 ) from exc
 
+        # result persistence settings
+        if persist_result is None:
+            if result_storage is not None or result_serializer is not None:
+                persist_result = True
+
         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.cache_result_in_memory = cache_result_in_memory
@@ -767,8 +784,7 @@ class Flow(Generic[P, R]):
         self.on_failure_hooks.append(fn)
         return fn
 
-    @sync_compatible
-    async def serve(
+    def serve(
         self,
         name: Optional[str] = None,
         interval: Optional[
@@ -873,7 +889,7 @@ class Flow(Generic[P, R]):
             name = Path(name).stem
 
         runner = Runner(name=name, pause_on_shutdown=pause_on_shutdown, limit=limit)
-        deployment_id = await runner.add_flow(
+        deployment_id = runner.add_flow(
             self,
             name=name,
             triggers=triggers,
@@ -906,15 +922,27 @@ class Flow(Generic[P, R]):
 
             console = Console()
             console.print(help_message, soft_wrap=True)
-        await runner.start(webserver=webserver)
+
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError as exc:
+            if "no running event loop" in str(exc):
+                loop = None
+            else:
+                raise
+
+        if loop is not None:
+            loop.run_until_complete(runner.start(webserver=webserver))
+        else:
+            asyncio.run(runner.start(webserver=webserver))
 
     @classmethod
     @sync_compatible
     async def from_source(
-        cls: Type[F],
+        cls: Type["Flow[P, R]"],
         source: Union[str, "RunnerStorage", ReadableDeploymentStorage],
         entrypoint: str,
-    ) -> F:
+    ) -> "Flow[P, R]":
         """
         Loads a flow from a remote source.
 
@@ -960,6 +988,29 @@ class Flow(Generic[P, R]):
 
             my_flow()
             ```
+
+            Load a flow from a local directory:
+
+            ``` python
+            # from_local_source.py
+
+            from pathlib import Path
+            from prefect import flow
+
+            @flow(log_prints=True)
+            def my_flow(name: str = "world"):
+                print(f"Hello {name}! I'm a flow from a Python script!")
+
+            if __name__ == "__main__":
+                my_flow.from_source(
+                    source=str(Path(__file__).parent),
+                    entrypoint="from_local_source.py:my_flow",
+                ).deploy(
+                    name="my-deployment",
+                    parameters=dict(name="Marvin"),
+                    work_pool_name="local",
+                )
+            ```
         """
 
         from prefect.runner.storage import (
@@ -969,7 +1020,9 @@ class Flow(Generic[P, R]):
             create_storage_from_source,
         )
 
-        if isinstance(source, str):
+        if isinstance(source, (Path, str)):
+            if isinstance(source, Path):
+                source = str(source)
             storage = create_storage_from_source(source)
         elif isinstance(source, RunnerStorage):
             storage = source
@@ -988,7 +1041,7 @@ class Flow(Generic[P, R]):
             await storage.pull_code()
 
             full_entrypoint = str(storage.destination / entrypoint)
-            flow: "Flow" = await from_async.wait_for_call_in_new_thread(
+            flow: Flow = await from_async.wait_for_call_in_new_thread(
                 create_call(load_flow_from_entrypoint, full_entrypoint)
             )
             flow._storage = storage
@@ -1115,7 +1168,13 @@ class Flow(Generic[P, R]):
                 )
             ```
         """
-        work_pool_name = work_pool_name or PREFECT_DEFAULT_WORK_POOL_NAME.value()
+        if not (
+            work_pool_name := work_pool_name or PREFECT_DEFAULT_WORK_POOL_NAME.value()
+        ):
+            raise ValueError(
+                "No work pool name provided. Please provide a `work_pool_name` or set the"
+                " `PREFECT_DEFAULT_WORK_POOL_NAME` environment variable."
+            )
 
         try:
             async with get_client() as client:
@@ -1146,9 +1205,9 @@ class Flow(Generic[P, R]):
             entrypoint_type=entrypoint_type,
         )
 
-        from prefect.deployments import runner
+        from prefect.deployments.runner import deploy
 
-        deployment_ids = await runner.deploy(
+        deployment_ids = await deploy(
             deployment,
             work_pool_name=work_pool_name,
             image=image,
@@ -1672,6 +1731,14 @@ def load_flow_from_entrypoint(
         raise MissingFlowError(
             f"Flow function with name {func_name!r} not found in {path!r}. "
         ) from exc
+    except ScriptError as exc:
+        # If the flow has dependencies that are not installed in the current
+        # environment, fallback to loading the flow via AST parsing. The
+        # drawback of this approach is that we're unable to actually load the
+        # function, so we create a placeholder flow that will re-raise this
+        # exception when called.
+
+        flow = load_placeholder_flow(entrypoint=entrypoint, raises=exc)
 
     if not isinstance(flow, Flow):
         raise MissingFlowError(
@@ -1682,14 +1749,13 @@ def load_flow_from_entrypoint(
     return flow
 
 
-@sync_compatible
-async def serve(
+def serve(
     *args: "RunnerDeployment",
     pause_on_shutdown: bool = True,
     print_starting_message: bool = True,
     limit: Optional[int] = None,
     **kwargs,
-) -> NoReturn:
+):
     """
     Serve the provided list of deployments.
 
@@ -1739,7 +1805,7 @@ async def serve(
 
     runner = Runner(pause_on_shutdown=pause_on_shutdown, limit=limit, **kwargs)
     for deployment in args:
-        await runner.add_deployment(deployment)
+        runner.add_deployment(deployment)
 
     if print_starting_message:
         help_message_top = (
@@ -1770,7 +1836,18 @@ async def serve(
             Group(help_message_top, table, help_message_bottom), soft_wrap=True
         )
 
-    await runner.start()
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError as exc:
+        if "no running event loop" in str(exc):
+            loop = None
+        else:
+            raise
+
+    if loop is not None:
+        loop.run_until_complete(runner.start())
+    else:
+        asyncio.run(runner.start())
 
 
 @client_injector
@@ -1852,24 +1929,138 @@ async def load_flow_from_flow_run(
     return flow
 
 
-def load_flow_argument_from_entrypoint(
-    entrypoint: str, arg: str = "name"
-) -> Optional[str]:
+def load_placeholder_flow(entrypoint: str, raises: Exception):
     """
-    Extract a flow argument from an entrypoint string.
+    Load a placeholder flow that is initialized with the same arguments as the
+    flow specified in the entrypoint. If called the flow will raise `raises`.
 
-    Loads the source code of the entrypoint and extracts the flow argument from the
-    `flow` decorator.
+    This is useful when a flow can't be loaded due to missing dependencies or
+    other issues but the base metadata defining the flow is still needed.
 
     Args:
-        entrypoint: a string in the format `<path_to_script>:<flow_func_name>` or a module path
-            to a flow function
+        entrypoint: a string in the format `<path_to_script>:<flow_func_name>`
+          or a module path to a flow function
+        raises: an exception to raise when the flow is called
+    """
+
+    def _base_placeholder():
+        raise raises
+
+    def sync_placeholder_flow(*args, **kwargs):
+        _base_placeholder()
+
+    async def async_placeholder_flow(*args, **kwargs):
+        _base_placeholder()
+
+    placeholder_flow = (
+        async_placeholder_flow
+        if is_entrypoint_async(entrypoint)
+        else sync_placeholder_flow
+    )
+
+    arguments = load_flow_arguments_from_entrypoint(entrypoint)
+    arguments["fn"] = placeholder_flow
+
+    return Flow(**arguments)
+
+
+def load_flow_arguments_from_entrypoint(
+    entrypoint: str, arguments: Optional[Union[List[str], Set[str]]] = None
+) -> dict[str, Any]:
+    """
+    Extract flow arguments from an entrypoint string.
+
+    Loads the source code of the entrypoint and extracts the flow arguments
+    from the `flow` decorator.
+
+    Args:
+        entrypoint: a string in the format `<path_to_script>:<flow_func_name>`
+          or a module path to a flow function
+    """
+
+    func_def, source_code = _entrypoint_definition_and_source(entrypoint)
+
+    if arguments is None:
+        # If no arguments are provided default to known arguments that are of
+        # built-in types.
+        arguments = {
+            "name",
+            "version",
+            "retries",
+            "retry_delay_seconds",
+            "description",
+            "timeout_seconds",
+            "validate_parameters",
+            "persist_result",
+            "cache_result_in_memory",
+            "log_prints",
+        }
+
+    result = {}
+
+    for decorator in func_def.decorator_list:
+        if (
+            isinstance(decorator, ast.Call)
+            and getattr(decorator.func, "id", "") == "flow"
+        ):
+            for keyword in decorator.keywords:
+                if keyword.arg not in arguments:
+                    continue
+
+                if isinstance(keyword.value, ast.Constant):
+                    # Use the string value of the argument
+                    result[keyword.arg] = str(keyword.value.value)
+                    continue
+
+                # if the arg value is not a raw str (i.e. a variable or expression),
+                # then attempt to evaluate it
+                namespace = safe_load_namespace(source_code)
+                literal_arg_value = ast.get_source_segment(source_code, keyword.value)
+                cleaned_value = (
+                    literal_arg_value.replace("\n", "") if literal_arg_value else ""
+                )
+
+                try:
+                    evaluated_value = eval(cleaned_value, namespace)  # type: ignore
+                    result[keyword.arg] = str(evaluated_value)
+                except Exception as e:
+                    logger.info(
+                        "Failed to parse @flow argument: `%s=%s` due to the following error. Ignoring and falling back to default behavior.",
+                        keyword.arg,
+                        literal_arg_value,
+                        exc_info=e,
+                    )
+                    # ignore the decorator arg and fallback to default behavior
+                    continue
+
+    if "name" in arguments and "name" not in result:
+        # If no matching decorator or keyword argument for `name' is found
+        # fallback to the function name.
+        result["name"] = func_def.name.replace("_", "-")
+
+    return result
+
+
+def is_entrypoint_async(entrypoint: str) -> bool:
+    """
+    Determine if the function specified in the entrypoint is asynchronous.
+
+    Args:
+        entrypoint: A string in the format `<path_to_script>:<func_name>` or
+          a module path to a function.
 
     Returns:
-        The flow argument value
+        True if the function is asynchronous, False otherwise.
     """
+    func_def, _ = _entrypoint_definition_and_source(entrypoint)
+    return isinstance(func_def, ast.AsyncFunctionDef)
+
+
+def _entrypoint_definition_and_source(
+    entrypoint: str,
+) -> Tuple[Union[ast.FunctionDef, ast.AsyncFunctionDef], str]:
     if ":" in entrypoint:
-        # split by the last colon once to handle Windows paths with drive letters i.e C:\path\to\file.py:do_stuff
+        # Split by the last colon once to handle Windows paths with drive letters i.e C:\path\to\file.py:do_stuff
         path, func_name = entrypoint.rsplit(":", maxsplit=1)
         source_code = Path(path).read_text()
     else:
@@ -1878,6 +2069,7 @@ def load_flow_argument_from_entrypoint(
         if not spec or not spec.origin:
             raise ValueError(f"Could not find module {path!r}")
         source_code = Path(spec.origin).read_text()
+
     parsed_code = ast.parse(source_code)
     func_def = next(
         (
@@ -1894,42 +2086,8 @@ def load_flow_argument_from_entrypoint(
         ),
         None,
     )
+
     if not func_def:
         raise ValueError(f"Could not find flow {func_name!r} in {path!r}")
-    for decorator in func_def.decorator_list:
-        if (
-            isinstance(decorator, ast.Call)
-            and getattr(decorator.func, "id", "") == "flow"
-        ):
-            for keyword in decorator.keywords:
-                if keyword.arg == arg:
-                    if isinstance(keyword.value, ast.Constant):
-                        return (
-                            keyword.value.value
-                        )  # Return the string value of the argument
-
-                    # if the arg value is not a raw str (i.e. a variable or expression),
-                    # then attempt to evaluate it
-                    namespace = safe_load_namespace(source_code)
-                    literal_arg_value = ast.get_source_segment(
-                        source_code, keyword.value
-                    )
-                    try:
-                        evaluated_value = eval(literal_arg_value, namespace)  # type: ignore
-                    except Exception as e:
-                        logger.info(
-                            "Failed to parse @flow argument: `%s=%s` due to the following error. Ignoring and falling back to default behavior.",
-                            arg,
-                            literal_arg_value,
-                            exc_info=e,
-                        )
-                        # ignore the decorator arg and fallback to default behavior
-                        break
-                    return str(evaluated_value)
-
-    if arg == "name":
-        return func_name.replace(
-            "_", "-"
-        )  # If no matching decorator or keyword argument is found
 
-    return None
+    return func_def, source_code