Repository: PrefectHQ/prefect

Issue:
Allow Flow States to be determined by tasks other than the terminal tasks
Currently, the `State` of a flow run is determined exclusively by that Flow's terminal tasks using this waterfall:

* If any terminal task fails, the flow fails
* If all terminal tasks succeed, the flow succeeds
* If all terminal tasks finish, the flow succeeds (this is a separate waterfall step incase we decide that `SKIPs` - a valid "finished" state - should be handled differently
* Otherwise the flow remains pending

However, we can imagine situations where this logic breaks down:

* A flow consists of a task that does something (X) and a terminal clean up task that only runs if X fails.
* If X succeeds, the clean up task will be skipped and the flow will be marked as successful (this is probably the correct behavior)
* If X fails, the clean up task will run and be successful -- the flow will be marked as successful. This might be the wrong behavior; here are two arguments why:
  * The flow failed its intended purpose (which was to run X), and should be considered an overall failure.
  * The workflow ran properly -- catching an error and cleaning up -- and should be considered an overall success.

I propose allowing users to specify the tasks that should determine flow state. By default, it's the terminal tasks, but in the above example a user might have said that the X task should be used to determine flow state. No matter what tasks are chosen, the same waterfall described at the top of this issue would apply to those tasks. This would be a property of the `Flow` object, not a runtime configuration, and would be provided when the flow itself was created:

```python
my_flow = prefect.Flow(name='demo flow', run_state_tasks = [X_task])
```

I think `run_state_tasks` is a bad name for this option -- open to suggestions.




Discussion:
Brain-monsoon for the name of this `kwarg`:
- `reference_tasks`
- `characteristic_tasks`
- `indicative_tasks`
- `representative_tasks`
- `diagnostic_tasks`
- `marked_tasks`
what about very explicit along the lines of:
```python
prefect.Flow(set_flow_state_from=[x, y])
```
Brings up another point: often, Flows will be instantiated before tasks exist. This option will probably be set post-init (even though it is an init option)... wonder if that changes things. 

```python
with Flow() as flow:
    x = x()
    y = y()

# ?????
flow.indicative_tasks = [x, y]
flow.set_flow_state_from = [x, y]
flow.set_flow_state_from(x, y)
```

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
@@ -47,6 +47,7 @@ def __init__(
         environment: prefect.environments.Environment = None,
         tasks: Iterable[Task] = None,
         edges: Iterable[Edge] = None,
+        key_tasks: Iterable[Task] = None,
     ) -> None:
 
         self.name = name or type(self).__name__
@@ -69,14 +70,30 @@ def __init__(
                 key=e.key,
             )
 
+        self.set_key_tasks(key_tasks or [])
+
         self._prefect_version = prefect.__version__
 
         super().__init__()
 
     def __eq__(self, other: Any) -> bool:
         if type(self) == type(other):
-            s = (self.name, self.version, self.tasks, self.edges)
-            o = (other.name, other.version, other.tasks, other.edges)
+            s = (
+                self.project,
+                self.name,
+                self.version,
+                self.tasks,
+                self.edges,
+                self.key_tasks(),
+            )
+            o = (
+                other.project,
+                other.name,
+                other.version,
+                other.tasks,
+                other.edges,
+                other.key_tasks(),
+            )
             return s == o
         return False
 
@@ -95,6 +112,7 @@ def copy(self) -> "Flow":
         new = copy.copy(self)
         new.tasks = self.tasks.copy()
         new.edges = self.edges.copy()
+        new.set_key_tasks(self._key_tasks)
         return new
 
     # Context Manager ----------------------------------------------------------
@@ -137,6 +155,41 @@ def parameters(self, only_required=False) -> Dict[str, ParameterDetails]:
             if isinstance(t, Parameter) and (t.required if only_required else True)
         }
 
