Repository: PrefectHQ/prefect

Issue:
LOOPing a Task depends on the recursion limit for Python
## Description
*A clear description of the bug*
Right now, Looping a task relies on recursion. This can cause the user to experience the following error: `RecursionError: maximum recursion depth exceeded in comparison`. Unless the user updates the system recursion limit for python (which may not be a great idea) the maximum number of times their task can loop is capped.



## Expected Behavior
*What did you expect to happen instead?*

I expected looping to support more than a limited loop count.


## Reproduction
*A minimal example that exhibits the behavior.*

```python
import prefect
from prefect import task, Flow
from prefect.engine.signals import LOOP

@task
def example():
    loop_payload = prefect.context.get("task_loop_result", 0)
    if loop_payload < 4000:
        loop_payload += 1
        raise LOOP(result=loop_payload)

    return loop_payload

with Flow("Example") as flow:
    result = example()
```

I run into the following error:
```bash
In [16]: flow.run()
[2019-12-10 23:16:25,571] INFO - prefect.FlowRunner | Beginning Flow run for 'Example'
[2019-12-10 23:16:25,575] INFO - prefect.FlowRunner | Starting flow run.
[2019-12-10 23:16:25,583] INFO - prefect.TaskRunner | Task 'example': Starting task run...
[2019-12-10 23:16:29,646] ERROR - prefect.TaskRunner | Task 'example': unexpected error while running task: RecursionError('maximum recursion depth exceeded in comparison')
Traceback (most recent call last):
  File "/Users/dylanhughes/dev/prefect/src/prefect/engine/task_runner.py", line 229, in run
    with prefect.context(context):
  File "/Users/dylanhughes/miniconda3/envs/product_flows/lib/python3.7/contextlib.py", line 112, in __enter__
    return next(self.gen)
  File "/Users/dylanhughes/dev/prefect/src/prefect/utilities/context.py", line 111, in __call__
    previous_context = self.copy()
  File "/Users/dylanhughes/dev/prefect/src/prefect/utilities/collections.py", line 107, in copy
    return type(self)(self.__dict__.copy())
  File "/Users/dylanhughes/dev/prefect/src/prefect/utilities/context.py", line 78, in __init__
    super().__init__(*args, **kwargs)
  File "/Users/dylanhughes/dev/prefect/src/prefect/utilities/collections.py", line 62, in __init__
    super().update(init_dict)
  File "/Users/dylanhughes/miniconda3/envs/product_flows/lib/python3.7/_collections_abc.py", line 839, in update
    if isinstance(other, Mapping):
  File "/Users/dylanhughes/miniconda3/envs/product_flows/lib/python3.7/abc.py", line 139, in __instancecheck__
    return _abc_instancecheck(cls, instance)
RecursionError: maximum recursion depth exceeded in comparison
[2019-12-10 23:16:29,750] INFO - prefect.TaskRunner | Task 'example': finished task run for task with final state: 'Failed'
[2019-12-10 23:16:29,753] INFO - prefect.TaskRunner | Task 'example': finished task run for task with final state: 'Failed'
[2019-12-10 23:16:29,753] INFO - prefect.FlowRunner | Flow run FAILED: some reference tasks failed.
Out[16]: <Failed: "Some reference tasks failed.">
```


## Environment
*Any additional information about your environment*
Standard python environment (python 3.7 docker container)


Discussion:
Specific task run for reference https://cloud.prefect.io/prefect-qa2/task-run/fa0d004f-df37-40c4-a2b7-a71318ce0724

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
@@ -3,13 +3,14 @@
 import _thread
 import time
 import warnings
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple
 
 import pendulum
 
 import prefect
 from prefect.client import Client
 from prefect.core import Edge, Task
+from prefect.utilities.executors import tail_recursive
 from prefect.engine.cloud.utilities import prepare_state_for_cloud
 from prefect.engine.result import NoResult, Result
 from prefect.engine.result_handlers import ResultHandler
@@ -249,6 +250,7 @@ def check_task_is_cached(self, state: State, inputs: Dict[str, Result]) -> State
 
         return state
 
+    @tail_recursive
     def run(
         self,
         state: State = None,
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
@@ -46,7 +46,11 @@
     TimedOut,
     TriggerFailed,
 )
-from prefect.utilities.executors import run_with_heartbeat
+from prefect.utilities.executors import (
+    run_with_heartbeat,
+    tail_recursive,
+    RecursiveCall,
+)
 
 if TYPE_CHECKING:
     from prefect.engine.result_handlers import ResultHandler
