prefect.flow_engine

FlowRunEngine dataclass

Bases: Generic[P, R]

Source code in src/prefect/flow_engine.py

@dataclass
class FlowRunEngine(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
    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
    _flow_run_name_set: bool = False

    def __post_init__(self):
        if self.flow is None and self.flow_run_id is None:
            raise ValueError("Either a flow or a flow_run_id must be provided.")

        if self.parameters is None:
            self.parameters = {}

    @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 {}

        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

        self.parameters = resolved_parameters

    def _wait_for_dependencies(self):
        if not self.wait_for:
            return

        visit_collection(
            self.wait_for,
            visit_fn=resolve_to_final_result,
            return_data=False,
            max_depth=-1,
            remove_annotations=True,
            context={},
        )

    def begin_run(self) -> State:
        try:
            self._resolve_parameters()
            self._wait_for_dependencies()
        except UpstreamTaskError as upstream_exc:
            state = 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

        # 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)
                self.handle_exception(
                    exc,
                    msg=message,
                    result_store=get_result_store().update_for_flow(
                        self.flow, _sync=True
                    ),
                )
                self.short_circuit = True
                self.call_hooks()

        new_state = Running()
        state = self.set_state(new_state)
        while state.is_pending():
            time.sleep(0.2)
            state = self.set_state(new_state)
        return state

    def set_state(self, state: State, force: bool = False) -> State:
        """ """
        # prevents any state-setting activity
        if self.short_circuit:
            return self.state

        state = propose_state_sync(
            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
        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 asyncio.iscoroutine(_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
        if asyncio.iscoroutine(_result):
            _result = run_coro_as_sync(_result)
        return _result

    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 = run_coro_as_sync(
            return_value_to_state(
                resolved_result,
                result_store=result_store,
                write_result=should_persist_result(),
            )
        )
        self.set_state(terminal_state)
        self._return_value = resolved_result
        return result

    def handle_exception(
        self,
        exc: Exception,
        msg: Optional[str] = None,
        result_store: Optional[ResultStore] = None,
    ) -> State:
        context = FlowRunContext.get()
        terminal_state = cast(
            State,
            run_coro_as_sync(
                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 = 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 = self.set_state(Running())
        self._raised = exc
        return state

    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",
        )
        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,
        parent_task_run: TaskRun,
        client: SyncPrefectClient,
        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 = 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

    def create_flow_run(self, client: SyncPrefectClient) -> 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 = run_coro_as_sync(
                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 := self.load_subflow_run(
                parent_task_run=parent_task_run, client=client, context=flow_run_ctx
            ):
                return subflow_run

        flow_run = 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

    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):
                    run_coro_as_sync(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")

    @contextmanager
    def setup_run_context(self, client: Optional[SyncPrefectClient] = 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 = 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
                )
                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

    @contextmanager
    def initialize_run(self):
        """
        Enters a client context and creates a flow run if needed.
        """
        with SyncClientContext.get_or_create() as client_ctx:
            self._client = client_ctx.client
            self._is_started = True

            if not self.flow_run:
                self.flow_run = 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

                self.client.update_flow_run(
                    flow_run_id=self.flow_run.id,
                    flow_version=self.flow.version,
                    empirical_policy=self.flow_run.empirical_policy,
                )
            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
            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
                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

    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
    #
    # --------------------------

    @contextmanager
    def start(self) -> Generator[None, None, None]:
        with self.initialize_run():
            self.begin_run()

            if self.state.is_running():
                self.call_hooks()
            yield

    @contextmanager
    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
        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:
                self.handle_timeout(exc)
            except Exception as 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]]:
        """
        Convenience method to call the flow function. Returns a coroutine if the
        flow is async.
        """
        if self.flow.isasync:

            async def _call_flow_fn():
                result = await call_with_parameters(self.flow.fn, self.parameters)
                self.handle_success(result)

            return _call_flow_fn()
        else:
            result = call_with_parameters(self.flow.fn, self.parameters)
            self.handle_success(result)

call_flow_fn()

Convenience method to call the flow function. Returns a coroutine if the flow is async.

Source code in src/prefect/flow_engine.py

def call_flow_fn(self) -> Union[R, Coroutine[Any, Any, R]]:
    """
    Convenience method to call the flow function. Returns a coroutine if the
    flow is async.
    """
    if self.flow.isasync:

        async def _call_flow_fn():
            result = await call_with_parameters(self.flow.fn, self.parameters)
            self.handle_success(result)

        return _call_flow_fn()
    else:
        result = call_with_parameters(self.flow.fn, self.parameters)
        self.handle_success(result)

initialize_run()

Enters a client context and creates a flow run if needed.

Source code in src/prefect/flow_engine.py

@contextmanager
def initialize_run(self):
    """
    Enters a client context and creates a flow run if needed.
    """
    with SyncClientContext.get_or_create() as client_ctx:
        self._client = client_ctx.client
        self._is_started = True

        if not self.flow_run:
            self.flow_run = 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

            self.client.update_flow_run(
                flow_run_id=self.flow_run.id,
                flow_version=self.flow.version,
                empirical_policy=self.flow_run.empirical_policy,
            )
        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
        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
            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

load_subflow_run(parent_task_run, client, context)

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.

Source code in src/prefect/flow_engine.py

def load_subflow_run(
    self,
    parent_task_run: TaskRun,
    client: SyncPrefectClient,
    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 = 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

set_state(state, force=False)

Source code in src/prefect/flow_engine.py

def set_state(self, state: State, force: bool = False) -> State:
    """ """
    # prevents any state-setting activity
    if self.short_circuit:
        return self.state

    state = propose_state_sync(
        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
    return state

FlowRunTimeoutError

Bases: TimeoutError

Raised when a flow run exceeds its defined timeout.

Source code in src/prefect/flow_engine.py

class FlowRunTimeoutError(TimeoutError):
    """Raised when a flow run exceeds its defined timeout."""