+    def key_tasks(self) -> Set[Task]:
+        """
+        A flow's "key tasks" are used to determine its state when it runs. If all the key
+        tasks are successful, then the flow run is considered successful. However, if
+        any of the key tasks fail, the flow is considered to fail. (Note that skips are
+        counted as successes.)
+
+        By default, a flow's key tasks are its terminal tasks. This means the state of a
+        flow is determined by the last tasks that run.
+
+        In some situations, users may want to customize that behavior; for example, if a
+        flow's terminal tasks are "clean up" tasks for the rest of the flow. The
+        flow.set_key_tasks() method can be used to set custom key_tasks.
+
+        Please note that even if key_tasks are provided that are not terminal tasks, the flow
+        will not be considered "finished" until all terminal tasks have completed. Only then
+        will state be determined, using the key tasks.
+        """
+        if self._key_tasks:
+            return set(self._key_tasks)
+        else:
+            return self.terminal_tasks()
+
+    def set_key_tasks(self, tasks: Iterable[Task]) -> None:
+        """
+        Sets the "key tasks" for the flow. See flow.key_tasks() for more details.
+
+        Args:
+            - tasks ([Task]): the tasks that should be set as a flow's key tasks.
+        """
+        key_tasks = set(tasks)
+        if any(t not in self.tasks for t in key_tasks):
+            raise ValueError("Key tasks must be part of the flow.")
+        self._key_tasks = key_tasks
+
     # Graph --------------------------------------------------------------------
 
     def add_task(self, task: Task) -> Task:
@@ -264,10 +317,19 @@ def downstream_tasks(self, task: Task) -> Set[Task]:
 
     def validate(self) -> None:
         """
-        Checks the flow for cycles and raises an error if one is found.
+        Checks that the flow is valid.
         """
+
+        if any(e.upstream_task not in self.tasks for e in self.edges) or any(
+            e.downstream_task not in self.tasks for e in self.edges
+        ):
+            raise ValueError("Some edges refer to tasks not contained in this flow.")
+
         self.sorted_tasks()
 
+        if any(t not in self.tasks for t in self.key_tasks()):
+            raise ValueError("Some key tasks are not contained in this flow.")
+
     def sorted_tasks(self, root_tasks: Iterable[Task] = None) -> Tuple[Task, ...]:
 
         # begin by getting all tasks under consideration (root tasks and all
@@ -428,6 +490,7 @@ def serialize(self) -> dict:
             environment=self.environment,
             parameters=self.parameters(),
             schedule=self.schedule,
