Repository: PrefectHQ/prefect

Issue:
State Handler Doesn’t Respect Retry Signals
## Opened from the [Prefect Public Slack Community](https://join.slack.com/t/prefect-public/shared_invite/enQtNzE5OTU3OTQwNzc1LTQ5M2FkZmQzZjI0ODg1ZTBmOTc0ZjVjYWFjMWExZDAyYzBmYjVmMTE1NTQ1Y2IxZTllOTc4MmI3NzYxMDlhYWU)

**jeffbrainerd**: Hello Prefect team, a question about retrying tasks conditionally. In our task state handler we look at the failure message every time a task fails and we want to retry on some of those errors. Per this doc entry:
<https://docs.prefect.io/core/concepts/execution.html#state-signals> we perform a `raise signals.RETRY` in those cases. Problem is this does not appear to work in the task handler. This code gets called: <https://github.com/PrefectHQ/prefect/blob/5b7015a75ea60bc5a8ffb676b3be2ff406103583/src/prefect/engine/cloud/task_runner.py#L96-L105> and the task ultimately is failed, not retried. So should our code be something more like this:
`return Retrying(msg=msg, context=new_state.context)`
or should this in fact work? Thanks!

**dylan**: Hey <@USBCBEJDV>!

Great question, let me take a look

**dylan**: Hey <@USBCBEJDV>

**dylan**: This is a bug! Thank you for finding it

**dylan**: I’ll open an issue :thumbsup:

**dylan**: <@ULVA73B9P> issue “State Handler Doesn’t Respect Retry Signals”

**dylan**: <@ULVA73B9P> open “State Handler Doesn’t Respect Retry Signals”

Original thread can be found [here](https://prefect-community.slack.com/archives/CL09KU1K7/p1599158497252600?thread_ts=1599158497.252600&cid=CL09KU1K7).




Discussion:

Resolving patch:
diff --git a/src/prefect/engine/cloud/task_runner.py b/src/prefect/engine/cloud/task_runner.py
--- a/src/prefect/engine/cloud/task_runner.py
+++ b/src/prefect/engine/cloud/task_runner.py
@@ -97,6 +97,18 @@ def call_runner_target_handlers(self, old_state: State, new_state: State) -> Sta
             new_state = super().call_runner_target_handlers(
                 old_state=old_state, new_state=new_state
             )
+
+        # PrefectStateSignals are trapped and turned into States
+        except prefect.engine.signals.PrefectStateSignal as exc:
+            self.logger.info(
+                "{name} signal raised: {rep}".format(
+                    name=type(exc).__name__, rep=repr(exc)
+                )
+            )
+            if raise_on_exception:
+                raise exc
+            new_state = exc.state
+
         except Exception as exc:
             msg = "Exception raised while calling state handlers: {}".format(repr(exc))
             self.logger.exception(msg)
@@ -337,7 +349,8 @@ def run(
         Returns:
             - `State` object representing the final post-run state of the Task
         """
-        with prefect.context(context or {}):
+        context = context or {}
+        with prefect.context(context):
             end_state = super().run(
                 state=state,
                 upstream_states=upstream_states,
diff --git a/src/prefect/engine/runner.py b/src/prefect/engine/runner.py
--- a/src/prefect/engine/runner.py
+++ b/src/prefect/engine/runner.py
@@ -52,7 +52,7 @@ def inner(self: "Runner", state: State, *args: Any, **kwargs: Any) -> State:
 
         # PrefectStateSignals are trapped and turned into States
         except signals.PrefectStateSignal as exc:
-            self.logger.debug(
+            self.logger.info(
                 "{name} signal raised: {rep}".format(
                     name=type(exc).__name__, rep=repr(exc)
                 )


Test patch:
diff --git a/tests/engine/cloud/test_cloud_task_runner.py b/tests/engine/cloud/test_cloud_task_runner.py
--- a/tests/engine/cloud/test_cloud_task_runner.py
+++ b/tests/engine/cloud/test_cloud_task_runner.py
@@ -603,6 +603,32 @@ def add(x, y):
         assert states[1]._result.location == "2"
 
 
+def test_task_handlers_handle_retry_signals(client):
+    def state_handler(t, o, n):
+        if n.is_failed():
+            raise prefect.engine.signals.RETRY(
+                "Will retry.", start_time=pendulum.now("utc").add(days=1)
+            )
+
+    @prefect.task(state_handlers=[state_handler])
+    def fn():
+        1 / 0
+
+    state = CloudTaskRunner(task=fn).run()
+
+    assert state.is_retrying()
+    assert state.run_count == 1
+
+    # to make it run
+    state.start_time = pendulum.now("utc")
+    new_state = CloudTaskRunner(task=fn).run(state=state)
+    assert new_state.is_retrying()
+    assert new_state.run_count == 2
+
+    states = [call[1]["state"] for call in client.set_task_run_state.call_args_list]
+    assert [type(s).__name__ for s in states] == ["Running", "Retrying"] * 2
+
+
 def test_state_handler_failures_are_handled_appropriately(client, caplog):
     def bad(*args, **kwargs):
         raise SyntaxError("Syntax Errors are nice because they're so unique")
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
@@ -1686,6 +1686,36 @@ def fn():
         # the task changed state two times: Pending -> Running -> Failed
         assert task_handler.call_count == 2
 
+    def test_task_handlers_respect_signals(self):
+        def state_handler(t, o, n):
+            if n.is_failed():
+                raise prefect.engine.signals.PAUSE("Pausing.")
+
+        @prefect.task(state_handlers=[state_handler])
+        def fn():
+            1 / 0
+
+        state = TaskRunner(task=fn).run()
+        assert isinstance(state, Paused)
+
+    def test_task_handlers_handle_retry_signals(self):
+        def state_handler(t, o, n):
+            if n.is_failed():
+                raise prefect.engine.signals.RETRY("Will retry.")
+
+        @prefect.task(state_handlers=[state_handler])
+        def fn():
+            1 / 0
+
+        state = TaskRunner(task=fn).run()
+
+        assert state.is_retrying()
+        assert state.run_count == 1
+
+        new_state = TaskRunner(task=fn).run(state=state)
+        assert new_state.is_retrying()
+        assert new_state.run_count == 2
+
     def test_multiple_task_handlers_are_called(self):
         task_handler = MagicMock(side_effect=lambda t, o, n: n)
         task = Task(state_handlers=[task_handler, task_handler])