Repository: PrefectHQ/prefect

Issue:
"Flow" is not exported from module "prefect.flows"
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar issue and didn't find it.
- [X] I searched the Prefect documentation for this issue.
- [X] I checked that this issue is related to Prefect and not one of its dependencies.

### Bug summary

pyright/pylance errors with:
```
error: "Flow" is not exported from module "prefect.flows" (reportPrivateImportUsage)
```

### Reproduction

```python3
from prefect.flows import Flow
```


### Error

_No response_

### Versions

```Text
2.8.3
```


### Additional context

This error was [introduced here](https://github.com/PrefectHQ/prefect/pull/8575/files#diff-835f620a789907a6183853543ab11ba098f897f57876cd97f75d19b6b41ce3cdR45) by imports that shadow the Flow class.


Discussion:
@tekumara thanks for the report! We should definitely not be importing those into the namespace like that. We should just import the `schemas` module if we need to have type hints for that, `FlowRun` should even come from `prefect.client.schemas` as well.

Resolving patch:
diff --git a/src/prefect/engine.py b/src/prefect/engine.py
--- a/src/prefect/engine.py
+++ b/src/prefect/engine.py
@@ -357,7 +357,7 @@ async def begin_flow_run(
 
     async with AsyncExitStack() as stack:
         await stack.enter_async_context(
-            report_flow_run_crashes(flow_run=flow_run, client=client)
+            report_flow_run_crashes(flow_run=flow_run, client=client, flow=flow)
         )
 
         # Create a task group for background tasks
@@ -532,7 +532,7 @@ async def create_and_begin_subflow_run(
         if terminal_state is None or not terminal_state.is_final():
             async with AsyncExitStack() as stack:
                 await stack.enter_async_context(
-                    report_flow_run_crashes(flow_run=flow_run, client=client)
+                    report_flow_run_crashes(flow_run=flow_run, client=client, flow=flow)
                 )
 
                 task_runner = flow.task_runner.duplicate()
@@ -1709,7 +1709,7 @@ async def wait_for_task_runs_and_report_crashes(
 
 
 @asynccontextmanager
-async def report_flow_run_crashes(flow_run: FlowRun, client: PrefectClient):
+async def report_flow_run_crashes(flow_run: FlowRun, client: PrefectClient, flow: Flow):
     """
     Detect flow run crashes during this context and update the run to a proper final
     state.
@@ -1746,6 +1746,14 @@ def cancel_flow_run(*args):
                 f"Reported crashed flow run {flow_run.name!r} successfully!"
             )
 
+            # Only `on_crashed` flow run state change hook is called here
+            # We call the hook after the state is set to `CRASHED`
+            await _run_flow_hooks(
+                flow=flow,
+                flow_run=flow_run,
+                state=state,
+            )
+
         if isinstance(exc, TerminationSignal):
             # Termination signals are swapped out during a flow run to perform
             # a graceful shutdown and raise this exception. This `os.kill` call
@@ -2158,7 +2166,7 @@ async def _run_task_hooks(task: Task, task_run: TaskRun, state: State) -> None:
 
 
 async def _run_flow_hooks(flow: Flow, flow_run: FlowRun, state: State) -> None:
-    """Run the on_failure and on_completion hooks for a flow, making sure to
+    """Run the on_failure, on_completion, and on_crashed hooks for a flow, making sure to
     catch and log any errors that occur.
     """
     hooks = None
@@ -2166,6 +2174,8 @@ async def _run_flow_hooks(flow: Flow, flow_run: FlowRun, state: State) -> None:
         hooks = flow.on_failure
     elif state.is_completed() and flow.on_completion:
         hooks = flow.on_completion
+    elif state.is_crashed() and flow.on_crashed:
+        hooks = flow.on_crashed
 
     if hooks:
         logger = flow_run_logger(flow_run)
diff --git a/src/prefect/flows.py b/src/prefect/flows.py
--- a/src/prefect/flows.py
+++ b/src/prefect/flows.py
@@ -42,7 +42,8 @@
 from prefect.futures import PrefectFuture
 from prefect.logging import get_logger
 from prefect.results import ResultSerializer, ResultStorage
-from prefect.server.schemas.core import Flow, FlowRun, raise_on_invalid_name
+import prefect.server.schemas as schemas
+from prefect.client.schemas import FlowRun
 from prefect.states import State
 from prefect.task_runners import BaseTaskRunner, ConcurrentTaskRunner
 from prefect.utilities.annotations import NotSet
@@ -117,6 +118,7 @@ class Flow(Generic[P, R]):
             loaded from the parent flow.
         on_failure: An optional list of callables to run when the flow enters a failed state.
         on_completion: An optional list of callables to run when the flow enters a completed state.
+        on_crashed: An optional list of callables to run when the flow enters a crashed state.
     """
 
     # NOTE: These parameters (types, defaults, and docstrings) should be duplicated
@@ -138,15 +140,22 @@ def __init__(
         result_serializer: Optional[ResultSerializer] = None,
         cache_result_in_memory: bool = True,
         log_prints: Optional[bool] = None,
-        on_completion: Optional[List[Callable[[Flow, FlowRun, State], None]]] = None,
-        on_failure: Optional[List[Callable[[Flow, FlowRun, State], None]]] = None,
+        on_completion: Optional[
+            List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+        ] = None,
+        on_failure: Optional[
+            List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+        ] = None,
+        on_crashed: Optional[
+            List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+        ] = None,
     ):
         if not callable(fn):
             raise TypeError("'fn' must be callable")
 
         # Validate name if given
         if name:
-            raise_on_invalid_name(name)
+            schemas.core.raise_on_invalid_name(name)
 
         self.name = name or fn.__name__.replace("_", "-")
 
@@ -228,6 +237,7 @@ def __init__(
             )
         self.on_completion = on_completion
         self.on_failure = on_failure
+        self.on_crashed = on_crashed
 
     def with_options(
         self,
@@ -246,8 +256,15 @@ def with_options(
         result_serializer: Optional[ResultSerializer] = NotSet,
         cache_result_in_memory: bool = None,
         log_prints: Optional[bool] = NotSet,
-        on_completion: Optional[List[Callable[[Flow, FlowRun, State], None]]] = None,
-        on_failure: Optional[List[Callable[[Flow, FlowRun, State], None]]] = None,
+        on_completion: Optional[
+            List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+        ] = None,
+        on_failure: Optional[
+            List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+        ] = None,
+        on_crashed: Optional[
+            List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+        ] = None,
     ):
         """
         Create a new flow from the current object, updating provided options.
@@ -274,6 +291,7 @@ def with_options(
                 be cached in memory.
             on_failure: A new list of callables to run when the flow enters a failed state.
             on_completion: A new list of callables to run when the flow enters a completed state.
+            on_crashed: A new list of callables to run when the flow enters a crashed state.
 
         Returns:
             A new `Flow` instance.
@@ -337,6 +355,7 @@ def with_options(
             log_prints=log_prints if log_prints is not NotSet else self.log_prints,
             on_completion=on_completion or self.on_completion,
             on_failure=on_failure or self.on_failure,
+            on_crashed=on_crashed or self.on_crashed,
         )
 
     def validate_parameters(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
@@ -550,8 +569,15 @@ def flow(
     result_serializer: Optional[ResultSerializer] = None,
     cache_result_in_memory: bool = True,
     log_prints: Optional[bool] = None,
-    on_completion: Optional[List[Callable[[Flow, FlowRun, State], None]]] = None,
-    on_failure: Optional[List[Callable[[Flow, FlowRun, State], None]]] = None,
+    on_completion: Optional[
+        List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+    ] = None,
+    on_failure: Optional[
+        List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+    ] = None,
+    on_crashed: Optional[
+        List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+    ] = None,
 ) -> Callable[[Callable[P, R]], Flow[P, R]]:
     ...
 
@@ -573,8 +599,15 @@ def flow(
     result_serializer: Optional[ResultSerializer] = None,
     cache_result_in_memory: bool = True,
     log_prints: Optional[bool] = None,
-    on_completion: Optional[List[Callable[[Flow, FlowRun, State], None]]] = None,
-    on_failure: Optional[List[Callable[[Flow, FlowRun, State], None]]] = None,
+    on_completion: Optional[
+        List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+    ] = None,
+    on_failure: Optional[
+        List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+    ] = None,
+    on_crashed: Optional[
+        List[Callable[[schemas.core.Flow, FlowRun, State], None]]
+    ] = None,
 ):
     """
     Decorator to designate a function as a Prefect workflow.
@@ -625,6 +658,15 @@ def flow(
             Prefect logger for the flow run. Defaults to `None`, which indicates that
             the value from the parent flow should be used. If this is a parent flow,
             the default is pulled from the `PREFECT_LOGGING_LOG_PRINTS` setting.
+        on_completion: An optional list of functions to call when the flow run is
+            completed. Each function should accept three arguments: the flow, the flow
+            run, and the final state of the flow run.
+        on_failure: An optional list of functions to call when the flow run fails. Each
+            function should accept three arguments: the flow, the flow run, and the
+            final state of the flow run.
+        on_crashed: An optional list of functions to call when the flow run crashes. Each
+            function should accept three arguments: the flow, the flow run, and the
+            final state of the flow run.
 
     Returns:
         A callable `Flow` object which, when called, will run the flow and return its
@@ -685,6 +727,7 @@ def flow(
                 log_prints=log_prints,
                 on_completion=on_completion,
                 on_failure=on_failure,
+                on_crashed=on_crashed,
             ),
         )
     else:
@@ -708,6 +751,7 @@ def flow(
                 log_prints=log_prints,
                 on_completion=on_completion,
                 on_failure=on_failure,
+                on_crashed=on_crashed,
             ),
         )
 


Test patch:
diff --git a/tests/test_engine.py b/tests/test_engine.py
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -1262,13 +1262,15 @@ async def my_flow():
         assert flow_run.state.type != StateType.CRASHED
 
     async def test_report_flow_run_crashes_handles_sigterm(
-        self, flow_run, orion_client, monkeypatch
+        self, flow_run, orion_client, monkeypatch, parameterized_flow
     ):
         original_handler = Mock()
         signal.signal(signal.SIGTERM, original_handler)
 
         with pytest.raises(TerminationSignal):
-            async with report_flow_run_crashes(flow_run=flow_run, client=orion_client):
+            async with report_flow_run_crashes(
+                flow_run=flow_run, client=orion_client, flow=parameterized_flow
+            ):
                 assert signal.getsignal(signal.SIGTERM) != original_handler
                 os.kill(os.getpid(), signal.SIGTERM)
 
diff --git a/tests/test_flows.py b/tests/test_flows.py
--- a/tests/test_flows.py
+++ b/tests/test_flows.py
@@ -223,6 +223,7 @@ def test_with_options_allows_override_of_flow_settings(self):
             cache_result_in_memory=False,
             on_completion=[],
             on_failure=[],
+            on_crashed=[],
         )
         def initial_flow():
             pass
@@ -233,6 +234,9 @@ def failure_hook(task, task_run, state):
         def success_hook(task, task_run, state):
             return print("Meow!")
 
+        def crash_hook(task, task_run, state):
+            return print("Crash!")
+
         flow_with_options = initial_flow.with_options(
             name="Copied flow",
             description="A copied flow",
@@ -248,6 +252,7 @@ def success_hook(task, task_run, state):
             cache_result_in_memory=True,
             on_completion=[success_hook],
             on_failure=[failure_hook],
+            on_crashed=[crash_hook],
         )
 
         assert flow_with_options.name == "Copied flow"
@@ -264,6 +269,7 @@ def success_hook(task, task_run, state):
         assert flow_with_options.cache_result_in_memory is True
         assert flow_with_options.on_completion == [success_hook]
         assert flow_with_options.on_failure == [failure_hook]
+        assert flow_with_options.on_crashed == [crash_hook]
 
     def test_with_options_uses_existing_settings_when_no_override(self):
         @flow(
@@ -2392,3 +2398,115 @@ def my_flow():
         state = my_flow._run()
         assert state.type == StateType.FAILED
         assert my_mock.call_args_list == [call(), call()]
+
+
+class TestFlowHooksOnCrashed:
+    def test_on_crashed_hooks_run_on_crashed_state(self):
+        my_mock = MagicMock()
+
+        def crashed_hook1(flow, flow_run, state):
+            my_mock("crashed_hook1")
+
+        def crashed_hook2(flow, flow_run, state):
+            my_mock("crashed_hook2")
+
+        @flow(on_crashed=[crashed_hook1, crashed_hook2])
+        def my_flow():
+            return State(type=StateType.CRASHED)
+
+        my_flow._run()
+        assert my_mock.mock_calls == [call("crashed_hook1"), call("crashed_hook2")]
+
+    def test_on_crashed_hooks_are_ignored_if_terminal_state_completed(self):
+        my_mock = MagicMock()
+
+        def crashed_hook1(flow, flow_run, state):
+            my_mock("crashed_hook1")
+
+        def crashed_hook2(flow, flow_run, state):
+            my_mock("crashed_hook2")
+
+        @flow(on_crashed=[crashed_hook1, crashed_hook2])
+        def my_passing_flow():
+            pass
+
+        state = my_passing_flow._run()
+        assert state.type == StateType.COMPLETED
+        assert my_mock.mock_calls == []
+
+    def test_on_crashed_hooks_are_ignored_if_terminal_state_failed(self):
+        my_mock = MagicMock()
+
+        def crashed_hook1(flow, flow_run, state):
+            my_mock("crashed_hook1")
+
+        def crashed_hook2(flow, flow_run, state):
+            my_mock("crashed_hook2")
+
+        @flow(on_crashed=[crashed_hook1, crashed_hook2])
+        def my_failing_flow():
+            raise Exception("Failing flow")
+
+        state = my_failing_flow._run()
+        assert state.type == StateType.FAILED
+        assert my_mock.mock_calls == []
+
+    def test_other_crashed_hooks_run_if_one_hook_fails(self):
+        my_mock = MagicMock()
+
+        def crashed1(flow, flow_run, state):
+            my_mock("crashed1")
+
+        def crashed2(flow, flow_run, state):
+            raise Exception("Failing flow")
+
+        def crashed3(flow, flow_run, state):
+            my_mock("crashed3")
+
+        @flow(on_crashed=[crashed1, crashed2, crashed3])
+        def my_flow():
+            return State(type=StateType.CRASHED)
+
+        my_flow._run()
+        assert my_mock.mock_calls == [call("crashed1"), call("crashed3")]
+
+    @pytest.mark.parametrize(
+        "hook1, hook2",
+        [
+            (create_hook, create_hook),
+            (create_hook, create_async_hook),
+            (create_async_hook, create_hook),
+            (create_async_hook, create_async_hook),
+        ],
+    )
+    def test_on_crashed_hooks_work_with_sync_and_async(self, hook1, hook2):
+        my_mock = MagicMock()
+        hook1_with_mock = hook1(my_mock)
+        hook2_with_mock = hook2(my_mock)
+
+        @flow(on_crashed=[hook1_with_mock, hook2_with_mock])
+        def my_flow():
+            return State(type=StateType.CRASHED)
+
+        my_flow._run()
+        assert my_mock.mock_calls == [call(), call()]
+
+    def test_on_crashed_hook_on_subflow_succeeds(self):
+        my_mock = MagicMock()
+
+        def crashed1(flow, flow_run, state):
+            my_mock("crashed1")
+
+        def failed1(flow, flow_run, state):
+            my_mock("failed1")
+
+        @flow(on_crashed=[crashed1])
+        def subflow():
+            return State(type=StateType.CRASHED)
+
+        @flow(on_failure=[failed1])
+        def my_flow():
+            subflow()
+
+        my_flow._run()
+        assert my_mock.mock_calls == [call("crashed1"), call("failed1")]