Repository: PrefectHQ/prefect

Issue:
`sorted_return_tasks` should be implemented in `Flow.run`
It's only for pretty printing and thus should be implemented in `Flow.run()`; having it implemented in `FlowRunner.run()` is causing weird bugs / requiring too much logic to be maintained


Discussion:

Resolving patch:
diff --git a/src/prefect/core/flow.py b/src/prefect/core/flow.py
--- a/src/prefect/core/flow.py
+++ b/src/prefect/core/flow.py
@@ -800,7 +800,7 @@ def visualize(self):
         for t in self.tasks:
             is_mapped = any(edge.mapped for edge in self.edges_to(t))
             shape = "box" if is_mapped else "ellipse"
-            name = "mapped: {}".format(t.name) if is_mapped else t.name
+            name = "{} <map>".format(t.name) if is_mapped else t.name
             graph.node(str(id(t)), name, shape=shape)
 
         for e in self.edges:
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
@@ -10,7 +10,7 @@
 from prefect.core import Flow, Task
 from prefect.engine import signals
 from prefect.engine.executors import DEFAULT_EXECUTOR
-from prefect.engine.state import Failed, Pending, Running, State, Success
+from prefect.engine.state import Failed, Pending, Retrying, Running, State, Success
 from prefect.engine.task_runner import TaskRunner
 from prefect.utilities.collections import flatten_seq
 
@@ -103,6 +103,7 @@ def run(
         task_states: Dict[Task, State] = None,
         start_tasks: Iterable[Task] = None,
         return_tasks: Iterable[Task] = None,
+        return_failed: bool = False,
         parameters: Dict[str, Any] = None,
         executor: "prefect.engine.executors.Executor" = None,
         context: Dict[str, Any] = None,
@@ -123,6 +124,9 @@ def run(
                 Defaults to `self.flow.root_tasks()`
             - 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
             - executor (Executor, optional): executor to use when performing
@@ -173,6 +177,7 @@ def run(
                     task_states=task_states,
                     start_tasks=start_tasks,
                     return_tasks=return_tasks,
+                    return_failed=return_failed,
                     executor=executor,
                     task_contexts=task_contexts,
                     throttle=throttle,
@@ -210,6 +215,7 @@ def get_run_state(
         return_tasks: Iterable[Task],
         task_contexts: Dict[Task, Dict[str, Any]],
         executor: "prefect.engine.executors.base.Executor",
+        return_failed: bool = False,
         throttle: Dict[str, int] = None,
     ) -> State:
 
@@ -218,7 +224,6 @@ def get_run_state(
         )
         start_tasks = start_tasks or []
         return_tasks = set(return_tasks or [])
-        sorted_return_tasks = []
         task_contexts = task_contexts or {}
         throttle = throttle or {}
 
@@ -238,9 +243,6 @@ def get_run_state(
 
             for task in self.flow.sorted_tasks(root_tasks=start_tasks):
 
-                if task in return_tasks:
-                    sorted_return_tasks.append(task)
-
                 upstream_states = {}
                 task_inputs = {}
                 maps = {}  # a dictionary of key -> task(s) for mapping
@@ -302,17 +304,29 @@ def get_run_state(
             # reference tasks determine flow state
             reference_tasks = self.flow.reference_tasks()
 
-            final_states = executor.wait(
-                {
-                    t: task_states[t]
-                    for t in terminal_tasks.union(reference_tasks).union(return_tasks)
-                }
-            )
+            if return_failed:
+                final_states = executor.wait(dict(task_states))
+                failed_tasks = [
+                    t
+                    for t, state in final_states.items()
+                    if isinstance(state, (Failed, Retrying))
+                ]
+                return_tasks.update(failed_tasks)
+            else:
+                final_states = executor.wait(
+                    {
+                        t: task_states[t]
+                        for t in terminal_tasks.union(reference_tasks).union(
+                            return_tasks
+                        )
+                    }
+                )
+
             terminal_states = set(
                 flatten_seq([final_states[t] for t in terminal_tasks])
             )
             key_states = set(flatten_seq([final_states[t] for t in reference_tasks]))
-            return_states = {t: final_states[t] for t in sorted_return_tasks}
+            return_states = {t: final_states[t] for t in return_tasks}
 
             # check that the flow is finished
             if not all(s.is_finished() for s in terminal_states):


Test patch:
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
@@ -26,7 +26,7 @@
     Success,
     TriggerFailed,
 )
-from prefect.triggers import manual_only
+from prefect.triggers import manual_only, any_failed
 from prefect.utilities.tests import raise_on_exception
 from unittest.mock import MagicMock
 
@@ -132,19 +132,6 @@ def test_flow_runner_prevents_bad_throttle_values():
     ]  # for py34
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 6), reason="Depends on ordered dictionaries of Python 3.6+"
-)
-def test_return_tasks_are_sorted():
-    flow = prefect.Flow()
-    a, b, c = SuccessTask(), SuccessTask(), SuccessTask()
-    flow.add_edge(a, b)
-    flow.add_edge(b, c)
-    flow_runner = FlowRunner(flow=flow)
-    result = flow_runner.run(return_tasks=[c, b, a])
-    assert list(result.result.keys()) == [a, b, c]
-
-
 def test_flow_runner_runs_basic_flow_with_2_independent_tasks():
     flow = prefect.Flow()
     task1 = SuccessTask()
@@ -530,6 +517,47 @@ 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_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)
+            s.set_upstream(e)
+        state = FlowRunner(flow=f).run(return_failed=True)
+        assert state.is_pending()
+        assert e in state.result
+        assert isinstance(state.result[e], Retrying)
+
+
 @pytest.mark.skipif(
     sys.version_info < (3, 5), reason="dask.distributed does not support Python 3.4"
 )