Repository: PrefectHQ/prefect

Issue:
RFC: remove `return_failed` kwarg from FlowRunner.run()?
The more that we run Flows in distributed contexts, and use helpers like `raise_on_exception`, the less useful I think the `return_failed` kwarg is for `FlowRunner.run()`. In particular, it requires some messy code because it can only be satisfied by gathering every single task state to the FlowRunner process, then filtering. If people are finding it useful, then I retract this, but IMO it is not as useful as it used to be.


Discussion:
I only use `raise_on_exception` personally, so I'm OK with this
I think an improvement would be to create a debug state_handler that added failed tasks to a global dictionary. 🤔 

Resolving patch:
diff --git a/conftest.py b/conftest.py
--- a/conftest.py
+++ b/conftest.py
@@ -5,7 +5,7 @@
 
 import prefect
 from prefect.engine.executors import LocalExecutor, SynchronousExecutor
-from prefect.utilities import tests
+from prefect.utilities import debug
 
 if sys.version_info >= (3, 5):
     from prefect.engine.executors import DaskExecutor
diff --git a/src/prefect/engine/flow_runner.py b/src/prefect/engine/flow_runner.py
--- a/src/prefect/engine/flow_runner.py
+++ b/src/prefect/engine/flow_runner.py
@@ -134,8 +134,8 @@ def run(
         start_tasks: Iterable[Task] = None,
         start_task_ids: Iterable[str] = None,
         return_tasks: Iterable[Task] = None,
-        return_failed: bool = False,
         parameters: Dict[str, Any] = None,
+        task_runner_state_handlers: Iterable[Callable] = None,
         executor: "prefect.engine.executors.Executor" = None,
         context: Dict[str, Any] = None,
     ) -> State:
