Repository: PrefectHQ/prefect

Issue:
Add `on_failure_callback` kwarg to tasks
This kwarg would accept a list or a single function with call signature `fn(state: State) -> None` (the return value is not important, but will be ignored) and be called anytime a task goes into a `Failed` state (including if a Task goes into a `Failed` state prior to a retry).  These callback functions would then be converted to state handlers under the hood, and called at the appropriate times.

Note: if the callback fails, we should trap the error and log it.


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
@@ -33,6 +33,7 @@
 from prefect.core.task import Parameter, Task
 from prefect.environments import Environment
 from prefect.utilities import logging
+from prefect.utilities.notifications import callback_factory
 from prefect.utilities.serialization import to_qualified_name
 from prefect.utilities.tasks import as_task, unmapped
 
@@ -116,6 +117,8 @@ def my_task():
                 `state_handler(flow: Flow, old_state: State, new_state: State) -> State`
             If multiple functions are passed, then the `new_state` argument will be the
             result of the previous handler.
+        - on_failure (Callable, optional): A function with signature `fn(flow: Flow, state: State) -> None`
+            which will be called anytime this Flow enters a failure state
         - validate (bool, optional): Whether or not to check the validity of
             the flow (e.g., presence of cycles and illegal keys) after adding the edges passed
             in the `edges` argument. Defaults to the value of `eager_edge_validation` in
@@ -133,7 +136,8 @@ def __init__(
         tasks: Iterable[Task] = None,
         edges: Iterable[Edge] = None,
         reference_tasks: Iterable[Task] = None,
-        state_handlers: Iterable[Callable] = None,
+        state_handlers: List[Callable] = None,
+        on_failure: Callable = None,
         validate: bool = None,
         result_handler: ResultHandler = None,
     ):