+            key_tasks=self.key_tasks(),
             tasks=[dict(obj_id=obj_ids[t], **t.serialize()) for t in self.tasks],
             edges=[
                 dict(
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
@@ -202,29 +202,44 @@ def get_run_state(
             # Collect results
             # ---------------------------------------------
 
+            # terminal tasks determine if the flow is finished
             terminal_tasks = self.flow.terminal_tasks()
-            final_tasks = terminal_tasks.union(return_tasks)
-            final_states = executor.wait({t: task_states[t] for t in final_tasks})
+            # key tasks determine flow state
+            key_tasks = self.flow.key_tasks()
+
+            final_states = executor.wait(
+                {
+                    t: task_states[t]
+                    for t in terminal_tasks.union(key_tasks).union(return_tasks)
+                }
+            )
             terminal_states = {final_states[t] for t in terminal_tasks}
+            key_states = {final_states[t] for t in key_tasks}
             return_states = {t: final_states[t] for t in return_tasks}
 
-            if any(s.is_failed() for s in terminal_states):
-                self.logger.info("Flow run FAILED: some terminal tasks failed.")
-                state = Failed(
-                    message="Some terminal tasks failed.", result=return_states
+            # check that the flow is finished
+            if not all(s.is_finished() for s in terminal_states):
+                self.logger.info("Flow run PENDING: terminal tasks are incomplete.")
+                state = Pending(
+                    message="Some terminal tasks are still pending.",
+                    result=return_states,
                 )
 
-            elif all(s.is_successful() for s in terminal_states):
-                self.logger.info("Flow run SUCCESS: all terminal tasks succeeded")
+            # check if any key task failed
+            elif any(s.is_failed() for s in key_states):
+                self.logger.info("Flow run FAILED: some key tasks failed.")
+                state = Failed(message="Some key tasks failed.", result=return_states)
+
+            # check if all key tasks succeeded
+            elif all(s.is_successful() for s in key_states):
+                self.logger.info("Flow run SUCCESS: all key tasks succeeded")
                 state = Success(
-                    message="All terminal tasks succeeded.", result=return_states
+                    message="All key tasks succeeded.", result=return_states
                 )
 
+            # check for any unanticipated state that is finished but neither success nor failed
             else:
-                self.logger.info("Flow run PENDING: terminal tasks are incomplete.")
-                state = Pending(
-                    message="Some terminal tasks are still pending.",
-                    result=return_states,
-                )
+                self.logger.info("Flow run SUCCESS: no key tasks failed")
+                state = Success(message="No key tasks failed.", result=return_states)
 
             return state


Test patch:
diff --git a/src/prefect/utilities/tests.py b/src/prefect/utilities/tests.py
--- a/src/prefect/utilities/tests.py
+++ b/src/prefect/utilities/tests.py
@@ -7,6 +7,7 @@
 import prefect
 from prefect import Flow, Task
 from prefect.engine.state import State
+from prefect.configuration import Config
 
 
 @contextmanager
@@ -25,12 +26,12 @@ def set_config(key: str, value: Any):
         ```
     """
     try:
-        old_config = copy.copy(prefect.config.__dict__)
+        old_config = copy.deepcopy(prefect.config.__dict__)
 
         config = prefect.config
         keys = key.split(".")
         for key in keys[:-1]:
-            config = getattr(config, key)
+            config = config.setdefault(key, Config())
         setattr(config, keys[-1], value)
         yield
     finally:
diff --git a/tests/core/test_flow.py b/tests/core/test_flow.py
--- a/tests/core/test_flow.py
+++ b/tests/core/test_flow.py
@@ -252,6 +252,27 @@ def test_eager_cycle_detection_works():
         with pytest.raises(ValueError):
             f.add_edge(t2, t1)
 
+    assert not prefect.config.flows.eager_edge_validation
+
+
+def test_copy():
+    with Flow() as f:
+        t1 = Task()
+        t2 = Task()
+        t3 = Task()
+
+    f.add_edge(t1, t2)
+    f.add_edge(t2, t3)
+    f.set_key_tasks([t1])
+
+    f2 = f.copy()
+    assert f2 == f
+
+    f.add_edge(Task(), Task())
+    assert len(f2.tasks) == len(f.tasks) - 2
+    assert len(f2.edges) == len(f.edges) - 1
+    assert f.key_tasks() == f2.key_tasks() == set([t1])
+
 
 def test_infer_root_tasks():
     with Flow() as f:
@@ -266,6 +287,67 @@ def test_infer_root_tasks():
 
 
 def test_infer_terminal_tasks():
+    with Flow() as f:
+        t1 = Task()
+        t2 = Task()
+        t3 = Task()
+        t4 = Task()
+
+    f.add_edge(t1, t2)
+    f.add_edge(t2, t3)
+    f.add_task(t4)
+
+    assert f.terminal_tasks() == set([t3, t4])
+
+
+def test_key_tasks_are_terminal_tasks_by_default():
+    with Flow() as f:
+        t1 = Task()
+        t2 = Task()
+        t3 = Task()
+        t4 = Task()
+
+    f.add_edge(t1, t2)
+    f.add_edge(t2, t3)
+    f.add_task(t4)
+
+    assert f.key_tasks() == f.terminal_tasks() == set([t3, t4])
+
+
+def test_set_key_tasks():
+    with Flow() as f:
+        t1 = Task()
+        t2 = Task()
+        t3 = Task()
+
+    f.add_edge(t1, t2)
+    f.add_edge(t2, t3)
+
+    f.set_key_tasks([])
+    assert f.key_tasks() == f.terminal_tasks()
+    f.set_key_tasks([t2])
+    assert f.key_tasks() == set([t2])
+
+
+def test_set_key_tasks_at_init_with_empty_flow_raises_error():
+
+    with pytest.raises(ValueError) as exc:
+        Flow(key_tasks=[Task()])
+    assert "must be part of the flow" in str(exc.value)
+
+
+def test_set_key_tasks_at_init():
+    t1 = Task()
+    f = Flow(key_tasks=[t1], tasks=[t1])
+    assert f.key_tasks() == set([t1]) == f.tasks == f.terminal_tasks()
+
+    t2 = Task()
+    f = Flow(key_tasks=[t2], tasks=[t1, t2])
+    assert f.key_tasks() == set([t2])
+
+
+def test_reset_key_tasks_to_terminal_tasks():
+
     with Flow() as f:
         t1 = Task()
         t2 = Task()
@@ -274,7 +356,95 @@ def test_infer_terminal_tasks():
     f.add_edge(t1, t2)
     f.add_edge(t2, t3)
 
-    assert f.terminal_tasks() == set([t3])
+    f.set_key_tasks([t2])
+    assert f.key_tasks() == set([t2])
+    f.set_key_tasks([])
+    assert f.key_tasks() == f.terminal_tasks()
+
+
+def test_key_states_raises_error_if_not_part_of_flow():
+    f = Flow()
+    t1 = Task()
+    with pytest.raises(ValueError):
+        f.set_key_tasks([t1])
+
+
+def test_key_states_raises_error_if_not_iterable():
+    f = Flow()
+    t1 = Task()
+    f.add_task(t1)
+    with pytest.raises(TypeError):
+        f.set_key_tasks(t1)
+
+
+class TestEquality:
+    def test_equality_based_on_tasks(self):
+        f1 = Flow()
+        f2 = Flow()
+
+        t1 = Task()
+        t2 = Task()
+        t3 = Task()
+
+        for f in [f1, f2]:
+            f.add_task(t1)
+            f.add_task(t2)
+        assert f1 == f2
+
+        f2.add_task(t3)
+        assert f1 != f2
+
+    def test_equality_based_on_edges(self):
+        f1 = Flow()
+        f2 = Flow()
+
+        t1 = Task()
+        t2 = Task()
+        t3 = Task()
+
+        for f in [f1, f2]:
+            f.add_edge(t1, t2)
+            f.add_edge(t1, t3)
+        assert f1 == f2
+
+        f2.add_edge(t2, t3)
+        assert f1 != f2
+
+    def test_equality_based_on_name(self):
+        f1 = Flow("hi")
+        f2 = Flow("bye")
+        assert f1 != f2
+
+    def test_equality_based_on_project(self):
+        f1 = Flow("flow", project="1")
+        f2 = Flow("flow", project="1")
+        f3 = Flow("flow", project="2")
+        assert f1 == f2
+        assert f2 != f3
+
+    def test_equality_based_on_version(self):
+        f1 = Flow("flow", version="1")
+        f2 = Flow("flow", version="1")
+        f3 = Flow("flow", version="2")
+        assert f1 == f2
+        assert f2 != f3
+
+    def test_equality_based_on_key_tasks(self):
+        f1 = Flow()
+        f2 = Flow()
+
+        t1 = Task()
+        t2 = Task()
+        t3 = Task()
+
+        for f in [f1, f2]:
+            f.add_edge(t1, t2)
+            f.add_edge(t1, t3)
+
+        f1.set_key_tasks([t2])
+        assert f1 != f2
+        f2.set_key_tasks([t2])
+        assert f1 == f2
 
 
 def test_merge():
@@ -392,3 +562,49 @@ def run(self):
 
     state = f.run(return_tasks=[t], parameters=dict(x=10, y=3, z=9))
     assert state.result[t].result == dict(x=10)
+
+
+def test_validate_cycles():
+    f = Flow()
+    t1 = Task()
+    t2 = Task()
+    f.add_edge(t1, t2)
+    f.add_edge(t2, t1)
+    with pytest.raises(ValueError) as exc:
+        f.validate()
+    assert "cycle found" in str(exc.value).lower()
+
+
+def test_validate_missing_edge_downstream_tasks():
+    f = Flow()
+    t1 = Task()
+    t2 = Task()
+    f.add_edge(t1, t2)
+    f.tasks.remove(t2)
+    with pytest.raises(ValueError) as exc:
+        f.validate()
+    assert "edges refer to tasks" in str(exc.value).lower()
+
+
+def test_validate_missing_edge_upstream_tasks():
+    f = Flow()
+    t1 = Task()
+    t2 = Task()
+    f.add_edge(t1, t2)
+    f.tasks.remove(t1)
+    with pytest.raises(ValueError) as exc:
+        f.validate()
+    assert "edges refer to tasks" in str(exc.value).lower()
+
+
+def test_validate_missing_key_tasks():
+    f = Flow()
+    t1 = Task()
+    t2 = Task()
+    f.add_task(t1)
+    f.add_task(t2)
+    f.set_key_tasks([t1])
+    f.tasks.remove(t1)
+    with pytest.raises(ValueError) as exc:
+        f.validate()
+    assert "key tasks are not contained" in str(exc.value).lower()
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
@@ -298,6 +298,32 @@ def test_missing_parameter_error_is_surfaced():
     assert "required parameter" in str(msg).lower()
 
 
+def test_flow_run_state_determined_by_key_tasks():
+    flow = prefect.Flow()
+    t1 = ErrorTask()
+    t2 = SuccessTask(trigger=prefect.triggers.all_finished)
+    flow.add_edge(t1, t2)
+
+    flow.set_key_tasks([t1])
+    flow_state = flow.run(return_tasks=[t1, t2])
+    assert isinstance(flow_state, Failed)
+    assert isinstance(flow_state.result[t1], Failed)
+    assert isinstance(flow_state.result[t2], Success)
+
+
+def test_flow_run_state_not_determined_by_key_tasks_if_terminal_tasks_are_not_finished():
+    flow = prefect.Flow()
+    t1 = ErrorTask()
+    t2 = RaiseRetryTask(trigger=prefect.triggers.all_finished)
+    flow.add_edge(t1, t2)
+
+    flow.set_key_tasks([t1])
+    flow_state = flow.run(return_tasks=[t1, t2])
+    assert isinstance(flow_state, Pending)
+    assert isinstance(flow_state.result[t1], Failed)
+    assert isinstance(flow_state.result[t2], Retrying)
+
+
 class TestFlowRunner_get_pre_run_state:
     def test_runs_as_expected(self):
         flow = prefect.Flow()
diff --git a/tests/utilities/test_tests.py b/tests/utilities/test_tests.py
--- a/tests/utilities/test_tests.py
+++ b/tests/utilities/test_tests.py
@@ -5,7 +5,7 @@
 import prefect
 from prefect.core import Flow, Task
 from prefect.engine import FlowRunner, TaskRunner, state
-from prefect.utilities.tests import raise_on_exception
+from prefect.utilities.tests import raise_on_exception, set_config
 
 
 class SuccessTask(Task):
@@ -105,3 +105,21 @@ def test_raise_on_exception_raises_basic_error():
         assert "_raise_on_exception" not in prefect.context
     except ZeroDivisionError:
         pass
+
+
+def test_set_config_is_temporary():
+    f1 = Flow()
+    assert f1.version == "1"
+    with set_config("flows.default_version", "5"):
+        f2 = Flow()
+        assert f2.version == "5"
+    f3 = Flow()
+    assert f3.version == "1"
+
+
+def test_set_config_can_invent_new_settings():
+    with set_config("flows.nested.nested_again.val", "5"):
+        assert prefect.config.flows.nested.nested_again.val == "5"
+
+    with pytest.raises(AttributeError):
+        assert prefect.config.flows.nested.nested_again.val == "5"