Repository: PrefectHQ/prefect

Issue:
Catching TimeoutError obscures traceback from underlying task
## Opened from the [Prefect Public Slack Community](https://join.slack.com/t/prefect-public/shared_invite/enQtNzE5OTU3OTQwNzc1LTQ5M2FkZmQzZjI0ODg1ZTBmOTc0ZjVjYWFjMWExZDAyYzBmYjVmMTE1NTQ1Y2IxZTllOTc4MmI3NzYxMDlhYWU)

**sarieddine.marwan**: Hi folks, one of our flow runs just failed due to a task run failing with the following message: `Task timed out during execution.` - this is quite unexpected given we don’t have any timeout set on the prefect task, has anyone else encountered this behavior before ?

**chris**: Hi Marwan, if you don’t have a timeout set on your task then my guess is that your underlying task raised a `TimeoutError` somehow (independently of Prefect); we broadly capture this error and convert it to a `TimedOut` state regardless of whether Prefect raised the error or the underlying task logic raised the error

**sarieddine.marwan**: Hi <@UKNSNMUE6> - thank you for the prompt response  interesting - so we do have a timeout error in place (as part of the underlying function) but I don’t see the error traceback leading back to our timeout

**chris**: I think that actually makes sense given the Prefect code path — we capture that error and don’t log it or re-raise it, instead we somewhat quietly wrap it up into a `TimedOut` state.

Since this is a built-in exception, we probably should create a custom exception to raise and catch instead so we don’t obscure underlying errors

**sarieddine.marwan**: I see - yep creating a custom exception to raise and catch would certainly help when it comes to debugging

**chris**: awesome, i’ll open an issue from this thread - should be straightforward to adjust

**sarieddine.marwan**: thank you :+1:

**chris**: <@ULVA73B9P> open “Catching TimeoutError obscures traceback from underlying task”

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



Discussion:

Patch:
diff --git a/changes/issue4091.yaml b/changes/issue4091.yaml
new file mode 100644
index 0000000000..129f1da216
--- /dev/null
+++ b/changes/issue4091.yaml
@@ -0,0 +1,2 @@
+enhancement:
+  - "Raise custom TimeoutError for task timeouts to allow for more granular user control - [#4091](https://github.com/PrefectHQ/prefect/issues/4091)"
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
index 9a601df55d..c75750aaf6 100644
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -42,6 +42,7 @@ from prefect.utilities.executors import (
     tail_recursive,
 )
 from prefect.utilities.compatibility import nullcontext
+from prefect.utilities.exceptions import TaskTimeoutError
 
 
 TaskRunnerInitializeResult = NamedTuple(
@@ -869,7 +870,7 @@ class TaskRunner(Runner):
                 )
 
         # inform user of timeout
-        except TimeoutError as exc:
+        except TaskTimeoutError as exc:
             if prefect.context.get("raise_on_exception"):
                 raise exc
             state = TimedOut("Task timed out during execution.", result=exc)
diff --git a/src/prefect/utilities/exceptions.py b/src/prefect/utilities/exceptions.py
index 361749df82..4ced1f149b 100644
--- a/src/prefect/utilities/exceptions.py
+++ b/src/prefect/utilities/exceptions.py
@@ -2,6 +2,10 @@ class PrefectError(Exception):
     pass
 
 
+class TaskTimeoutError(PrefectError):
+    pass
+
+
 class ContextError(KeyError, PrefectError):
     pass
 
diff --git a/src/prefect/utilities/executors.py b/src/prefect/utilities/executors.py
index 726feeb86f..86c2d3744d 100644
--- a/src/prefect/utilities/executors.py
+++ b/src/prefect/utilities/executors.py
@@ -15,6 +15,7 @@ from logging import Logger
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Union, Sequence, Mapping
 
 import prefect
+from prefect.utilities.exceptions import TaskTimeoutError
 from prefect.utilities.logging import get_logger
 
 if TYPE_CHECKING:
