Repository: PrefectHQ/prefect

Issue:
Add Looping documentation
- [x] API references
- [ ] Example Flow
- [ ] Blog Post


Discussion:

Resolving patch:
diff --git a/examples/task_looping.py b/examples/task_looping.py
new file mode 100644
--- /dev/null
+++ b/examples/task_looping.py
@@ -0,0 +1,38 @@
+"""
+This Flow uses Prefect's Task Looping feature for computing
+the largest Fibonacci number less than a given number M. Each
+Fibonacci number is computed using a black-box external API.
+"""
+import requests
+from datetime import timedelta
+
+import prefect
+from prefect import task, Flow, Parameter
+from prefect.engine.signals import LOOP
+
+
+@task(max_retries=5, retry_delay=timedelta(seconds=2))
+def compute_large_fibonacci(M):
+    # we extract the accumulated task loop result from context
+    loop_payload = prefect.context.get("task_loop_result", {})
+
+    n = loop_payload.get("n", 1)
+    fib = loop_payload.get("fib", 1)
+
+    next_fib = requests.post(
+        "https://nemo.api.stdlib.com/fibonacci@0.0.1/", data={"nth": n}
+    ).json()
+
+    if next_fib > M:
+        return fib  # return statements end the loop
+
+    raise LOOP(message=f"Fib {n}={next_fib}", result=dict(n=n + 1, fib=next_fib))
+
+
+with Flow("fibonacci") as flow:
+    M = Parameter("M")
+    fib_num = compute_large_fibonacci(M)
+
+
+flow_state = flow.run(M=100)
+print(flow_state.result[fib_num].result)  # 89
diff --git a/src/prefect/utilities/context.py b/src/prefect/utilities/context.py
--- a/src/prefect/utilities/context.py
+++ b/src/prefect/utilities/context.py
@@ -36,6 +36,8 @@
 | `task_slug` | the slug of the current task |
 | `task_tags` | the tags on the current task |
 | `task_run_count` | the run count of the task run - typically only interesting for retrying tasks |
+| `task_loop_count` | if the Task utilizes looping, the loop count of the task run |
+| `task_loop_result` | if the Task is looping, the current loop result |
 
 In addition, Prefect Cloud supplies some additional context variables:
 


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
@@ -2201,6 +2201,33 @@ def downstream(l):
     assert flow_state.result[final].result == 200 ** 2
 
 
+def test_looping_with_retries_resets_run_count():
+    run_counts = []
+
+    @task(max_retries=1, retry_delay=datetime.timedelta(seconds=0))
+    def looper(x):
+        run_counts.append(prefect.context.get("task_run_count"))
+
+        if (
+            prefect.context.get("task_loop_count") == 2
+            and prefect.context.get("task_run_count", 1) == 1
+        ):
+            raise ValueError("err")
+
+        if prefect.context.get("task_loop_count", 1) < 20:
+            raise LOOP(result=prefect.context.get("task_loop_result", 0) + x)
+        return prefect.context.get("task_loop_result") + x
+
+    with Flow(name="looping") as f:
+        inter = looper(1)
+
+    flow_state = f.run()
+
+    assert flow_state.is_successful()
+    assert flow_state.result[inter].result == 20
+    assert list(filter(lambda x: x == 2, run_counts)) == [2]
+
+
 def test_starting_at_arbitrary_loop_index():
     @task
     def looper(x):