Repository: PrefectHQ/prefect

Issue:
Task output does not seem to be cached (despite `cache_for` argument being present)
_(I'm using prefect version `0.5.5`.)_

Here is a simple flow containing a single task that has the `cache_for` argument set.

I would expect this flow to take ~2s the first time it runs and then pretty much complete instantaneously on subsequent runs.
```python
from prefect import task, Flow
from datetime import timedelta
import time

@task(cache_for=timedelta(hours=1))
def return_the_answer():
    time.sleep(2)
    return 42

with Flow("answer") as flow:
    return_the_answer()

flow.run()
```

Here is the output from running in a Jupyter notebook using `%%time`, which confirms that the flow takes ~2s to complete.
```
[2019-07-10 20:45:09,227] INFO - prefect.FlowRunner | Beginning Flow run for 'answer'
[2019-07-10 20:45:09,229] INFO - prefect.FlowRunner | Starting flow run.
[2019-07-10 20:45:09,235] INFO - prefect.TaskRunner | Task 'return_the_answer': Starting task run...
[2019-07-10 20:45:11,240] INFO - prefect.TaskRunner | Task 'return_the_answer': finished task run for task with final state: 'Cached'
[2019-07-10 20:45:11,241] INFO - prefect.FlowRunner | Flow run SUCCESS: all reference tasks succeeded
CPU times: user 14.6 ms, sys: 4.24 ms, total: 18.8 ms
Wall time: 2.02 s
<Success: "All reference tasks succeeded.">
```

The strange thing is that subsequent runs take exactly the same amount of time but I'd expect the result to be cached so they should complete much more quickly.

It's quite possible that this is a misunderstanding on my part rather than a bug. Am I supplying the `cache_for` argument in the wrong place? (I couldn't find an example in the docs which applies it to tasks defined via the `@task` decorator.) Any clarifications would be appreciated. Many thanks!


Discussion:
Hi @maxalbert - this is an excellent question.  There is some misunderstanding, but it is very subtle.

First off, you are providing the `cache_for` argument in the correct place - the `@task` decorator accepts all Task initialization keyword arguments.

The subtlety is that caching behaves a little differently when you are using only the open source engine.  Because the engine alone doesn't have a permanent place to store data, your cache only lives in memory for the duration of the `flow.run()` call.  This means that if you were to run this flow _on a schedule_ and call `flow.run()`, you will observe that caching is occurring in the way you expect.

On the Prefect Cloud platform, your cache will (essentially) live in the database, and so it persists across runs.  

I haven't found a satisfactory middle ground yet that both satisfies the expectations of open source users and is compatible with our API.  Since you want to use caching, let me ask you: where would you expect your cached data to live in this example, and how would you expect to be able to configure that location?
Thanks for the super quick reply and the clarification. 👍 

That's an excellent point about where the cache should live. I hadn't considered that of course if the workflow is defined in a script which is executed at different times then the different flow objects that are created during those executions need to be able to have a shared location to check the results from previous runs.

The context in which I ran into this issue was during prototyping, where I was defining and executing everything within the same interactive Python session - and thus calling `flow.run()` on flows which referenced the _same_ task object multiple times. Therefore my expectation was that the task would simply cache results in-memory and then re-use it while the Python session is still active (effectively emulating `functools.lru_cache()` or similar).