@@ -96,7 +97,7 @@ def run_with_thread_timeout(
         - args (Sequence): arguments to pass to the function
         - kwargs (Mapping): keyword arguments to pass to the function
         - timeout (int): the length of time to allow for execution before raising a
-            `TimeoutError`, represented as an integer in seconds
+            `TaskTimeoutError`, represented as an integer in seconds
         - logger (Logger): an optional logger to use. If not passed, a logger for the
             `prefect.executors.run_with_thread_timeout` namespace will be created.
         - name (str): an optional name to attach to logs for this function run, defaults
@@ -107,7 +108,7 @@ def run_with_thread_timeout(
         - the result of `fn(*args, **kwargs)`
 
     Raises:
-        - TimeoutError: if function execution exceeds the allowed timeout
+        - TaskTimeoutError: if function execution exceeds the allowed timeout
         - ValueError: if run from outside the main thread
     """
     logger = logger or get_logger()
@@ -118,7 +119,7 @@ def run_with_thread_timeout(
         return fn(*args, **kwargs)
 
     def error_handler(signum, frame):  # type: ignore
-        raise TimeoutError("Execution timed out.")
+        raise TaskTimeoutError("Execution timed out.")
 
     try:
         # Set the signal handler for alarms
@@ -197,7 +198,7 @@ def run_with_multiprocess_timeout(
         - args (Sequence): arguments to pass to the function
         - kwargs (Mapping): keyword arguments to pass to the function
         - timeout (int): the length of time to allow for execution before raising a
-            `TimeoutError`, represented as an integer in seconds
+            `TaskTimeoutError`, represented as an integer in seconds
         - logger (Logger): an optional logger to use. If not passed, a logger for the
             `prefect.` namespace will be created.
         - name (str): an optional name to attach to logs for this function run, defaults
@@ -209,7 +210,7 @@ def run_with_multiprocess_timeout(
 
     Raises:
         - AssertionError: if run from a daemonic process
-        - TimeoutError: if function execution exceeds the allowed timeout
+        - TaskTimeoutError: if function execution exceeds the allowed timeout
     """
     logger = logger or get_logger()
     name = name or f"Function '{fn.__name__}'"
@@ -252,7 +253,7 @@ def run_with_multiprocess_timeout(
             raise result
         return result
     else:
-        raise TimeoutError(f"Execution timed out for {name}.")
+        raise TaskTimeoutError(f"Execution timed out for {name}.")
 
 
 def run_task_with_timeout(
@@ -267,7 +268,7 @@ def run_task_with_timeout(
     The exact implementation varies depending on whether this function is being
     run in the main thread or a non-daemonic subprocess.  If this is run from a
     daemonic subprocess or on Windows, the task is run in a `ThreadPoolExecutor`
-    and only a soft timeout is enforced, meaning a `TimeoutError` is raised at the
+    and only a soft timeout is enforced, meaning a `TaskTimeoutError` is raised at the
     appropriate time but the task continues running in the background.
 
     The task is passed instead of a function so we can give better logs and messages.
@@ -287,7 +288,7 @@ def run_task_with_timeout(
         - the result of `f(*args, **kwargs)`
 
     Raises:
-        - TimeoutError: if function execution exceeds the allowed timeout
+        - TaskTimeoutError: if function execution exceeds the allowed timeout
     """
     logger = logger or get_logger()
     name = prefect.context.get("task_full_name", task.name)
@@ -361,7 +362,7 @@ def run_task_with_timeout(
     try:
         return fut.result(timeout=task.timeout)
     except FutureTimeout as exc:
-        raise TimeoutError(
+        raise TaskTimeoutError(
             f"Execution timed out but was executed {soft_timeout_reason} and will "
             "continue to run in the background."
         ) from exc


Test patch:
diff --git a/tests/core/test_flow.py b/tests/core/test_flow.py
index 6f868bce68..f7df8b2f84 100644
--- a/tests/core/test_flow.py
+++ b/tests/core/test_flow.py
@@ -48,6 +48,7 @@ from prefect.run_configs import LocalRun, UniversalRun
 from prefect.schedules.clocks import ClockEvent
 from prefect.tasks.core.function import FunctionTask
 from prefect.utilities.configuration import set_temporary_config
+from prefect.utilities.exceptions import TaskTimeoutError
 from prefect.utilities.serialization import from_qualified_name
 from prefect.utilities.tasks import task
 from prefect.utilities.edges import unmapped
@@ -3114,7 +3115,7 @@ def test_timeout_actually_stops_execution(
 
     assert state.is_failed()
     assert isinstance(state.result[slow_fn], TimedOut)
-    assert isinstance(state.result[slow_fn].result, TimeoutError)
+    assert isinstance(state.result[slow_fn].result, TaskTimeoutError)
     # We cannot capture the UserWarning because it is being run by a Dask worker
     # but we can make sure the TimeoutError includes a note about it
     assert (
diff --git a/tests/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
index 69dc566805..47749cd250 100644
--- a/tests/engine/test_flow_runner.py
+++ b/tests/engine/test_flow_runner.py
@@ -36,6 +36,7 @@ from prefect.engine.state import (
 from prefect.tasks.secrets import PrefectSecret
 from prefect.triggers import manual_only
 from prefect.utilities.debug import raise_on_exception
+from prefect.utilities.exceptions import TaskTimeoutError
 
 
 class SuccessTask(Task):
@@ -842,7 +843,7 @@ def test_flow_runner_handles_timeouts(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)
+    assert isinstance(state.result[res].result, TaskTimeoutError)
 
 
 def test_flow_runner_handles_timeout_error_with_mproc(mproc):
@@ -854,7 +855,7 @@ def test_flow_runner_handles_timeout_error_with_mproc(mproc):
     state = FlowRunner(flow=flow).run(return_tasks=[res], executor=mproc)
     assert state.is_failed()
     assert isinstance(state.result[res], TimedOut)
-    assert isinstance(state.result[res].result, TimeoutError)
+    assert isinstance(state.result[res].result, TaskTimeoutError)
 
 
 handler_results = collections.defaultdict(lambda: 0)
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
index 9be34ec8fd..7e4e69735a 100644
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -46,6 +46,7 @@ from prefect.engine.state import (
 from prefect.engine.task_runner import ENDRUN, TaskRunner
 from prefect.utilities.configuration import set_temporary_config
 from prefect.utilities.debug import raise_on_exception
+from prefect.utilities.exceptions import TaskTimeoutError
 from prefect.utilities.tasks import pause_task
 
 
@@ -343,7 +344,7 @@ def test_timeout_actually_stops_execution():
 
     assert state.is_failed()
     assert isinstance(state, TimedOut)
-    assert isinstance(state.result, TimeoutError)
+    assert isinstance(state.result, TaskTimeoutError)
 
 
 def test_task_runner_can_handle_timeouts_by_default():
@@ -354,7 +355,7 @@ def test_task_runner_can_handle_timeouts_by_default():
     )
     assert isinstance(state, TimedOut)
     assert "timed out" in state.message
-    assert isinstance(state.result, TimeoutError)
+    assert isinstance(state.result, TaskTimeoutError)
 
 
 def test_task_runner_handles_secrets():
diff --git a/tests/utilities/test_executors.py b/tests/utilities/test_executors.py
index 4c1fc568a1..d67ba63a7f 100644
--- a/tests/utilities/test_executors.py
+++ b/tests/utilities/test_executors.py
@@ -7,6 +7,7 @@ import time
 import pytest
 
 import prefect
+from prefect.utilities.exceptions import TaskTimeoutError
 from prefect.utilities.executors import (
     run_with_thread_timeout,
     run_with_multiprocess_timeout,
@@ -27,7 +28,7 @@ TIMEOUT_HANDLERS = [run_with_thread_timeout, run_with_multiprocess_timeout]
 @pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
 def test_timeout_handler_times_out(timeout_handler):
     slow_fn = lambda: time.sleep(2)
-    with pytest.raises(TimeoutError):
+    with pytest.raises(TaskTimeoutError):
         timeout_handler(slow_fn, timeout=1)
 
 
@@ -57,7 +58,7 @@ def test_timeout_handler_actually_stops_execution(timeout_handler, tmpdir):
 
     start_time = time.time()
     stop_time = start_time + wait_time + max_overhead
-    with pytest.raises(TimeoutError):
+    with pytest.raises(TaskTimeoutError):
         timeout_handler(
             slow_fn, args=(start_path, finish_path, wait_time), timeout=timeout
         )