@@ -155,11 +155,12 @@ def run(
                 a list of task IDs. The two options may be used simultaneously.
             - return_tasks ([Task], optional): list of Tasks to include in the
                 final returned Flow state. Defaults to `None`
-            - return_failed (bool, optional): whether to return all tasks
-                which fail, regardless of whether they are terminal tasks or in `return_tasks`.
-                Defaults to `False`
             - parameters (dict, optional): dictionary of any needed Parameter
-                values, with keys being strings representing Parameter names and values being their corresponding values
+                values, with keys being strings representing Parameter names and values being
+                their corresponding values
+            - task_runner_state_handlers (Iterable[Callable], optional): A list of state change
+                handlers that will be provided to the task_runner, and called whenever a task changes
+                state.
             - executor (Executor, optional): executor to use when performing
                 computation; defaults to the executor specified in your prefect configuration
             - context (dict, optional): prefect.Context to use for execution
@@ -173,14 +174,10 @@ def run(
         self.logger.info("Beginning Flow run for '{}'".format(self.flow.name))
 
         context = context or {}
-        return_tasks = set(return_tasks or [])
         parameters = parameters or {}
         if executor is None:
             executor = prefect.engine.get_default_executor_class()()
 
-        if return_tasks.difference(self.flow.tasks):
-            raise ValueError("Some tasks in return_tasks were not found in the flow.")
-
         try:
             state, context = self.initialize_run(state, context, parameters)
 
@@ -196,7 +193,7 @@ def run(
                     start_tasks=start_tasks,
                     start_task_ids=start_task_ids,
                     return_tasks=return_tasks,
-                    return_failed=return_failed,
+                    task_runner_state_handlers=task_runner_state_handlers,
                     executor=executor,
                 )
 
@@ -274,8 +271,8 @@ def get_flow_run_state(
         start_tasks: Iterable[Task],
         start_task_ids: Iterable[str],
         return_tasks: Set[Task],
+        task_runner_state_handlers: Iterable[Callable],
         executor: "prefect.engine.executors.base.Executor",
-        return_failed: bool = False,
     ) -> State:
         """
         Runs the flow.
@@ -292,11 +289,11 @@ def get_flow_run_state(
                 a list of task IDs. The two options may be used simultaneously.
             - return_tasks ([Task], optional): list of Tasks to include in the
                 final returned Flow state. Defaults to `None`
+            - task_runner_state_handlers (Iterable[Callable], optional): A list of state change
+                handlers that will be provided to the task_runner, and called whenever a task changes
+                state.
             - executor (Executor, optional): executor to use when performing
                 computation; defaults to the executor provided in your prefect configuration
-            - return_failed (bool, optional): whether to return all tasks
-                which fail, regardless of whether they are terminal tasks or in `return_tasks`.
-                Defaults to `False`
 
         Returns:
             - State: `State` representing the final post-run state of the `Flow`.
@@ -307,14 +304,19 @@ def get_flow_run_state(
             self.logger.info("Flow is not in a Running state.")
             raise ENDRUN(state)
 
-        # make a copy to avoid modifying the user-supplied task_states dict
+        # make copies to avoid modifying the user-supplied values
         task_states = dict(task_states or {})
-        return_tasks = set(return_tasks or [])
         start_tasks = list(start_tasks or [])
         if any(i not in self.flow.task_ids for i in start_task_ids or []):
             raise ValueError("Invalid start_task_ids.")
         start_tasks.extend(self.flow.task_ids[i] for i in start_task_ids or [])
 
+        # don't make a copy to allow dynamic modification
+        if return_tasks is None:
+            return_tasks = set()
+        if set(return_tasks).difference(self.flow.tasks):
+            raise ValueError("Some tasks in return_tasks were not found in the flow.")
+
         # -- process each task in order
 
         with executor.start():
@@ -343,7 +345,9 @@ def get_flow_run_state(
 
                 # -- run the task
                 task_runner = self.task_runner_cls(
-                    task=task, result_handler=self.flow.result_handler
+                    task=task,
+                    result_handler=self.flow.result_handler,
+                    state_handlers=task_runner_state_handlers,
                 )
 
                 task_states[task] = executor.submit(
@@ -367,45 +371,25 @@ def get_flow_run_state(
             # reference tasks determine flow state
             reference_tasks = self.flow.reference_tasks()
 
-            if return_failed:
-                final_states = executor.wait(task_states)
-                all_final_states = final_states.copy()
-                failed_tasks = []
-                for t, s in final_states.items():
-                    if isinstance(s, (Failed, Retrying)):
-                        failed_tasks.append(t)
-                    elif s.is_mapped():
-                        s.map_states = executor.wait(s.map_states)
-                        s.result = [ms.result for ms in s.map_states]
-                        all_final_states[t] = s.map_states
-                        if any(
-                            [isinstance(r, (Failed, Retrying)) for r in s.map_states]
-                        ):
-                            failed_tasks.append(t)
-
-                return_tasks.update(failed_tasks)
-            else:
-                # wait until all terminal tasks are finished
-                final_tasks = terminal_tasks.union(reference_tasks).union(return_tasks)
-                final_states = executor.wait(
-                    {
-                        t: task_states.get(
-                            t, Pending("Task not evaluated by FlowRunner.")
-                        )
-                        for t in final_tasks
-                    }
-                )
+            # wait until all terminal tasks are finished
+            final_tasks = terminal_tasks.union(reference_tasks).union(return_tasks)
+            final_states = executor.wait(
+                {
+                    t: task_states.get(t, Pending("Task not evaluated by FlowRunner."))
+                    for t in final_tasks
+                }
+            )
 
-                # also wait for any children of Mapped tasks to finish, and add them
-                # to the dictionary to determine flow state
-                all_final_states = final_states.copy()
-                for t, s in list(final_states.items()):
-                    if s.is_mapped():
-                        s.map_states = executor.wait(s.map_states)
-                        s.result = [ms.result for ms in s.map_states]
-                        all_final_states[t] = s.map_states
+            # also wait for any children of Mapped tasks to finish, and add them
+            # to the dictionary to determine flow state
+            all_final_states = final_states.copy()
+            for t, s in list(final_states.items()):
+                if s.is_mapped():
+                    s.map_states = executor.wait(s.map_states)
+                    s.result = [ms.result for ms in s.map_states]
+                    all_final_states[t] = s.map_states
 
-                assert isinstance(final_states, dict)
+            assert isinstance(final_states, dict)
 
         key_states = set(flatten_seq([all_final_states[t] for t in reference_tasks]))
         terminal_states = set(


Test patch:
diff --git a/src/prefect/utilities/tests.py b/src/prefect/utilities/debug.py
similarity index 67%
rename from src/prefect/utilities/tests.py
rename to src/prefect/utilities/debug.py
--- a/src/prefect/utilities/tests.py
+++ b/src/prefect/utilities/debug.py
@@ -6,9 +6,11 @@
 import tempfile
 
 from contextlib import contextmanager
+import textwrap
 from typing import Any, Iterator
 
 import prefect
+from prefect.engine import state
 
 
 def is_serializable(obj: Any, raise_on_error: bool = False) -> bool:
@@ -28,12 +30,14 @@ def is_serializable(obj: Any, raise_on_error: bool = False) -> bool:
         - subprocess.CalledProcessError: if `raise_on_error=True` and the object is not deployable
     """
 
-    template = """
-import cloudpickle
+    template = textwrap.dedent(
+        """
+        import cloudpickle
 
-with open('{}', 'rb') as z76123:
-    res = cloudpickle.load(z76123)
-    """
+        with open('{}', 'rb') as z76123:
+            res = cloudpickle.load(z76123)
+        """
+    )
     bd, binary_file = tempfile.mkstemp()
     sd, script_file = tempfile.mkstemp()
     os.close(bd)
@@ -70,7 +74,7 @@ def raise_on_exception() -> Iterator:
     Example:
         ```python
         from prefect import Flow, task
-        from prefect.utilities.tests import raise_on_exception
+        from prefect.utilities.debug import raise_on_exception
 
         @task
         def div(x):
@@ -85,3 +89,33 @@ def div(x):
     """
     with prefect.context(raise_on_exception=True):
         yield
+
+
+def make_return_failed_handler(failed_tasks_set: set):
+    """
+    This state handler can be used to automatically return any tasks that failed or retried
+    from the FlowRunner.
+
+    NOTE: this will only work with the LocalExecutor, as it depends on sharing a global set
+
+    It can be used like this:
+
+    ```
+    flow = prefect.Flow(...)
+    return_tasks = set()
+    state = flow.run(
+        return_tasks=return_tasks,
+        task_runner_state_handlers=[make_return_failed_handler(return_tasks)])
+    ```
+    """
+
+    def handler(
+        task_runner: "prefect.engine.task_runner.TaskRunner",
+        old_state: state.State,
+        new_state: state.State,
+    ) -> state.State:
+        if isinstance(new_state, (state.Failed, state.Retrying)):
+            failed_tasks_set.add(task_runner.task)
+        return new_state
+
+    return handler
diff --git a/tests/core/test_task_map.py b/tests/core/test_task_map.py
--- a/tests/core/test_task_map.py
+++ b/tests/core/test_task_map.py
@@ -6,7 +6,7 @@
 import prefect
 from prefect.core import Edge, Flow, Parameter, Task
 from prefect.utilities.tasks import task, unmapped
-from prefect.utilities.tests import raise_on_exception
+from prefect.utilities.debug import raise_on_exception
 from prefect.engine.state import Mapped, Success
 
 
diff --git a/tests/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
--- a/tests/engine/test_flow_runner.py
+++ b/tests/engine/test_flow_runner.py
@@ -31,7 +31,7 @@
     TriggerFailed,
 )
 from prefect.triggers import manual_only, any_failed
-from prefect.utilities.tests import raise_on_exception
+from prefect.utilities.debug import raise_on_exception
 from unittest.mock import MagicMock
 
 
@@ -122,8 +122,8 @@ def test_flow_runner_with_invalid_return_tasks():
     task = SuccessTask()
     flow.add_task(task)
     flow_runner = FlowRunner(flow=flow)
-    with pytest.raises(ValueError):
-        flow_runner.run(return_tasks=[1])
+    state = flow_runner.run(return_tasks=[1])
+    assert state.is_failed()
 
 
 def test_flow_runner_with_return_tasks_that_arent_evaluated():
@@ -502,6 +502,7 @@ def test_running_state_finishes(self):
             start_tasks=[],
             start_task_ids=[],
             return_tasks=set(),
+            task_runner_state_handlers=[],
             executor=LocalExecutor(),
         )
         assert new_state.is_successful()
@@ -518,6 +519,7 @@ def test_other_states_raise_endrun(self, state):
                 start_tasks=[],
                 start_task_ids=[],
                 return_tasks=set(),
+                task_runner_state_handlers=[],
                 executor=Executor(),
             )
 
@@ -533,6 +535,7 @@ def determine_final_state(self, *args, **kwargs):
             start_tasks=[],
             start_task_ids=[],
             return_tasks=set(),
+            task_runner_state_handlers=[],
             executor=LocalExecutor(),
         )
         assert new_state.is_failed()
@@ -770,64 +773,6 @@ def run(self, x, s):
         assert flow_state.result[y].result == 100
 
 
-class TestReturnFailed:
-    def test_return_failed_works_when_all_fail(self):
-        with Flow() as f:
-            s = SuccessTask()
-            e = ErrorTask()
-            s.set_upstream(e)
-        state = FlowRunner(flow=f).run(return_failed=True)
-        assert state.is_failed()
-        assert e in state.result
-        assert s in state.result
-
-    def test_return_failed_works_when_non_terminal_fails(self):
-        with Flow() as f:
-            s = SuccessTask(trigger=any_failed)
-            e = ErrorTask()
-            s.set_upstream(e)
-        state = FlowRunner(flow=f).run(return_failed=True)
-        assert state.is_successful()
-        assert e in state.result
-        assert s not in state.result
-
-    def test_return_failed_works_with_mapping(self):
-        @prefect.task
-        def div(x):
-            return 1 / x
-
-        @prefect.task
-        def gimme(x):
-            return x
-
-        with Flow() as f:
-            res = gimme.map(div.map(x=[1, 0, 42]))
-
-        state = FlowRunner(flow=f).run(return_failed=True)
-        assert state.is_failed()
-        assert len(state.result) == 2
-        assert all([len(v.result) == 3 for v in state.result.values()])
-
-    def test_return_failed_doesnt_duplicate(self):
-        with Flow() as f:
-            s = SuccessTask()
-            e = ErrorTask()
-            s.set_upstream(e)
-        state = FlowRunner(flow=f).run(return_tasks=f.tasks, return_failed=True)
-        assert state.is_failed()
-        assert len(state.result) == 2
-
-    def test_return_failed_includes_retries(self):
-        with Flow() as f:
-            s = SuccessTask()
-            e = ErrorTask(max_retries=1, retry_delay=datetime.timedelta(0))
-            s.set_upstream(e)
-        state = FlowRunner(flow=f).run(return_failed=True)
-        assert state.is_running()
-        assert e in state.result
-        assert isinstance(state.result[e], Retrying)
-
-
 class TestRunCount:
     def test_run_count_updates_after_each_retry(self):
         flow = Flow()
@@ -1085,7 +1030,7 @@ def test_multiple_flow_handlers_are_called_in_sequence(self):
         # previous result
         flow = Flow(state_handlers=[lambda *a: None, flow_handler])
         with pytest.raises(AssertionError):
-            with prefect.utilities.tests.raise_on_exception():
+            with prefect.utilities.debug.raise_on_exception():
                 FlowRunner(flow=flow).run()
 
     def test_task_handler_that_doesnt_return_state(self):
@@ -1093,7 +1038,7 @@ def test_task_handler_that_doesnt_return_state(self):
         # raises an attribute error because it tries to access a property of the state that
         # doesn't exist on None
         with pytest.raises(AttributeError):
-            with prefect.utilities.tests.raise_on_exception():
+            with prefect.utilities.debug.raise_on_exception():
                 FlowRunner(flow=flow).run()
 
 
@@ -1115,7 +1060,7 @@ def test_multiple_task_runner_handlers_are_called_in_sequence(self):
         # and raise an error, as long as the flow_handlers are called in sequence on the
         # previous result
         with pytest.raises(AssertionError):
-            with prefect.utilities.tests.raise_on_exception():
+            with prefect.utilities.debug.raise_on_exception():
                 FlowRunner(
                     flow=Flow(), state_handlers=[lambda *a: None, flow_runner_handler]
                 ).run()
@@ -1124,7 +1069,7 @@ def test_task_runner_handler_that_doesnt_return_state(self):
         # raises an attribute error because it tries to access a property of the state that
         # doesn't exist on None
         with pytest.raises(AttributeError):
-            with prefect.utilities.tests.raise_on_exception():
+            with prefect.utilities.debug.raise_on_exception():
                 FlowRunner(flow=Flow(), state_handlers=[lambda *a: None]).run()
 
     def test_task_handler_that_raises_signal_is_trapped(self):
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -40,7 +40,7 @@
 from prefect.engine.task_runner import ENDRUN, TaskRunner
 from prefect.utilities.configuration import set_temporary_config
 from prefect.utilities.tasks import pause_task
-from prefect.utilities.tests import raise_on_exception
+from prefect.utilities.debug import raise_on_exception
 
 
 class SuccessTask(Task):
@@ -913,14 +913,14 @@ def task_handler(task_runner, old_state, new_state):
         # previous result
         task = Task(state_handlers=[lambda *a: None, task_handler])
         with pytest.raises(AssertionError):
-            with prefect.utilities.tests.raise_on_exception():
+            with prefect.utilities.debug.raise_on_exception():
                 TaskRunner(task=task).run()
 
     def test_task_handler_that_doesnt_return_state(self):
         # this will raise an error because no state is returned
         task = Task(state_handlers=[lambda *a: None])
         with pytest.raises(AttributeError):
-            with prefect.utilities.tests.raise_on_exception():
+            with prefect.utilities.debug.raise_on_exception():
                 TaskRunner(task=task).run()
 
 
@@ -985,7 +985,7 @@ def task_runner_handler(task_runner, old_state, new_state):
             assert isinstance(new_state, State)
 
         with pytest.raises(AssertionError):
-            with prefect.utilities.tests.raise_on_exception():
+            with prefect.utilities.debug.raise_on_exception():
                 TaskRunner(
                     task=Task(),
                     state_handlers=[lambda *a: Ellipsis, task_runner_handler],
@@ -994,7 +994,7 @@ def task_runner_handler(task_runner, old_state, new_state):
     def test_task_runner_handler_that_doesnt_return_state(self):
         # raises an error because the state handler doesn't return a state
         with pytest.raises(AttributeError):
-            with prefect.utilities.tests.raise_on_exception():
+            with prefect.utilities.debug.raise_on_exception():
                 TaskRunner(task=Task(), state_handlers=[lambda *a: None]).run()
 
     def test_task_handler_that_raises_signal_is_trapped(self):
diff --git a/tests/tasks/test_shell.py b/tests/tasks/test_shell.py
--- a/tests/tasks/test_shell.py
+++ b/tests/tasks/test_shell.py
@@ -6,7 +6,7 @@
 from prefect import Flow
 from prefect.engine import signals
 from prefect.tasks.shell import ShellTask
-from prefect.utilities.tests import raise_on_exception
+from prefect.utilities.debug import raise_on_exception
 
 
 def test_shell_initializes_and_runs_basic_cmd():
diff --git a/tests/tasks/test_sqlite.py b/tests/tasks/test_sqlite.py
--- a/tests/tasks/test_sqlite.py
+++ b/tests/tasks/test_sqlite.py
@@ -5,7 +5,7 @@
 
 from prefect import Flow
 from prefect.tasks.database import SQLiteQueryTask, SQLiteScriptTask
-from prefect.utilities.tests import raise_on_exception
+from prefect.utilities.debug import raise_on_exception
 
 
 sql_script = """
diff --git a/tests/tasks/test_strings.py b/tests/tasks/test_strings.py
--- a/tests/tasks/test_strings.py
+++ b/tests/tasks/test_strings.py
@@ -6,7 +6,7 @@
 from prefect import Flow
 from prefect.engine import signals
 from prefect.tasks.templates import StringFormatterTask, JinjaTemplateTask
-from prefect.utilities.tests import raise_on_exception
+from prefect.utilities.debug import raise_on_exception
 
 
 def test_string_formatter_simply_formats():
diff --git a/tests/utilities/test_debug.py b/tests/utilities/test_debug.py
new file mode 100644
--- /dev/null
+++ b/tests/utilities/test_debug.py
@@ -0,0 +1,283 @@
+from datetime import timedelta
+import os
+import pytest
+import subprocess
+import tempfile
+import textwrap
+
+import prefect
+from prefect.core import Flow, Task
+from prefect.engine import FlowRunner, TaskRunner, state
+from prefect.utilities.configuration import set_temporary_config
+from prefect.utilities.debug import (
+    is_serializable,
+    raise_on_exception,
+    make_return_failed_handler,
+)
+
+
+class SuccessTask(Task):
+    def run(self):
+        return 1
+
+
+class ErrorTask(Task):
+    def run(self):
+        raise ValueError("I had a problem.")
+
+
+class BusinessTask(Task):
+    def run(self):
+        raise prefect.engine.signals.FAIL("needs more blockchain!")
+
+
+class MathTask(Task):
+    def run(self):
+        1 / 0
+
+
+def assert_script_runs(script):
+    sd, script_file = tempfile.mkstemp()
+    os.close(sd)
+    try:
+        with open(script_file, "w") as sf:
+            sf.write(script)
+        subprocess.check_output(
+            "python {}".format(script_file), shell=True, stderr=subprocess.STDOUT
+        )
+    except Exception as exc:
+        raise exc
+    finally:
+        os.unlink(script_file)
+
+
+def test_raise_on_exception_raises_basic_error():
+    flow = Flow()
+    flow.add_task(MathTask())
+    with pytest.raises(ZeroDivisionError):
+        with raise_on_exception():
+            flow.run()
+
+
+def test_raise_on_exception_raises_basic_prefect_signal():
+    flow = Flow()
+    flow.add_task(BusinessTask())
+    with pytest.raises(prefect.engine.signals.FAIL) as error:
+        with raise_on_exception():
+            flow.run()
+    assert "needs more blockchain!" in str(error)
+
+
+def test_raise_on_exception_works_at_the_task_level_with_error():
+    taskrunner = TaskRunner(task=MathTask())
+    with pytest.raises(ZeroDivisionError):
+        with raise_on_exception():
+            taskrunner.run()
+
+
+def test_raise_on_exception_works_at_the_task_level_with_signal():
+    taskrunner = TaskRunner(task=BusinessTask())
+    with pytest.raises(prefect.engine.signals.FAIL) as error:
+        with raise_on_exception():
+            taskrunner.run()
+    assert "needs more blockchain!" in str(error)
+
+
+def test_that_bad_code_in_flow_runner_is_caught():
+    """
+    Test that an error in the actual FlowRunner code itself (not the execution of user code)
+    is caught
+    """
+
+    class BadFlowRunner(FlowRunner):
+        def get_flow_run_state(self, *args, **kwargs):
+            raise RuntimeError("I represent bad code in the flow runner.")
+
+    flow = Flow()
+    flow.add_task(MathTask())
+    flow_state = BadFlowRunner(flow=flow).run()
+    assert isinstance(flow_state, state.Failed)
+    assert isinstance(flow_state.result, RuntimeError)
+    assert "I represent bad code in the flow runner." == str(flow_state.result)
+
+
+def test_that_bad_code_in_task_runner_is_caught():
+    """
+    Test that an error in the actual TaskRunner code itself (not the execution of user code)
+    is caught
+    """
+
+    class BadTaskRunner(TaskRunner):
+        def get_task_run_state(self, *args, **kwargs):
+            raise RuntimeError("I represent bad code in the task runner.")
+
+    flow = Flow()
+    math_task = MathTask()
+    flow.add_task(math_task)
+    flow_state = FlowRunner(flow=flow, task_runner_cls=BadTaskRunner).run(
+        return_tasks=[math_task]
+    )
+    assert isinstance(flow_state, state.Failed)
+    assert flow_state.result[math_task].is_failed()
+    caught_exc = flow_state.result[math_task].result
+    assert isinstance(caught_exc, RuntimeError)
+    assert "I represent bad code in the task runner." == str(caught_exc)
+
+
+def test_raise_on_exception_plays_well_with_context():
+    flow = Flow()
+    flow.add_task(MathTask())
+    try:
+        assert "raise_on_exception" not in prefect.context
+        with raise_on_exception():
+            assert "raise_on_exception" in prefect.context
+            flow.run()
+    except ZeroDivisionError:
+        assert "raise_on_exception" not in prefect.context
+        pass
+
+
+@pytest.mark.parametrize("obj", [5, "string", lambda x, y, z: None, bool])
+def test_is_serializable_returns_true_for_basic_objects(obj):
+    assert is_serializable(obj) is True
+
+
+def test_is_serializable_returns_false_for_curried_functions_defined_in_main():
+    script = textwrap.dedent(
+        """
+        from toolz import curry
+        from prefect.utilities.debug import is_serializable
+
+        @curry
+        def f(x, y):
+            pass
+
+        g = f(y=5)
+        assert is_serializable(g) is False
+        """
+    )
+    assert_script_runs(script)
+
+
+def test_is_serializable_with_raise_is_informative():
+    script = textwrap.dedent(
+        """
+        import subprocess
+        from toolz import curry
+        from prefect.utilities.debug import is_serializable
+
+        @curry
+        def f(x, y):
+            pass
+
+        g = f(y=5)
+        try:
+            is_serializable(g, raise_on_error=True)
+        except subprocess.CalledProcessError as exc:
+            assert "has no attribute \'f\'" in exc.output.decode()
+        """
+    )
+    assert_script_runs(script)
+
+
+class TestReturnFailedStateHandler:
+    @pytest.fixture(autouse=True)
+    def use_local_executor(self):
+        with set_temporary_config(
+            {"engine.executor.default_class": prefect.engine.executors.LocalExecutor}
+        ):
+            yield
+
+    def test_return_failed_works_when_all_fail(self):
+        with Flow() as f:
+            e = ErrorTask()
+            s1 = Task()
+            s2 = Task()
+            s1.set_upstream(e)
+            s2.set_upstream(s1)
+        return_tasks = set()
+        state = FlowRunner(flow=f).run(
+            return_tasks=return_tasks,
+            task_runner_state_handlers=[make_return_failed_handler(return_tasks)],
+        )
+        assert state.is_failed()
+        assert e in state.result
+        assert s1 in state.result
+        assert s2 in state.result
+        assert return_tasks == set([e, s1, s2])
+
+    def test_return_failed_works_when_all_fail_from_flow_run_method(self):
+        with Flow() as f:
+            e = ErrorTask()
+            s1 = Task()
+            s2 = Task()
+            s1.set_upstream(e)
+            s2.set_upstream(s1)
+        return_tasks = set()
+        state = f.run(
+            return_tasks=return_tasks,
+            task_runner_state_handlers=[make_return_failed_handler(return_tasks)],
+        )
+        assert state.is_failed()
+        assert e in state.result
+        assert s1 in state.result
+        assert s2 in state.result
+        assert return_tasks == set([e, s1, s2])
+
+    def test_return_failed_works_when_non_terminal_fails(self):
+        with Flow() as f:
+            e = ErrorTask()
+            t = Task(trigger=prefect.triggers.always_run)
+            t.set_upstream(e)
+        return_tasks = set()
+        state = FlowRunner(flow=f).run(
+            return_tasks=return_tasks,
+            task_runner_state_handlers=[make_return_failed_handler(return_tasks)],
+        )
+        assert state.is_successful()
+        assert e in state.result
+        assert t not in state.result
+        assert return_tasks == set([e])
+
+    def test_return_failed_works_with_mapping(self):
+        @prefect.task
+        def div(x):
+            return 1 / x
+
+        @prefect.task
+        def gimme(x):
+            return x
+
+        with Flow() as f:
+            d = div.map(x=[1, 0, 42])
+            res = gimme.map(d)
+
+        return_tasks = set()
+        state = FlowRunner(flow=f).run(
+            return_tasks=return_tasks,
+            task_runner_state_handlers=[make_return_failed_handler(return_tasks)],
+        )
+        assert state.is_failed()
+        assert res in state.result
+        assert d in state.result
+
+        # we return the full mapped task
+        assert state.result[d].is_mapped()
+        assert any(s.is_failed() for s in state.result[d].map_states)
+        assert state.result[res].is_mapped()
+        assert any(s.is_failed() for s in state.result[res].map_states)
+
+    def test_return_failed_works_for_retries(self):
+
+        with Flow() as f:
+            e = ErrorTask(max_retries=1, retry_delay=timedelta(0))
+            s1 = Task(max_retries=1, retry_delay=timedelta(0))
+            s1.set_upstream(e)
+        return_tasks = set()
+        state = FlowRunner(flow=f).run(
+            return_tasks=return_tasks,
+            task_runner_state_handlers=[make_return_failed_handler(return_tasks)],
+        )
+        assert state.is_running()
+        assert e in state.result
+        assert s1 not in state.result
diff --git a/tests/utilities/test_tests.py b/tests/utilities/test_tests.py
deleted file mode 100644
--- a/tests/utilities/test_tests.py
+++ /dev/null
@@ -1,164 +0,0 @@
-import os
-import pytest
-import subprocess
-import tempfile
-
-import prefect
-from prefect.core import Flow, Task
-from prefect.engine import FlowRunner, TaskRunner, state
-from prefect.utilities.tests import is_serializable, raise_on_exception
-
-
-class SuccessTask(Task):
-    def run(self):
-        return 1
-
-
-class BusinessTask(Task):
-    def run(self):
-        raise prefect.engine.signals.FAIL("needs more blockchain!")
-
-
-class MathTask(Task):
-    def run(self):
-        1 / 0
-
-
-def assert_script_runs(script):
-    sd, script_file = tempfile.mkstemp()
-    os.close(sd)
-    try:
-        with open(script_file, "w") as sf:
-            sf.write(script)
-        subprocess.check_output(
-            "python {}".format(script_file), shell=True, stderr=subprocess.STDOUT
-        )
-    except Exception as exc:
-        raise exc
-    finally:
-        os.unlink(script_file)
-
-
-def test_raise_on_exception_raises_basic_error():
-    flow = Flow()
-    flow.add_task(MathTask())
-    with pytest.raises(ZeroDivisionError):
-        with raise_on_exception():
-            flow.run()
-
-
-def test_raise_on_exception_raises_basic_prefect_signal():
-    flow = Flow()
-    flow.add_task(BusinessTask())
-    with pytest.raises(prefect.engine.signals.FAIL) as error:
-        with raise_on_exception():
-            flow.run()
-    assert "needs more blockchain!" in str(error)
-
-
-def test_raise_on_exception_works_at_the_task_level_with_error():
-    taskrunner = TaskRunner(task=MathTask())
-    with pytest.raises(ZeroDivisionError):
-        with raise_on_exception():
-            taskrunner.run()
-
-
-def test_raise_on_exception_works_at_the_task_level_with_signal():
-    taskrunner = TaskRunner(task=BusinessTask())
-    with pytest.raises(prefect.engine.signals.FAIL) as error:
-        with raise_on_exception():
-            taskrunner.run()
-    assert "needs more blockchain!" in str(error)
-
-
-def test_that_bad_code_in_flow_runner_is_caught():
-    """
-    Test that an error in the actual FlowRunner code itself (not the execution of user code)
-    is caught
-    """
-
-    class BadFlowRunner(FlowRunner):
-        def get_flow_run_state(self, *args, **kwargs):
-            raise RuntimeError("I represent bad code in the flow runner.")
-
-    flow = Flow()
-    flow.add_task(MathTask())
-    flow_state = BadFlowRunner(flow=flow).run()
-    assert isinstance(flow_state, state.Failed)
-    assert isinstance(flow_state.result, RuntimeError)
-    assert "I represent bad code in the flow runner." == str(flow_state.result)
-
-
-def test_that_bad_code_in_task_runner_is_caught():
-    """
-    Test that an error in the actual TaskRunner code itself (not the execution of user code)
-    is caught
-    """
-
-    class BadTaskRunner(TaskRunner):
-        def get_task_run_state(self, *args, **kwargs):
-            raise RuntimeError("I represent bad code in the task runner.")
-
-    flow = Flow()
-    math_task = MathTask()
-    flow.add_task(math_task)
-    flow_state = FlowRunner(flow=flow, task_runner_cls=BadTaskRunner).run(
-        return_tasks=[math_task]
-    )
-    assert isinstance(flow_state, state.Failed)
-    assert flow_state.result[math_task].is_failed()
-    caught_exc = flow_state.result[math_task].result
-    assert isinstance(caught_exc, RuntimeError)
-    assert "I represent bad code in the task runner." == str(caught_exc)
-
-
-def test_raise_on_exception_plays_well_with_context():
-    flow = Flow()
-    flow.add_task(MathTask())
-    try:
-        assert "raise_on_exception" not in prefect.context
-        with raise_on_exception():
-            assert "raise_on_exception" in prefect.context
-            flow.run()
-    except ZeroDivisionError:
-        assert "raise_on_exception" not in prefect.context
-        pass
-
-
-@pytest.mark.parametrize("obj", [5, "string", lambda x, y, z: None, bool])
-def test_is_serializable_returns_true_for_basic_objects(obj):
-    assert is_serializable(obj) is True
-
-
-def test_is_serializable_returns_false_for_curried_functions_defined_in_main():
-    script = """
-from toolz import curry
-from prefect.utilities.tests import is_serializable
-
-@curry
-def f(x, y):
-    pass
-
-g = f(y=5)
-assert is_serializable(g) is False
-    """
-    assert_script_runs(script)
-
-
-def test_is_serializable_with_raise_is_informative():
-    script = """
-import subprocess
-from toolz import curry
-from prefect.utilities.tests import is_serializable
-
-@curry
-def f(x, y):
-    pass
-
-g = f(y=5)
-try:
-    is_serializable(g, raise_on_error=True)
-except subprocess.CalledProcessError as exc:
-    assert "has no attribute \'f\'" in exc.output.decode()
-    """
-    assert_script_runs(script)