Repository: PrefectHQ/prefect

Issue:
Templated results are overwritten on raised LOOP
## Opened from the [Prefect Public Slack Community](https://join.slack.com/t/prefect-public/shared_invite/enQtNzE5OTU3OTQwNzc1LTQ5M2FkZmQzZjI0ODg1ZTBmOTc0ZjVjYWFjMWExZDAyYzBmYjVmMTE1NTQ1Y2IxZTllOTc4MmI3NzYxMDlhYWU)

**me1548**: Hey, I’m trying to use the `LOOP` and write the results on each iteration as documented <https://docs.prefect.io/core/advanced_tutorials/using-results.html#looping|here>. However, it doesn’t seem to work. It only writes the last iteration.


```
@task()
def log_output(result):
    logger = prefect.context.get('logger')
    <http://logger.info|logger.info>(result)


@task(result=LocalResult(dir='./results', location='test-{task_loop_count}.prefect'))
def loop_test():
    loop_payload = prefect.context.get("task_loop_result", {})

    n = loop_payload.get("n", 1)
    print(n)

    if n &gt; 5:
        return n

    raise LOOP(f'Iteration {n}', result=dict(n=n+1))


with Flow("Postgres -&gt; BigQuery") as flow:
    x = loop_test()
    log_output(x)
```

_See output logging and diagnostics in thread_

**me1548**: 
```
[2020-11-05 07:11:07] INFO - prefect.FlowRunner | Beginning Flow run for 'Postgres -&gt; BigQuery'
[2020-11-05 07:11:07] DEBUG - prefect.FlowRunner | Using executor type LocalExecutor
[2020-11-05 07:11:07] DEBUG - prefect.FlowRunner | Flow 'Postgres -&gt; BigQuery': Handling state change from Scheduled to Running
[2020-11-05 07:11:07] INFO - prefect.TaskRunner | Task 'loop_test': Starting task run...
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Pending to Running
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Calling task.run() method...
1
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Running to Looped
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Pending to Running
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Calling task.run() method...
2
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Running to Looped
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Pending to Running
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Calling task.run() method...
3
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Running to Looped
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Pending to Running
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Calling task.run() method...
4
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Running to Looped
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Pending to Running
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Calling task.run() method...
5
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Running to Looped
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Pending to Running
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Calling task.run() method...
6
[2020-11-05 07:11:07] DEBUG - prefect.LocalResult | Starting to upload result to test-6.prefect...
[2020-11-05 07:11:07] DEBUG - prefect.LocalResult | Finished uploading result to /Users/joell/joell.dev/Scraper/scraper-next/prefect-sync/results/test-6.prefect...
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'loop_test': Handling state change from Running to Success
[2020-11-05 07:11:07] INFO - prefect.TaskRunner | Task 'loop_test': Finished task run for task with final state: 'Success'
[2020-11-05 07:11:07] INFO - prefect.TaskRunner | Task 'log_output': Starting task run...
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'log_output': Handling state change from Pending to Running
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'log_output': Calling task.run() method...
[2020-11-05 07:11:07] INFO - prefect.log_output | 6
[2020-11-05 07:11:07] DEBUG - prefect.TaskRunner | Task 'log_output': Handling state change from Running to Success
[2020-11-05 07:11:07] INFO - prefect.TaskRunner | Task 'log_output': Finished task run for task with final state: 'Success'
[2020-11-05 07:11:07] INFO - prefect.FlowRunner | Flow run SUCCESS: all reference tasks succeeded
[2020-11-05 07:11:07] DEBUG - prefect.FlowRunner | Flow 'Postgres -&gt; BigQuery': Handling state change from Running to Success
{
  "config_overrides": {},
  "env_vars": [
    "PREFECT__LOGGING__LEVEL",
    "PREFECT__CONTEXT__SECRETS__POSTGRES",
    "PREFECT__CONTEXT__SECRETS__GCP_CREDENTIALS",
    "PREFECT__FLOWS__CHECKPOINTING"
  ],
  "flow_information": {
    "environment": {
      "executor": true,
      "labels": true,
      "logger": true,
      "metadata": {
        "image": true
      },
      "on_exit": false,
      "on_start": false,
      "type": "LocalEnvironment"
    },
    "result": {
      "type": "LocalResult"
    },
    "schedule": {},
    "storage": {
      "_flows": {
        "Postgres -&gt; BigQuery": true
      },
      "_labels": false,
      "add_default_labels": true,
      "directory": true,
      "flows": {
        "Postgres -&gt; BigQuery": true
      },
      "path": false,
      "result": true,
      "secrets": false,
      "stored_as_script": false,
      "type": "Local"
    },
    "task_count": 2
  },
  "system_information": {
    "platform": "macOS-10.15.7-x86_64-i386-64bit",
    "prefect_backend": "server",
    "prefect_version": "0.13.13",
    "python_version": "3.8.6"
  }
}
```


**znicholasbrown**: Hi <@U01DW286KGC> - thanks for reporting this - I can confirm each iteration of the loop is overwriting the previous result as <https://docs.prefect.io/core/advanced_tutorials/using-results.html#looping|described here>, despite a templated location being present (this holds true for strings and callables on the location kwarg).

**znicholasbrown**: <@ULVA73B9P> open "Templated results are overwritten on raised LOOP"

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



Discussion:

Patch:
diff --git a/changes/issue3619.yaml b/changes/issue3619.yaml
new file mode 100644
index 0000000000..8b23031c28
--- /dev/null
+++ b/changes/issue3619.yaml
@@ -0,0 +1,2 @@
+enhancement:
+  - "Checkpoint all iterations of Looped tasks - [#3619](https://github.com/PrefectHQ/prefect/issues/3619)"
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
index 03f4faa235..056df601db 100644
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -827,6 +827,7 @@ class TaskRunner(Runner):
 
         value = None
         raw_inputs = {k: r.value for k, r in inputs.items()}
+        new_state = None
         try:
             self.logger.debug(
                 "Task '{name}': Calling task.run() method...".format(
@@ -859,11 +860,10 @@ class TaskRunner(Runner):
         except signals.LOOP as exc:
             new_state = exc.state
             assert isinstance(new_state, Looped)
-            new_state.result = self.result.from_value(value=new_state.result)
+            value = new_state.result
             new_state.message = exc.state.message or "Task is looping ({})".format(
                 new_state.loop_count
             )
-            return new_state
 
         # checkpoint tasks if a result is present, except for when the user has opted out by
         # disabling checkpointing
@@ -884,6 +884,10 @@ class TaskRunner(Runner):
         else:
             result = self.result.from_value(value=value)
 
+        if new_state is not None:
+            new_state.result = result
+            return new_state
+
         state = Success(result=result, message="Task run succeeded.")
         return state
 


Test patch:
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
index f6f4551a13..841623656d 100644
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -2165,7 +2165,7 @@ class TestLooping:
         assert state.result == 3
 
     @pytest.mark.parametrize("checkpoint", [True, None])
-    def test_looping_only_checkpoints_the_final_result(self, checkpoint):
+    def test_looping_checkpoints_all_iterations(self, checkpoint):
         class Handler(ResultHandler):
             data = []
 
@@ -2189,7 +2189,7 @@ class TestLooping:
         state = TaskRunner(my_task).run(context={"checkpointing": True})
         assert state.is_successful()
         assert state.result == 3
-        assert handler.data == [3]
+        assert handler.data == [1, 2, 3]
 
     def test_looping_works_with_retries(self):
         @prefect.task(max_retries=2, retry_delay=timedelta(seconds=0))