prefect-client 3.1.4__py3-none-any.whl → 3.1.6__py3-none-any.whl

prefect/flow_engine.py

@@ -2,7 +2,7 @@ import asyncio
 import logging
 import os
 import time
-from contextlib import ExitStack, contextmanager
+from contextlib import ExitStack, asynccontextmanager, contextmanager, nullcontext
 from dataclasses import dataclass, field
 from typing import (
     Any,
@@ -22,19 +22,25 @@ from typing import (
 )
 from uuid import UUID
 
-from opentelemetry import trace
+from anyio import CancelScope
+from opentelemetry import propagate, trace
 from opentelemetry.trace import Tracer, get_tracer
 from typing_extensions import ParamSpec
 
 import prefect
 from prefect import Task
-from prefect.client.orchestration import SyncPrefectClient, get_client
+from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.filters import FlowRunFilter
 from prefect.client.schemas.sorting import FlowRunSort
 from prefect.concurrency.context import ConcurrencyContext
 from prefect.concurrency.v1.context import ConcurrencyContext as ConcurrencyContextV1
-from prefect.context import FlowRunContext, SyncClientContext, TagsContext
+from prefect.context import (
+    AsyncClientContext,
+    FlowRunContext,
+    SyncClientContext,
+    TagsContext,
+)
 from prefect.exceptions import (
     Abort,
     Pause,
@@ -66,6 +72,9 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
+from prefect.telemetry.run_telemetry import OTELSetter
+from prefect.types import KeyValueLabels
+from prefect.utilities._engine import get_hook_name, resolve_custom_flow_run_name
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import (
@@ -75,10 +84,9 @@ from prefect.utilities.callables import (
 )
 from prefect.utilities.collections import visit_collection
 from prefect.utilities.engine import (
-    _get_hook_name,
-    _resolve_custom_flow_run_name,
     capture_sigterm,
     link_state_to_result,
+    propose_state,
     propose_state_sync,
     resolve_to_final_result,
 )
@@ -87,6 +95,8 @@ from prefect.utilities.urls import url_for
 
 P = ParamSpec("P")
 R = TypeVar("R")
+LABELS_TRACEPARENT_KEY = "__OTEL_TRACEPARENT"
+TRACEPARENT_KEY = "traceparent"
 
 
 class FlowRunTimeoutError(TimeoutError):
@@ -112,7 +122,7 @@ def load_flow_and_flow_run(flow_run_id: UUID) -> Tuple[FlowRun, Flow]:
 
 
 @dataclass
-class FlowRunEngine(Generic[P, R]):
+class BaseFlowRunEngine(Generic[P, R]):
     flow: Union[Flow[P, R], Flow[P, Coroutine[Any, Any, R]]]
     parameters: Optional[Dict[str, Any]] = None
     flow_run: Optional[FlowRun] = None
@@ -124,7 +134,6 @@ class FlowRunEngine(Generic[P, R]):
     # 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
     _flow_run_name_set: bool = False
     _tracer: Tracer = field(
@@ -139,16 +148,81 @@ class FlowRunEngine(Generic[P, R]):
         if self.parameters is None:
             self.parameters = {}
 
+    @property
+    def state(self) -> State:
+        return self.flow_run.state  # type: ignore
+
+    def _end_span_on_success(self):
+        if not self._span:
+            return
+        self._span.set_status(trace.Status(trace.StatusCode.OK))
+        self._span.end(time.time_ns())
+        self._span = None
+
+    def _end_span_on_error(self, exc: BaseException, description: Optional[str]):
+        if not self._span:
+            return
+        self._span.record_exception(exc)
+        self._span.set_status(trace.Status(trace.StatusCode.ERROR, description))
+        self._span.end(time.time_ns())
+        self._span = None
+
+    def is_running(self) -> bool:
+        if getattr(self, "flow_run", None) is None:
+            return False
+        return getattr(self, "flow_run").state.is_running()
+
+    def is_pending(self) -> bool:
+        if getattr(self, "flow_run", None) is None:
+            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
+
+    def _update_otel_labels(
+        self, span: trace.Span, client: Union[SyncPrefectClient, PrefectClient]
+    ):
+        parent_flow_run_ctx = FlowRunContext.get()
+        if parent_flow_run_ctx and parent_flow_run_ctx.flow_run:
+            if traceparent := parent_flow_run_ctx.flow_run.labels.get(
+                LABELS_TRACEPARENT_KEY
+            ):
+                carrier: KeyValueLabels = {TRACEPARENT_KEY: traceparent}
+                propagate.get_global_textmap().inject(
+                    carrier={TRACEPARENT_KEY: traceparent},
+                    setter=OTELSetter(),
+                )
+            else:
+                carrier: KeyValueLabels = {}
+                propagate.get_global_textmap().inject(
+                    carrier,
+                    context=trace.set_span_in_context(span),
+                    setter=OTELSetter(),
+                )
+            if carrier.get(TRACEPARENT_KEY):
+                if self.flow_run:
+                    client.update_flow_run_labels(
+                        flow_run_id=self.flow_run.id,
+                        labels={LABELS_TRACEPARENT_KEY: carrier[TRACEPARENT_KEY]},
+                    )
+                else:
+                    self.logger.info(
+                        f"Tried to set traceparent {carrier[TRACEPARENT_KEY]} for flow run, but None was found"
+                    )
+
+
+@dataclass
+class FlowRunEngine(BaseFlowRunEngine[P, R]):
+    _client: Optional[SyncPrefectClient] = None
+
     @property
     def client(self) -> SyncPrefectClient:
         if not self._is_started or self._client is None:
             raise RuntimeError("Engine has not started.")
         return self._client
 
-    @property
-    def state(self) -> State:
-        return self.flow_run.state  # type: ignore
-
     def _resolve_parameters(self):
         if not self.parameters:
             return {}
@@ -243,7 +317,7 @@ class FlowRunEngine(Generic[P, R]):
 
         if self._span:
             self._span.add_event(
-                state.name,
+                state.name or state.type,
                 {
                     "prefect.state.message": state.message or "",
                     "prefect.state.type": state.type,
@@ -362,22 +436,7 @@ class FlowRunEngine(Generic[P, R]):
         self.set_state(state, force=True)
         self._raised = exc
 
-        self._end_span_on_error(exc, state.message)
-
-    def _end_span_on_success(self):
-        if not self._span:
-            return
-        self._span.set_status(trace.Status(trace.StatusCode.OK))
-        self._span.end(time.time_ns())
-        self._span = None
-
-    def _end_span_on_error(self, exc: BaseException, description: Optional[str]):
-        if not self._span:
-            return
-        self._span.record_exception(exc)
-        self._span.set_status(trace.Status(trace.StatusCode.ERROR, description))
-        self._span.end(time.time_ns())
-        self._span = None
+        self._end_span_on_error(exc, state.message if state else "")
 
     def load_subflow_run(
         self,
@@ -512,7 +571,7 @@ class FlowRunEngine(Generic[P, R]):
             hooks = None
 
         for hook in hooks or []:
-            hook_name = _get_hook_name(hook)
+            hook_name = get_hook_name(hook)
 
             try:
                 self.logger.info(
@@ -575,7 +634,7 @@ class FlowRunEngine(Generic[P, R]):
 
             # update the flow run name if necessary
             if not self._flow_run_name_set and self.flow.flow_run_name:
-                flow_run_name = _resolve_custom_flow_run_name(
+                flow_run_name = resolve_custom_flow_run_name(
                     flow=self.flow, parameters=self.parameters
                 )
                 self.client.set_flow_run_name(
@@ -622,7 +681,7 @@ class FlowRunEngine(Generic[P, R]):
                     empirical_policy=self.flow_run.empirical_policy,
                 )
 
-            self._span = self._tracer.start_span(
+            span = self._tracer.start_span(
                 name=self.flow_run.name,
                 attributes={
                     **self.flow_run.labels,
@@ -632,6 +691,9 @@ class FlowRunEngine(Generic[P, R]):
                     "prefect.flow.name": self.flow.name,
                 },
             )
+            self._update_otel_labels(span, self.client)
+
+            self._span = span
 
             try:
                 yield self
@@ -665,20 +727,6 @@ class FlowRunEngine(Generic[P, R]):
                 self._is_started = False
                 self._client = None
 
-    def is_running(self) -> bool:
-        if getattr(self, "flow_run", None) is None:
-            return False
-        return getattr(self, "flow_run").state.is_running()
-
-    def is_pending(self) -> bool:
-        if getattr(self, "flow_run", None) is None:
-            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
@@ -687,12 +735,13 @@ class FlowRunEngine(Generic[P, R]):
 
     @contextmanager
     def start(self) -> Generator[None, None, None]:
-        with self.initialize_run(), trace.use_span(self._span):
-            self.begin_run()
+        with self.initialize_run():
+            with trace.use_span(self._span) if self._span else nullcontext():
+                self.begin_run()
 
-            if self.state.is_running():
-                self.call_hooks()
-            yield
+                if self.state.is_running():
+                    self.call_hooks()
+                yield
 
     @contextmanager
     def run_context(self):
@@ -734,126 +783,694 @@ class FlowRunEngine(Generic[P, R]):
             self.handle_success(result)
 
 
-def run_flow_sync(
-    flow: Flow[P, R],
-    flow_run: Optional[FlowRun] = None,
-    parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
-    return_type: Literal["state", "result"] = "result",
-) -> Union[R, State, None]:
-    engine = FlowRunEngine[P, R](
-        flow=flow,
-        parameters=parameters,
-        flow_run=flow_run,
-        wait_for=wait_for,
-    )
+@dataclass
+class AsyncFlowRunEngine(BaseFlowRunEngine[P, R]):
+    """
+    Async version of the flow run engine.
 
-    with engine.start():
-        while engine.is_running():
-            with engine.run_context():
-                engine.call_flow_fn()
+    NOTE: This has not been fully asyncified yet which may lead to async flows
+    not being fully asyncified.
+    """
 
-    return engine.state if return_type == "state" else engine.result()
+    _client: Optional[PrefectClient] = None
 
+    @property
+    def client(self) -> PrefectClient:
+        if not self._is_started or self._client is None:
+            raise RuntimeError("Engine has not started.")
+        return self._client
 
-async def run_flow_async(
-    flow: Flow[P, R],
-    flow_run: Optional[FlowRun] = None,
-    parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
-    return_type: Literal["state", "result"] = "result",
-) -> Union[R, State, None]:
-    engine = FlowRunEngine[P, R](
-        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
-    )
+    def _resolve_parameters(self):
+        if not self.parameters:
+            return {}
 
-    with engine.start():
-        while engine.is_running():
-            with engine.run_context():
-                await engine.call_flow_fn()
+        resolved_parameters = {}
+        for parameter, value in self.parameters.items():
+            try:
+                resolved_parameters[parameter] = visit_collection(
+                    value,
+                    visit_fn=resolve_to_final_result,
+                    return_data=True,
+                    max_depth=-1,
+                    remove_annotations=True,
+                    context={},
+                )
+            except UpstreamTaskError:
+                raise
+            except Exception as exc:
+                raise PrefectException(
+                    f"Failed to resolve inputs in parameter {parameter!r}. If your"
+                    " parameter type is not supported, consider using the `quote`"
+                    " annotation to skip resolution of inputs."
+                ) from exc
 
-    return engine.state if return_type == "state" else engine.result()
+        self.parameters = resolved_parameters
 
+    def _wait_for_dependencies(self):
+        if not self.wait_for:
+            return
 
-def run_generator_flow_sync(
-    flow: Flow[P, R],
-    flow_run: Optional[FlowRun] = None,
-    parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
-    return_type: Literal["state", "result"] = "result",
-) -> Generator[R, None, None]:
-    if return_type != "result":
-        raise ValueError("The return_type for a generator flow must be 'result'")
+        visit_collection(
+            self.wait_for,
+            visit_fn=resolve_to_final_result,
+            return_data=False,
+            max_depth=-1,
+            remove_annotations=True,
+            context={},
+        )
 
-    engine = FlowRunEngine[P, R](
-        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
-    )
+    async def begin_run(self) -> State:
+        try:
+            self._resolve_parameters()
+            self._wait_for_dependencies()
+        except UpstreamTaskError as upstream_exc:
+            state = await self.set_state(
+                Pending(
+                    name="NotReady",
+                    message=str(upstream_exc),
+                ),
+                # if orchestrating a run already in a pending state, force orchestration to
+                # update the state name
+                force=self.state.is_pending(),
+            )
+            return state
 
-    with engine.start():
-        while engine.is_running():
-            with engine.run_context():
-                call_args, call_kwargs = parameters_to_args_kwargs(
-                    flow.fn, engine.parameters or {}
+        # validate prior to context so that context receives validated params
+        if self.flow.should_validate_parameters:
+            try:
+                self.parameters = self.flow.validate_parameters(self.parameters or {})
+            except Exception as exc:
+                message = "Validation of flow parameters failed with error:"
+                self.logger.error("%s %s", message, exc)
+                await self.handle_exception(
+                    exc,
+                    msg=message,
+                    result_store=get_result_store().update_for_flow(
+                        self.flow, _sync=True
+                    ),
                 )
-                gen = flow.fn(*call_args, **call_kwargs)
-                try:
-                    while True:
-                        gen_result = next(gen)
-                        # link the current state to the result for dependency tracking
-                        link_state_to_result(engine.state, gen_result)
-                        yield gen_result
-                except StopIteration as exc:
-                    engine.handle_success(exc.value)
-                except GeneratorExit as exc:
-                    engine.handle_success(None)
-                    gen.throw(exc)
+                self.short_circuit = True
+                await self.call_hooks()
 
-    return engine.result()
+        new_state = Running()
+        state = await self.set_state(new_state)
+        while state.is_pending():
+            await asyncio.sleep(0.2)
+            state = await self.set_state(new_state)
+        return state
 
+    async def set_state(self, state: State, force: bool = False) -> State:
+        """ """
+        # prevents any state-setting activity
+        if self.short_circuit:
+            return self.state
 
-async def run_generator_flow_async(
-    flow: Flow[P, R],
-    flow_run: Optional[FlowRun] = None,
-    parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
-    return_type: Literal["state", "result"] = "result",
-) -> AsyncGenerator[R, None]:
-    if return_type != "result":
-        raise ValueError("The return_type for a generator flow must be 'result'")
+        state = await propose_state(
+            self.client, state, flow_run_id=self.flow_run.id, force=force
+        )  # type: ignore
+        self.flow_run.state = state  # type: ignore
+        self.flow_run.state_name = state.name  # type: ignore
+        self.flow_run.state_type = state.type  # type: ignore
 
-    engine = FlowRunEngine[P, R](
-        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
-    )
+        if self._span:
+            self._span.add_event(
+                state.name or state.type,
+                {
+                    "prefect.state.message": state.message or "",
+                    "prefect.state.type": state.type,
+                    "prefect.state.name": state.name or state.type,
+                    "prefect.state.id": str(state.id),
+                },
+            )
+        return state
 
-    with engine.start():
-        while engine.is_running():
-            with engine.run_context():
-                call_args, call_kwargs = parameters_to_args_kwargs(
-                    flow.fn, engine.parameters or {}
-                )
-                gen = flow.fn(*call_args, **call_kwargs)
-                try:
-                    while True:
+    async 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 asyncio.iscoroutine(_result):
+                # getting the value for a BaseResult may return an awaitable
+                # depending on whether the parent frame is sync or not
+                _result = await _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
+        if asyncio.iscoroutine(_result):
+            _result = await _result
+        return _result
+
+    async def handle_success(self, result: R) -> R:
+        result_store = getattr(FlowRunContext.get(), "result_store", None)
+        if result_store is None:
+            raise ValueError("Result store is not set")
+        resolved_result = resolve_futures_to_states(result)
+        terminal_state = await return_value_to_state(
+            resolved_result,
+            result_store=result_store,
+            write_result=should_persist_result(),
+        )
+        await self.set_state(terminal_state)
+        self._return_value = resolved_result
+
+        self._end_span_on_success()
+
+        return result
+
+    async def handle_exception(
+        self,
+        exc: Exception,
+        msg: Optional[str] = None,
+        result_store: Optional[ResultStore] = None,
+    ) -> State:
+        context = FlowRunContext.get()
+        terminal_state = cast(
+            State,
+            await exception_to_failed_state(
+                exc,
+                message=msg or "Flow run encountered an exception:",
+                result_store=result_store or getattr(context, "result_store", None),
+                write_result=True,
+            ),
+        )
+        state = await self.set_state(terminal_state)
+        if self.state.is_scheduled():
+            self.logger.info(
+                (
+                    f"Received non-final state {state.name!r} when proposing final"
+                    f" state {terminal_state.name!r} and will attempt to run again..."
+                ),
+            )
+            state = await self.set_state(Running())
+        self._raised = exc
+
+        self._end_span_on_error(exc, state.message)
+
+        return state
+
+    async def handle_timeout(self, exc: TimeoutError) -> None:
+        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,
+            message=message,
+            name="TimedOut",
+        )
+        await self.set_state(state)
+        self._raised = exc
+
+        self._end_span_on_error(exc, message)
+
+    async def handle_crash(self, exc: BaseException) -> None:
+        # need to shield from asyncio cancellation to ensure we update the state
+        # on the server before exiting
+        with CancelScope(shield=True):
+            state = await exception_to_crashed_state(exc)
+            self.logger.error(f"Crash detected! {state.message}")
+            self.logger.debug("Crash details:", exc_info=exc)
+            await self.set_state(state, force=True)
+            self._raised = exc
+
+            self._end_span_on_error(exc, state.message)
+
+    async def load_subflow_run(
+        self,
+        parent_task_run: TaskRun,
+        client: PrefectClient,
+        context: FlowRunContext,
+    ) -> Union[FlowRun, None]:
+        """
+        This method attempts to load an existing flow run for a subflow task
+        run, if appropriate.
+
+        If the parent task run is in a final but not COMPLETED state, and not
+        being rerun, then we attempt to load an existing flow run instead of
+        creating a new one. This will prevent the engine from running the
+        subflow again.
+
+        If no existing flow run is found, or if the subflow should be rerun,
+        then no flow run is returned.
+        """
+
+        # check if the parent flow run is rerunning
+        rerunning = (
+            context.flow_run.run_count > 1
+            if getattr(context, "flow_run", None)
+            and isinstance(context.flow_run, FlowRun)
+            else False
+        )
+
+        # if the parent task run is in a final but not completed state, and
+        # not rerunning, then retrieve the most recent flow run instead of
+        # creating a new one. This effectively loads a cached flow run for
+        # situations where we are confident the flow should not be run
+        # again.
+        assert isinstance(parent_task_run.state, State)
+        if parent_task_run.state.is_final() and not (
+            rerunning and not parent_task_run.state.is_completed()
+        ):
+            # return the most recent flow run, if it exists
+            flow_runs = await client.read_flow_runs(
+                flow_run_filter=FlowRunFilter(
+                    parent_task_run_id={"any_": [parent_task_run.id]}
+                ),
+                sort=FlowRunSort.EXPECTED_START_TIME_ASC,
+                limit=1,
+            )
+            if flow_runs:
+                loaded_flow_run = flow_runs[-1]
+                self._return_value = loaded_flow_run.state
+                return loaded_flow_run
+
+    async def create_flow_run(self, client: PrefectClient) -> FlowRun:
+        flow_run_ctx = FlowRunContext.get()
+        parameters = self.parameters or {}
+
+        parent_task_run = None
+
+        # this is a subflow run
+        if flow_run_ctx:
+            # add a task to a parent flow run that represents the execution of a subflow run
+            parent_task = Task(
+                name=self.flow.name, fn=self.flow.fn, version=self.flow.version
+            )
+
+            parent_task_run = await parent_task.create_run(
+                flow_run_context=flow_run_ctx,
+                parameters=self.parameters,
+                wait_for=self.wait_for,
+            )
+
+            # check if there is already a flow run for this subflow
+            if subflow_run := await self.load_subflow_run(
+                parent_task_run=parent_task_run, client=client, context=flow_run_ctx
+            ):
+                return subflow_run
+
+        flow_run = await client.create_flow_run(
+            flow=self.flow,
+            parameters=self.flow.serialize_parameters(parameters),
+            state=Pending(),
+            parent_task_run_id=getattr(parent_task_run, "id", None),
+            tags=TagsContext.get().current_tags,
+        )
+        if flow_run_ctx:
+            parent_logger = get_run_logger(flow_run_ctx)
+            parent_logger.info(
+                f"Created subflow run {flow_run.name!r} for flow {self.flow.name!r}"
+            )
+        else:
+            self.logger.info(
+                f"Created flow run {flow_run.name!r} for flow {self.flow.name!r}"
+            )
+
+        return flow_run
+
+    async def call_hooks(self, state: Optional[State] = None):
+        if state is None:
+            state = self.state
+        flow = self.flow
+        flow_run = self.flow_run
+
+        if not flow_run:
+            raise ValueError("Flow run is not set")
+
+        enable_cancellation_and_crashed_hooks = (
+            os.environ.get(
+                "PREFECT__ENABLE_CANCELLATION_AND_CRASHED_HOOKS", "true"
+            ).lower()
+            == "true"
+        )
+
+        if state.is_failed() and flow.on_failure_hooks:
+            hooks = flow.on_failure_hooks
+        elif state.is_completed() and flow.on_completion_hooks:
+            hooks = flow.on_completion_hooks
+        elif (
+            enable_cancellation_and_crashed_hooks
+            and state.is_cancelling()
+            and flow.on_cancellation_hooks
+        ):
+            hooks = flow.on_cancellation_hooks
+        elif (
+            enable_cancellation_and_crashed_hooks
+            and state.is_crashed()
+            and flow.on_crashed_hooks
+        ):
+            hooks = flow.on_crashed_hooks
+        elif state.is_running() and flow.on_running_hooks:
+            hooks = flow.on_running_hooks
+        else:
+            hooks = None
+
+        for hook in hooks or []:
+            hook_name = get_hook_name(hook)
+
+            try:
+                self.logger.info(
+                    f"Running hook {hook_name!r} in response to entering state"
+                    f" {state.name!r}"
+                )
+                result = hook(flow, flow_run, state)
+                if asyncio.iscoroutine(result):
+                    await result
+            except Exception:
+                self.logger.error(
+                    f"An error was encountered while running hook {hook_name!r}",
+                    exc_info=True,
+                )
+            else:
+                self.logger.info(f"Hook {hook_name!r} finished running successfully")
+
+    @asynccontextmanager
+    async def setup_run_context(self, client: Optional[PrefectClient] = None):
+        from prefect.utilities.engine import (
+            should_log_prints,
+        )
+
+        if client is None:
+            client = self.client
+        if not self.flow_run:
+            raise ValueError("Flow run not set")
+
+        self.flow_run = await client.read_flow_run(self.flow_run.id)
+        log_prints = should_log_prints(self.flow)
+
+        with ExitStack() as stack:
+            # TODO: Explore closing task runner before completing the flow to
+            # wait for futures to complete
+            stack.enter_context(capture_sigterm())
+            if log_prints:
+                stack.enter_context(patch_print())
+            task_runner = stack.enter_context(self.flow.task_runner.duplicate())
+            stack.enter_context(
+                FlowRunContext(
+                    flow=self.flow,
+                    log_prints=log_prints,
+                    flow_run=self.flow_run,
+                    parameters=self.parameters,
+                    client=client,
+                    result_store=get_result_store().update_for_flow(
+                        self.flow, _sync=True
+                    ),
+                    task_runner=task_runner,
+                    persist_result=self.flow.persist_result
+                    if self.flow.persist_result is not None
+                    else should_persist_result(),
+                )
+            )
+            stack.enter_context(ConcurrencyContextV1())
+            stack.enter_context(ConcurrencyContext())
+
+            # set the logger to the flow run logger
+            self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
+
+            # update the flow run name if necessary
+            if not self._flow_run_name_set and self.flow.flow_run_name:
+                flow_run_name = resolve_custom_flow_run_name(
+                    flow=self.flow, parameters=self.parameters
+                )
+                await self.client.set_flow_run_name(
+                    flow_run_id=self.flow_run.id, name=flow_run_name
+                )
+                self.logger.extra["flow_run_name"] = flow_run_name
+                self.logger.debug(
+                    f"Renamed flow run {self.flow_run.name!r} to {flow_run_name!r}"
+                )
+                self.flow_run.name = flow_run_name
+                self._flow_run_name_set = True
+            yield
+
+    @asynccontextmanager
+    async def initialize_run(self):
+        """
+        Enters a client context and creates a flow run if needed.
+        """
+        async with AsyncClientContext.get_or_create() as client_ctx:
+            self._client = client_ctx.client
+            self._is_started = True
+
+            if not self.flow_run:
+                self.flow_run = await self.create_flow_run(self.client)
+                flow_run_url = url_for(self.flow_run)
+
+                if flow_run_url:
+                    self.logger.info(
+                        f"View at {flow_run_url}", extra={"send_to_api": False}
+                    )
+            else:
+                # Update the empirical policy to match the flow if it is not set
+                if self.flow_run.empirical_policy.retry_delay is None:
+                    self.flow_run.empirical_policy.retry_delay = (
+                        self.flow.retry_delay_seconds
+                    )
+
+                if self.flow_run.empirical_policy.retries is None:
+                    self.flow_run.empirical_policy.retries = self.flow.retries
+
+                await self.client.update_flow_run(
+                    flow_run_id=self.flow_run.id,
+                    flow_version=self.flow.version,
+                    empirical_policy=self.flow_run.empirical_policy,
+                )
+
+            span = self._tracer.start_span(
+                name=self.flow_run.name,
+                attributes={
+                    **self.flow_run.labels,
+                    "prefect.run.type": "flow",
+                    "prefect.run.id": str(self.flow_run.id),
+                    "prefect.tags": self.flow_run.tags,
+                    "prefect.flow.name": self.flow.name,
+                },
+            )
+            self._update_otel_labels(span, self.client)
+            self._span = span
+
+            try:
+                yield self
+
+            except TerminationSignal as exc:
+                self.cancel_all_tasks()
+                await self.handle_crash(exc)
+                raise
+            except Exception:
+                # regular exceptions are caught and re-raised to the user
+                raise
+            except (Abort, Pause):
+                raise
+            except GeneratorExit:
+                # Do not capture generator exits as crashes
+                raise
+            except BaseException as exc:
+                # BaseExceptions are caught and handled as crashes
+                await self.handle_crash(exc)
+                raise
+            finally:
+                # If debugging, use the more complete `repr` than the usual `str` description
+                display_state = (
+                    repr(self.state) if PREFECT_DEBUG_MODE else str(self.state)
+                )
+                self.logger.log(
+                    level=logging.INFO if self.state.is_completed() else logging.ERROR,
+                    msg=f"Finished in state {display_state}",
+                )
+
+                self._is_started = False
+                self._client = None
+
+    # --------------------------
+    #
+    # The following methods compose the main task run loop
+    #
+    # --------------------------
+
+    @asynccontextmanager
+    async def start(self) -> AsyncGenerator[None, None]:
+        async with self.initialize_run():
+            with trace.use_span(self._span) if self._span else nullcontext():
+                await self.begin_run()
+
+                if self.state.is_running():
+                    await self.call_hooks()
+                yield
+
+    @asynccontextmanager
+    async def run_context(self):
+        timeout_context = timeout_async if self.flow.isasync else timeout
+        # reenter the run context to ensure it is up to date for every run
+        async with self.setup_run_context():
+            try:
+                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}..."
+                    )
+                    yield self
+            except TimeoutError as exc:
+                await self.handle_timeout(exc)
+            except Exception as exc:
+                self.logger.exception("Encountered exception during execution: %r", exc)
+                await self.handle_exception(exc)
+            finally:
+                if self.state.is_final() or self.state.is_cancelling():
+                    await self.call_hooks()
+
+    async def call_flow_fn(self) -> Coroutine[Any, Any, R]:
+        """
+        Convenience method to call the flow function. Returns a coroutine if the
+        flow is async.
+        """
+        assert self.flow.isasync, "Flow must be async to be run with AsyncFlowRunEngine"
+
+        result = await call_with_parameters(self.flow.fn, self.parameters)
+        await self.handle_success(result)
+        return result
+
+
+def run_flow_sync(
+    flow: Flow[P, R],
+    flow_run: Optional[FlowRun] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    return_type: Literal["state", "result"] = "result",
+) -> Union[R, State, None]:
+    engine = FlowRunEngine[P, R](
+        flow=flow,
+        parameters=parameters,
+        flow_run=flow_run,
+        wait_for=wait_for,
+    )
+
+    with engine.start():
+        while engine.is_running():
+            with engine.run_context():
+                engine.call_flow_fn()
+
+    return engine.state if return_type == "state" else engine.result()
+
+
+async def run_flow_async(
+    flow: Flow[P, R],
+    flow_run: Optional[FlowRun] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    return_type: Literal["state", "result"] = "result",
+) -> Union[R, State, None]:
+    engine = AsyncFlowRunEngine[P, R](
+        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
+    )
+
+    async with engine.start():
+        while engine.is_running():
+            async with engine.run_context():
+                await engine.call_flow_fn()
+
+    return engine.state if return_type == "state" else await engine.result()
+
+
+def run_generator_flow_sync(
+    flow: Flow[P, R],
+    flow_run: Optional[FlowRun] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    return_type: Literal["state", "result"] = "result",
+) -> Generator[R, None, None]:
+    if return_type != "result":
+        raise ValueError("The return_type for a generator flow must be 'result'")
+
+    engine = FlowRunEngine[P, R](
+        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
+    )
+
+    with engine.start():
+        while engine.is_running():
+            with engine.run_context():
+                call_args, call_kwargs = parameters_to_args_kwargs(
+                    flow.fn, engine.parameters or {}
+                )
+                gen = flow.fn(*call_args, **call_kwargs)
+                try:
+                    while True:
+                        gen_result = next(gen)
+                        # link the current state to the result for dependency tracking
+                        link_state_to_result(engine.state, gen_result)
+                        yield gen_result
+                except StopIteration as exc:
+                    engine.handle_success(exc.value)
+                except GeneratorExit as exc:
+                    engine.handle_success(None)
+                    gen.throw(exc)
+
+    return engine.result()
+
+
+async def run_generator_flow_async(
+    flow: Flow[P, R],
+    flow_run: Optional[FlowRun] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
+    return_type: Literal["state", "result"] = "result",
+) -> AsyncGenerator[R, None]:
+    if return_type != "result":
+        raise ValueError("The return_type for a generator flow must be 'result'")
+
+    engine = AsyncFlowRunEngine[P, R](
+        flow=flow, parameters=parameters, flow_run=flow_run, wait_for=wait_for
+    )
+
+    async with engine.start():
+        while engine.is_running():
+            async with engine.run_context():
+                call_args, call_kwargs = parameters_to_args_kwargs(
+                    flow.fn, engine.parameters or {}
+                )
+                gen = flow.fn(*call_args, **call_kwargs)
+                try:
+                    while True:
                         # can't use anext in Python < 3.10
                         gen_result = await gen.__anext__()
                         # link the current state to the result for dependency tracking
                         link_state_to_result(engine.state, gen_result)
                         yield gen_result
                 except (StopAsyncIteration, GeneratorExit) as exc:
-                    engine.handle_success(None)
+                    await engine.handle_success(None)
                     if isinstance(exc, GeneratorExit):
                         gen.throw(exc)
 
     # async generators can't return, but we can raise failures here
     if engine.state.is_failed():
-        engine.result()
+        await engine.result()
 
 
 def run_flow(
     flow: Flow[P, R],
     flow_run: Optional[FlowRun] = None,
     parameters: Optional[Dict[str, Any]] = None,
-    wait_for: Optional[Iterable[PrefectFuture]] = None,
+    wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     return_type: Literal["state", "result"] = "result",
 ) -> Union[R, State, None]:
     kwargs = dict(