@@ -170,6 +174,10 @@ def __init__(
         if state_handlers and not isinstance(state_handlers, collections.Sequence):
             raise TypeError("state_handlers should be iterable.")
         self.state_handlers = state_handlers or []
+        if on_failure is not None:
+            self.state_handlers.append(
+                callback_factory(on_failure, check=lambda s: s.is_failed())
+            )
 
         super().__init__()
 
diff --git a/src/prefect/core/task.py b/src/prefect/core/task.py
--- a/src/prefect/core/task.py
+++ b/src/prefect/core/task.py
@@ -6,13 +6,14 @@
 import uuid
 import warnings
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Set, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Set, Tuple
 
 import prefect
 import prefect.engine.cache_validators
 import prefect.engine.signals
 import prefect.triggers
 from prefect.utilities import logging
+from prefect.utilities.notifications import callback_factory
 
 if TYPE_CHECKING:
     from prefect.core.flow import Flow  # pylint: disable=W0611
@@ -116,6 +117,8 @@ def run(self, x, y):
                 `state_handler(task: Task, old_state: State, new_state: State) -> State`
             If multiple functions are passed, then the `new_state` argument will be the
             result of the previous handler.
+        - on_failure (Callable, optional): A function with signature `fn(task: Task, state: State) -> None`
+            with will be called anytime this Task enters a failure state
 
     Raises:
         - TypeError: if `tags` is of type `str`
@@ -137,7 +140,8 @@ def __init__(
         skip_on_upstream_skip: bool = True,
         cache_for: timedelta = None,
         cache_validator: Callable = None,
-        state_handlers: Iterable[Callable] = None,
+        state_handlers: List[Callable] = None,
+        on_failure: Callable = None,
     ):
 
         self.name = name or type(self).__name__
@@ -200,6 +204,10 @@ def __init__(
         if state_handlers and not isinstance(state_handlers, collections.Sequence):
             raise TypeError("state_handlers should be iterable.")
         self.state_handlers = state_handlers or []
+        if on_failure is not None:
+            self.state_handlers.append(
+                callback_factory(on_failure, check=lambda s: s.is_failed())
+            )
 
     def __repr__(self) -> str:
         return "<Task: {self.name}>".format(self=self)
diff --git a/src/prefect/utilities/notifications.py b/src/prefect/utilities/notifications.py
--- a/src/prefect/utilities/notifications.py
+++ b/src/prefect/utilities/notifications.py
@@ -10,10 +10,60 @@
 
 import requests
 from toolz import curry
+from typing import Any, Callable
 
 import prefect
 
-__all__ = ["gmail_notifier", "slack_notifier"]
+__all__ = ["callback_factory", "gmail_notifier", "slack_notifier"]
+
+
+def callback_factory(
+    fn: Callable[[Any, "prefect.engine.state.State"], Any],
+    check: Callable[["prefect.engine.state.State"], bool],
+) -> Callable:
+    """
+    Utility for generating state handlers which serve as callbacks, under arbitrary
+    state-based checks.
+
+    Args:
+        - fn (Callable): a function with signature `fn(obj, state: State) -> None`
+            which will be called anytime the associated state-check passes; in general,
+            it is expected that this function will have side effects (e.g., sends an email).  The first
+            argument to this function is the `Task` or `Flow` it is attached to.
+        - check (Callable): a function with signature `check(state: State) -> bool`
+            which is used for determining when the callback function should be called
+
+    Returns:
+        - state_handler (Callable): a state handler function which can be attached to both Tasks and Flows
+
+    Example:
+        ```python
+        from prefect import Task, Flow
+        from prefect.utilities.notifications import callback_factory
+
+        fn = lambda obj, state: print(state)
+        check = lambda state: state.is_successful()
+        callback = callback_factory(fn, check)
+
+        t = Task(state_handlers=[callback])
+        f = Flow(tasks=[t], state_handlers=[callback])
+        f.run()
+        # prints:
+        # Success("Task run succeeded.")
+        # Success("All reference tasks succeeded.")
+        ```
+    """
+
+    def state_handler(
+        obj: Any,
+        old_state: "prefect.engine.state.State",
+        new_state: "prefect.engine.state.State",
+    ) -> "prefect.engine.state.State":
+        if check(new_state) is True:
+            fn(obj, new_state)
+        return new_state
+
+    return state_handler
 
 
 def email_message_formatter(tracked_obj, state, email_to):


Test patch:
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
@@ -67,6 +67,10 @@ def test_create_flow_with_schedule(self):
     def test_create_flow_without_state_handler(self):
         assert Flow().state_handlers == []
 
+    def test_create_flow_with_on_failure(self):
+        f = Flow(on_failure=lambda *args: None)
+        assert len(f.state_handlers) == 1
+
     @pytest.mark.parametrize("handlers", [[lambda *a: 1], [lambda *a: 1, lambda *a: 2]])
     def test_create_flow_with_state_handler(self, handlers):
         assert Flow(state_handlers=handlers).state_handlers == handlers
diff --git a/tests/core/test_task.py b/tests/core/test_task.py
--- a/tests/core/test_task.py
+++ b/tests/core/test_task.py
@@ -81,6 +81,10 @@ def test_create_task_without_state_handler(self):
     def test_create_task_with_state_handler(self, handlers):
         assert Task(state_handlers=handlers).state_handlers == handlers
 
+    def test_create_task_with_on_failure(self):
+        t = Task(on_failure=lambda *args: None)
+        assert len(t.state_handlers) == 1
+
     def test_create_task_illegal_handler(self):
         with pytest.raises(TypeError):
             Task(state_handlers=lambda *a: 1)
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
@@ -1027,6 +1027,20 @@ def test_flow_handlers_are_called(self):
         # the flow changed state twice: Pending -> Running -> Success
         assert handler_results["Flow"] == 2
 
+    def test_flow_on_failure_is_not_called(self):
+        on_failure = MagicMock()
+        flow = Flow(on_failure=on_failure, tasks=[Task()])
+        FlowRunner(flow=flow).run()
+        assert not on_failure.called
+
+    def test_task_on_failure_is_called(self):
+        on_failure = MagicMock()
+        flow = Flow(tasks=[ErrorTask()], on_failure=on_failure)
+        FlowRunner(flow=flow).run()
+        assert on_failure.call_count == 1
+        assert on_failure.call_args[0][0] is flow
+        assert on_failure.call_args[0][1].is_failed()
+
     def test_multiple_flow_handlers_are_called(self):
         flow = Flow(state_handlers=[flow_handler, flow_handler])
         FlowRunner(flow=flow).run()
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
@@ -953,6 +953,29 @@ def test_task_handlers_are_called(self):
         # the task changed state twice: Pending -> Running -> Success
         assert task_handler.call_count == 2
 
+    def test_task_on_failure_is_not_called(self):
+        on_failure = MagicMock()
+        task = Task(on_failure=on_failure)
+        TaskRunner(task=task).run()
+        assert not on_failure.called
+
+    def test_task_on_failure_is_called(self):
+        on_failure = MagicMock()
+        task = ErrorTask(on_failure=on_failure)
+        TaskRunner(task=task).run()
+        assert on_failure.call_count == 1
+        assert on_failure.call_args[0][0] is task
+        assert on_failure.call_args[0][1].is_failed()
+
+    def test_task_on_trigger_failure_is_called(self):
+        on_failure = MagicMock()
+        task = Task(on_failure=on_failure)
+        edge = Edge(Task(), task)
+        TaskRunner(task=task).run(upstream_states={edge: Failed()})
+        assert on_failure.call_count == 1
+        assert on_failure.call_args[0][0] is task
+        assert isinstance(on_failure.call_args[0][1], TriggerFailed)
+
     def test_task_handlers_are_called_on_retry(self):
         task_handler = MagicMock(side_effect=lambda t, o, n: n)
 
diff --git a/tests/utilities/test_notifications.py b/tests/utilities/test_notifications.py
--- a/tests/utilities/test_notifications.py
+++ b/tests/utilities/test_notifications.py
@@ -1,5 +1,9 @@
+import cloudpickle
 import json
+import os
 import sys
+import tempfile
+from multiprocessing.pool import ThreadPool
 from unittest.mock import MagicMock
 
 import pytest
@@ -21,12 +25,50 @@
 )
 from prefect.utilities.configuration import set_temporary_config
 from prefect.utilities.notifications import (
+    callback_factory,
     gmail_notifier,
     slack_message_formatter,
     slack_notifier,
 )
 
 
+def test_callback_factory_generates_pickleable_objs():
+    """
+    This test does some heavy lifting to create two lambda functions,
+    and then passes them through callback_factory to generate a state-handler.
+    We then pickle the state handler to a temporary file, delete the functions,
+    and try to unpickle the file in a new Process, asserting that the result is what we
+    think it should be.
+    """
+
+    def load_bytes(fname):
+        import cloudpickle
+
+        with open(fname, "rb") as f:
+            obj = cloudpickle.load(f)
+        return obj(1, 2, 3)
+
+    fn = lambda obj, state: None
+    check = lambda state: True
+    handler = callback_factory(fn, check)
+
+    sd, tmp_file = tempfile.mkstemp()
+    os.close(sd)
+    try:
+        with open(tmp_file, "wb") as bit_file:
+            cloudpickle.dump(handler, bit_file)
+        del fn
+        del check
+        pool = ThreadPool(processes=1)
+        result = pool.apply_async(load_bytes, (tmp_file,))
+        value = result.get()
+        assert value == 3
+    except Exception as exc:
+        raise exc
+    finally:
+        os.unlink(tmp_file)
+
+
 @pytest.mark.parametrize(
     "state",
     [