@@ -177,6 +181,7 @@ def initialize_run(  # type: ignore
 
         return TaskRunnerInitializeResult(state=state, context=context)
 
+    @tail_recursive
     def run(
         self,
         state: State = None,
@@ -310,6 +315,8 @@ def run(
             if exc.state.is_pending() or exc.state.is_failed():
                 exc.state.cached_inputs = task_inputs or {}  # type: ignore
             state = exc.state
+        except RecursiveCall as exc:
+            raise exc
 
         except Exception as exc:
             msg = "Task '{name}': unexpected error while running task: {exc}".format(
@@ -1028,7 +1035,9 @@ def check_task_is_looping(
             )
             context.update(task_run_version=prefect.context.get("task_run_version"))
             new_state = Pending(message=msg)
-            return self.run(
+            raise RecursiveCall(
+                self.run,
+                self,
                 new_state,
                 upstream_states=upstream_states,
                 context=context,
diff --git a/src/prefect/utilities/executors.py b/src/prefect/utilities/executors.py
--- a/src/prefect/utilities/executors.py
+++ b/src/prefect/utilities/executors.py
@@ -286,3 +286,56 @@ def run_with_ctx(*args: Any, _ctx_dict: dict, **kwargs: Any) -> Any:
         return fut.result(timeout=timeout)
     except FutureTimeout:
         raise TimeoutError("Execution timed out.")
+
+
+class RecursiveCall(Exception):
+    def __init__(self, func: Callable, *args: Any, **kwargs: Any):
+        self.func = func
+        self.args = args
+        self.kwargs = kwargs
+
+
+def tail_recursive(func: Callable) -> Callable:
+    """
+    Helper function to facilitate tail recursion of the wrapped function.
+
+    This allows for recursion with unlimited depth since a stack is not allocated for
+    each "nested" call. Note: instead of calling the target function in question, a 
+    `RecursiveCall` exception must be raised instead. 
+
+    Args:
+        - fn (callable): the function to execute
+
+    Returns:
+        - the result of `f(*args, **kwargs)`
+
+    Raises:
+        - RecursionError: if a recursive "call" (raised exception) is made with a function that is 
+            not decorated with `tail_recursive` decorator.
+    """
+
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        while True:
+            try:
+                return func(*args, **kwargs)
+            except RecursiveCall as exc:
+                try:
+                    call_func = getattr(exc.func, "__wrapped_func__")
+                except AttributeError:
+                    raise RecursionError(
+                        "function has not been wrapped to provide tail recursion (func={})".format(
+                            exc.func
+                        )
+                    )
+
+                # there may be multiple nested recursive calls, we should only respond to calls for the
+                # wrapped function explicitly, otherwise allow the call to continue to propagate
+                if call_func != func:
+                    raise exc
+                args = exc.args
+                kwargs = exc.kwargs
+                continue
+
+    setattr(wrapper, "__wrapped_func__", func)
+    return wrapper


Test patch:
diff --git a/tests/utilities/test_executors.py b/tests/utilities/test_executors.py
--- a/tests/utilities/test_executors.py
+++ b/tests/utilities/test_executors.py
@@ -11,7 +11,13 @@
 
 import prefect
 from prefect.utilities.configuration import set_temporary_config
-from prefect.utilities.executors import Heartbeat, timeout_handler, run_with_heartbeat
+from prefect.utilities.executors import (
+    Heartbeat,
+    timeout_handler,
+    run_with_heartbeat,
+    tail_recursive,
+    RecursiveCall,
+)
 
 
 @pytest.mark.parametrize("interval", [0.2, 1])
@@ -224,3 +230,95 @@ def my_fun(x, **kwargs):
 def test_timeout_handler_preserves_logging(caplog):
     timeout_handler(prefect.Flow("logs").run, timeout=2)
     assert len(caplog.records) >= 2  # 1 INFO to start, 1 INFO to end
+
+
+def test_recursion_go_case():
+    @tail_recursive
+    def my_func(a=0):
+        if a > 5:
+            return a
+        raise RecursiveCall(my_func, a + 2)
+
+    assert 6 == my_func()
+
+
+def test_recursion_beyond_python_limits():
+    RECURSION_LIMIT = sys.getrecursionlimit()
+
+    @tail_recursive
+    def my_func(calls=0):
+        if calls > RECURSION_LIMIT + 10:
+            return calls
+        raise RecursiveCall(my_func, calls + 1)
+
+    assert my_func() == RECURSION_LIMIT + 11
+
+
+def test_recursion_nested():
+    def utility_func(a):
+        if a > 5:
+            return a
+        raise RecursiveCall(my_func, a + 2)
+
+    @tail_recursive
+    def my_func(a=0):
+        return utility_func(a)
+
+    assert 6 == my_func()
+
+
+def test_recursion_multiple():
+    call_checkpoints = []
+
+    @tail_recursive
+    def a_func(a=0):
+        call_checkpoints.append(("a", a))
+        if a > 5:
+            return a
+        a = b_func(a + 1)
+        raise RecursiveCall(a_func, (a + 1) * 2)
+
+    @tail_recursive
+    def b_func(b=0):
+        call_checkpoints.append(("b", b))
+        if b > 5:
+            return b
+        b = a_func(b + 2)
+        raise RecursiveCall(b_func, b + 2)
+
+    assert a_func() == 42  # :)
+    assert call_checkpoints == [
+        ("a", 0),
+        ("b", 1),
+        ("a", 3),
+        ("b", 4),
+        ("a", 6),
+        ("b", 8),
+        ("a", 18),
+        ("b", 20),
+        ("a", 42),
+    ]
+
+
+def test_recursion_raises_when_not_decorated():
+    call_checkpoints = []
+
+    @tail_recursive
+    def a_func(a=0):
+        call_checkpoints.append(("a", a))
+        if a > 5:
+            return a
+        a = b_func(a + 1)
+        raise RecursiveCall(a_func, (a + 1) * 2)
+
+    def b_func(b=0):
+        call_checkpoints.append(("b", b))
+        if b > 5:
+            return b
+        b = a_func(b + 2)
+        raise RecursiveCall(b_func, b + 2)
+
+    with pytest.raises(RecursionError):
+        assert a_func()
+
+    assert call_checkpoints == [("a", 0), ("b", 1), ("a", 3), ("b", 4), ("a", 6)]