Repository: PrefectHQ/prefect

Issue:
Add Timeout State
Suggestion: subclass `Failed` to create a `Timeout` state (or `TimedOut` state?). The nature of the failure is specific enough that this could be interesting to trap specifically.


Discussion:
+1 and it might be useful to cache the inputs for this particular `Failed` state as well (even if no retry is set).
This is actually a weird state to think about.  Features I'd expect:
- cached inputs so I can see if there's an issue with them
- ability to rerun this task (using the cached inputs)
- it to be a failed state

Resolving patch:
diff --git a/src/prefect/engine/state.py b/src/prefect/engine/state.py
--- a/src/prefect/engine/state.py
+++ b/src/prefect/engine/state.py
@@ -15,7 +15,7 @@
 When a task is running it will enter a `Running` state which means that the task is
 currently being executed.
 
-The four types of `Finished` states are `Success`, `Failed`, `TriggerFailed`, `Mapped` and
+The six types of `Finished` states are `Success`, `Failed`, `TriggerFailed`, `TimedOut`, `Mapped` and
 `Skipped`.
 """
 import datetime
@@ -362,6 +362,30 @@ class Failed(Finished):
     color = "#FF0000"
 
 
+class TimedOut(Failed):
+    """
+    Finished state indicating failure due to execution timeout.
+
+    Args:
+        - message (str or Exception, optional): Defaults to `None`. A message about the
+            state, which could be an `Exception` (or [`Signal`](signals.html)) that caused it.
+        - result (Any, optional): Defaults to `None`. A data payload for the state.
+        - cached_inputs (dict): Defaults to `None`. A dictionary of input
+        keys to values.  Used / set if the Task requires Retries.
+    """
+
+    color = "#CDC9A5"
+
+    def __init__(
+        self,
+        message: str = None,
+        result: Any = None,
+        cached_inputs: Dict[str, Any] = None,
+    ) -> None:
+        super().__init__(message=message, result=result)
+        self.cached_inputs = cached_inputs
+
+
 class TriggerFailed(Failed):
     """
     Finished state indicating failure due to trigger.
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -22,6 +22,7 @@
     Skipped,
     State,
     Success,
+    TimedOut,
     TriggerFailed,
 )
 from prefect.engine.runner import ENDRUN, Runner, call_state_handlers
@@ -604,7 +605,9 @@ def get_task_run_state(
         except TimeoutError as exc:
             if raise_on_exception:
                 raise exc
-            return Failed("Task timed out during execution.", result=exc)
+            return TimedOut(
+                "Task timed out during execution.", result=exc, cached_inputs=inputs
+            )
 
         # Exceptions are trapped and turned into Failed states
         except Exception as exc:


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,6 +26,7 @@
     Skipped,
     State,
     Success,
+    TimedOut,
     TriggerFailed,
 )
 from prefect.triggers import manual_only, any_failed
@@ -362,7 +363,7 @@ def test_flow_with_multiple_retry_tasks_doesnt_run_them_early():
     state2 = flow.run(return_tasks=flow.tasks, task_states=state1.result)
 
     assert isinstance(state2.result[t2], Retrying)
-    assert state2.result[t2] is state1.result[t2]  # state is not modified at all
+    assert state2.result[t2] == state1.result[t2]  # state is not modified at all
     assert isinstance(state2.result[t3], Failed)  # this task ran
 
     state3 = flow.run(
@@ -844,6 +845,7 @@ def test_flow_runner_handles_timeouts(executor):
 
     state = FlowRunner(flow=flow).run(return_tasks=[res], executor=executor)
     assert state.is_failed()
+    assert isinstance(state.result[res], TimedOut)
     assert "timed out" in state.result[res].message
     assert isinstance(state.result[res].result, TimeoutError)
 
diff --git a/tests/engine/test_state.py b/tests/engine/test_state.py
--- a/tests/engine/test_state.py
+++ b/tests/engine/test_state.py
@@ -15,6 +15,7 @@
     Skipped,
     State,
     Success,
+    TimedOut,
     TriggerFailed,
 )
 from prefect.utilities import json
@@ -176,12 +177,18 @@ def test_failed_is_finished(self):
     def test_trigger_failed_is_finished(self):
         assert issubclass(TriggerFailed, Finished)
 
+    def test_timedout_is_finished(self):
+        assert issubclass(TimedOut, Finished)
+
     def test_skipped_is_finished(self):
         assert issubclass(Skipped, Finished)
 
     def test_skipped_is_success(self):
         assert issubclass(Skipped, Success)
 
+    def test_timedout_is_failed(self):
+        assert issubclass(TimedOut, Failed)
+
     def test_trigger_failed_is_failed(self):
         assert issubclass(TriggerFailed, Failed)
 
@@ -277,6 +284,16 @@ def test_state_type_methods_with_failed_state(self):
         assert not state.is_successful()
         assert state.is_failed()
 
+    def test_state_type_methods_with_timedout_state(self):
+        state = TimedOut(message="")
+        assert not state.is_pending()
+        assert not state.is_running()
+        assert state.is_finished()
+        assert not state.is_skipped()
+        assert not state.is_scheduled()
+        assert not state.is_successful()
+        assert state.is_failed()
+
     def test_state_type_methods_with_trigger_failed_state(self):
         state = TriggerFailed(message="")
         assert not state.is_pending()
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
@@ -33,6 +33,7 @@
     Skipped,
     State,
     Success,
+    TimedOut,
     TriggerFailed,
 )
 from prefect.utilities.tasks import pause_task
@@ -307,9 +308,10 @@ def test_task_runner_prioritizes_inputs():
 def test_task_runner_can_handle_timeouts_by_default():
     sleeper = SlowTask(timeout=timedelta(seconds=1))
     state = TaskRunner(sleeper).run(inputs=dict(secs=2))
-    assert state.is_failed()
+    assert isinstance(state, TimedOut)
     assert "timed out" in state.message
     assert isinstance(state.result, TimeoutError)
+    assert state.cached_inputs == dict(secs=2)
 
 
 def test_task_runner_handles_secrets():