(The lengthy task in question was the first one in a sequence of a few tasks, and I was tweaking those subsequent tasks (and thus re-defining the flow which contained them) but I was keeping the first task fixed, so I expected the result to still be cached. I hope this makes sense, I'm happy to provide an example of my interactive session if required.)

Having this in-memory caching functionality would be quite useful during prototyping where you want to avoid having to wait for a long-running task every time you do a test run of the flow while you are tweaking some tasks in it.

Would this be something worth considering to support in Prefect? (And if so, would it be difficult to implement?)
Ah that makes a lot of sense!  Yea I think there's room for some sort of global persistent dictionary that stores cached states in memory.  Let me try to prototype something and I'll ping you when I've got it working.
Awesome, thank you very much!

Resolving patch:
diff --git a/src/prefect/core/flow.py b/src/prefect/core/flow.py
--- a/src/prefect/core/flow.py
+++ b/src/prefect/core/flow.py
@@ -836,7 +836,7 @@ def _run_on_schedule(
             flow_state.result = {}
         task_states = kwargs.pop("task_states", {})
         flow_state.result.update(task_states)
-        prefect.context.caches = {}
+        prefect.context.setdefault("caches", {})
 
         ## run this flow indefinitely, so long as its schedule has future dates
         while True:
@@ -901,11 +901,11 @@ def _run_on_schedule(
 
                     fresh_states = [
                         s
-                        for s in prefect.context.caches.get(t.name, [])
+                        for s in prefect.context.caches.get(t.cache_key or t.name, [])
                         + cached_sub_states
                         if s.cached_result_expiration > now
                     ]
-                    prefect.context.caches[t.name] = fresh_states
+                    prefect.context.caches[t.cache_key or t.name] = fresh_states
                 if self.schedule is not None:
                     next_run_time = self.schedule.next(1)[0]
                 else:
diff --git a/src/prefect/core/task.py b/src/prefect/core/task.py
--- a/src/prefect/core/task.py
+++ b/src/prefect/core/task.py
@@ -112,7 +112,8 @@ def run(self, x, y):
             whether the cache for this task is still valid (only required if `cache_for`
             is provided; defaults to `prefect.engine.cache_validators.duration_only`)
         - cache_key (str, optional): if provided, a `cache_key` serves as a unique identifier for this Task's cache, and can
-            be shared across both Tasks _and_ Flows
+            be shared across both Tasks _and_ Flows; if not provided, the Task's _name_ will be used if running locally, or the
+            Task's database ID if running in Cloud
         - checkpoint (bool, optional): if this Task is successful, whether to
             store its result using the `result_handler` available during the run; defaults to the value of
             `tasks.defaults.checkpoint` in your user config
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
@@ -573,15 +573,6 @@ def check_task_is_cached(self, state: State, inputs: Dict[str, Result]) -> State
         Raises:
             - ENDRUN: if the task is not ready to run
         """
-        if self.task.cache_for is not None:
-            candidate_states = prefect.context.caches.get(self.task.name, [])
-            sanitized_inputs = {key: res.value for key, res in inputs.items()}
-            for candidate in candidate_states:
-                if self.task.cache_validator(
-                    candidate, sanitized_inputs, prefect.context.get("parameters")
-                ):
-                    candidate._result = candidate._result.to_result()
-                    return candidate
         if state.is_cached():
             assert isinstance(state, Cached)  # mypy assert
             sanitized_inputs = {key: res.value for key, res in inputs.items()}
@@ -591,14 +582,27 @@ def check_task_is_cached(self, state: State, inputs: Dict[str, Result]) -> State
                 state._result = state._result.to_result()
                 return state
             else:
-                self.logger.warning(
-                    "Task '{name}': can't use cache because it "
-                    "is now invalid".format(
-                        name=prefect.context.get("task_full_name", self.task.name)
-                    )
-                )
-                return Pending("Cache was invalid; ready to run.")
-        return state
+                state = Pending("Cache was invalid; ready to run.")
+
+        if self.task.cache_for is not None:
+            candidate_states = prefect.context.caches.get(
+                self.task.cache_key or self.task.name, []
+            )
+            sanitized_inputs = {key: res.value for key, res in inputs.items()}
+            for candidate in candidate_states:
+                if self.task.cache_validator(
+                    candidate, sanitized_inputs, prefect.context.get("parameters")
+                ):
+                    candidate._result = candidate._result.to_result()
+                    return candidate
+
+        self.logger.warning(
+            "Task '{name}': can't use cache because it "
+            "is now invalid".format(
+                name=prefect.context.get("task_full_name", self.task.name)
+            )
+        )
+        return state or Pending("Cache was invalid; ready to run.")
 
     def run_mapped_task(
         self,


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
@@ -47,6 +47,11 @@ def add_flow():
     return f
 
 
+@pytest.fixture
+def clear_context_cache():
+    prefect.context["caches"] = {}
+
+
 class TestCreateFlow:
     """ Test various Flow constructors """
 
@@ -1461,6 +1466,7 @@ def test_serialize_fails_with_no_storage(self):
             s_build = f.serialize(build=True)
 
 
+@pytest.mark.usefixtures("clear_context_cache")
 class TestFlowRunMethod:
     def test_flow_dot_run_runs_on_schedule(self):
         class MockSchedule(prefect.schedules.Schedule):
@@ -1645,6 +1651,44 @@ def store_y(y):
 
         assert storage == dict(y=[1, 1, 3])
 
+    def test_flow_dot_run_handles_cached_states_across_runs(self):
+        class StatefulTask(Task):
+            def __init__(self, maxit=False, **kwargs):
+                self.maxit = maxit
+                super().__init__(**kwargs)
+
+            call_count = 0
+
+            def run(self):
+                self.call_count += 1
+                if self.maxit:
+                    return max(self.call_count, 2)
+                else:
+                    return self.call_count
+
+        @task(
+            cache_for=datetime.timedelta(minutes=1),
+            cache_validator=partial_inputs_only(validate_on=["x"]),
+        )
+        def return_x(x, y):
+            return y
+
+        storage = {"y": []}
+
+        @task
+        def store_y(y):
+            storage["y"].append(y)
+
+        t1, t2 = StatefulTask(maxit=True), StatefulTask()
+        with Flow(name="test") as f:
+            res = store_y(return_x(x=t1, y=t2))
+
+        f.run()
+        f.run()
+        f.run()
+
+        assert storage == dict(y=[1, 1, 3])
+
     def test_flow_dot_run_handles_mapped_cached_states(self):
         class MockSchedule(prefect.schedules.Schedule):
             call_count = 0
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
@@ -820,6 +820,64 @@ def test_reads_result_if_cached_valid(self):
         assert new is state
         assert new.result == 2
 
+    def test_reads_result_from_context_if_cached_valid(self):
+        task = Task(
+            cache_for=timedelta(minutes=1),
+            cache_validator=cache_validators.duration_only,
+        )
+        result = SafeResult("2", result_handler=JSONResultHandler())
+        state = Cached(
+            result=result,
+            cached_result_expiration=pendulum.now("utc") + timedelta(minutes=1),
+        )
+
+        with prefect.context(caches={"Task": [state]}):
+            new = TaskRunner(task).check_task_is_cached(
+                state=Pending(), inputs={"a": Result(1)}
+            )
+        assert new is state
+        assert new.result == 2
+
+    def test_state_kwarg_is_prioritized_over_context_caches(self):
+        task = Task(
+            cache_for=timedelta(minutes=1),
+            cache_validator=cache_validators.duration_only,
+        )
+        state_a = Cached(
+            result=SafeResult("2", result_handler=JSONResultHandler()),
+            cached_result_expiration=pendulum.now("utc") + timedelta(minutes=1),
+        )
+        state_b = Cached(
+            result=SafeResult("99", result_handler=JSONResultHandler()),
+            cached_result_expiration=pendulum.now("utc") + timedelta(minutes=1),
+        )
+
+        with prefect.context(caches={"Task": [state_a]}):
+            new = TaskRunner(task).check_task_is_cached(
+                state=state_b, inputs={"a": Result(1)}
+            )
+        assert new is state_b
+        assert new.result == 99
+
+    def test_reads_result_from_context_with_cache_key_if_cached_valid(self):
+        task = Task(
+            cache_for=timedelta(minutes=1),
+            cache_validator=cache_validators.duration_only,
+            cache_key="FOO",
+        )
+        result = SafeResult("2", result_handler=JSONResultHandler())
+        state = Cached(
+            result=result,
+            cached_result_expiration=pendulum.now("utc") + timedelta(minutes=1),
+        )
+
+        with prefect.context(caches={"FOO": [state]}):
+            new = TaskRunner(task).check_task_is_cached(
+                state=Pending(), inputs={"a": Result(1)}
+            )
+        assert new is state
+        assert new.result == 2
+
 
 class TestSetTaskRunning:
     @pytest.mark.parametrize("state", [Pending()])