Repository: PrefectHQ/prefect

Issue:
Clarify Docs and Expand Testing for Map / Caching Behavior
Following up on #1071, just wanted to suggest adding a test to capture current map/cache behavior and expanding map caching documentation.

As far as I understand prefect mapping & caching, output caching does not work (should not work ???) in the following situations:

1. Multiple elements in a given map have the same inputs and the mapped task is setup to cache.
2. If the position of an item in a mapped inputs changes, then the cache does not get triggered, e.g. cache retrieval is not only based on the function arguments, but the position of an item in the map.

This means that flows which want to take advantage of caching with mapped inputs need to treat the mapped inputs as ordered lists (effectively each ordered element in a mapped input has its own separate cache).

Any thoughts about changing the map/cache architecture so that the redundant computation within 1) and the duplicate cache elements within 2) are eliminated?

```python
class MockSchedule(prefect.schedules.Schedule):
    call_count = 0

    def next(self, n):
        if self.call_count < 3:
            self.call_count += 1
            return [pendulum.now("utc")]
        else:
            return []

class StatefulTask(Task):
    def __init__(self, returnit=False, **kwargs):
        self.returnit = returnit
        super().__init__(**kwargs)

    call_count = 0
    default = [0, 0, 0]
    
    def run(self):
        self.call_count += 1

        if self.returnit:
            return [self.call_count] * 3
        else:
            self.default[self.call_count - 1] = self.default[self.call_count - 1] + 1
            print(self.default)
            return self.default


    
@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(trigger=prefect.triggers.always_run)
def store_y(y):
    storage["y"].append(y)

t1, t2 = StatefulTask(), StatefulTask(returnit=True)
schedule = MockSchedule()
with Flow(name="test", schedule=schedule) as f:
    res = store_y(return_x.map(x=t1, y=t2))

f.run()

first_run = storage["y"][0]
second_run = storage["y"][1]
third_run = storage["y"][2]

## first run: all elements return 1, e.g. updated in first run
## y: [1, 0, 0]
assert first_run == [1.0, 1.0, 1.0]

## second run: caching for first and third elements; no cross-element caching (e.g. second element not cached even though return_x(x=1) and return_x(x=0) is cached for other map elements)
## y: [1, 1, 0]
assert second_run == [1.0, 2.0, 1.0]

## third run: all tasks succeed, caching for first and second elements
## y: [1, 1, 1]
assert third_run == [1.0, 2.0, 3.0]
```


Discussion:
Interesting, it took me a bit to understand your train of thought but I see what you're getting at now.  I don't have a reply quite yet, I'll need to chew on this a little to think through the implications / how we might design this to also be compatible with our Cloud API.
Small update on this:  interestingly, caching across mapped tasks is already supported in Prefect Cloud; this is because in Cloud (which includes a database), `CachedState`s are stored and retrieved by `task_id` which is shared among all common mapped tasks.  The wrinkle is that in the `prefect` package alone there is no state management service that can be queried for all _potential_ valid caches of a task.  Consequently, the current no-Cloud implementation simply matches mapped tasks with their initial state ordering.

I definitely agree that this leads to undesirable behavior, but I haven't found a way around it (yet).

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,6 +836,7 @@ def _run_on_schedule(
             flow_state.result = {}
         task_states = kwargs.pop("task_states", {})
         flow_state.result.update(task_states)
+        prefect.context.caches = {}
 
         ## run this flow indefinitely, so long as its schedule has future dates
         while True:
@@ -883,25 +884,28 @@ def _run_on_schedule(
 
             ## create next scheduled run
             try:
-                cached_tasks = {}
-
+                # update context cache
                 for t, s in flow_state.result.items():
                     if s.is_cached():
-                        cached_tasks.update({t: s})
+                        cached_sub_states = [s]
                     elif s.is_mapped() and any(
                         sub_state.is_cached() for sub_state in s.map_states
                     ):
                         cached_sub_states = [
-                            sub_state if sub_state.is_cached() else None
+                            sub_state
                             for sub_state in s.map_states
+                            if sub_state.is_cached()
                         ]
-                        cached_tasks.update(
-                            {
-                                t: prefect.engine.state.Mapped(
-                                    map_states=cached_sub_states
-                                )
-                            }
-                        )
+                    else:
+                        cached_sub_states = []
+
+                    fresh_states = [
+                        s
+                        for s in prefect.context.caches.get(t.name, [])
+                        + cached_sub_states
+                        if s.cached_result_expiration > now
+                    ]
+                    prefect.context.caches[t.name] = fresh_states
                 if self.schedule is not None:
                     next_run_time = self.schedule.next(1)[0]
                 else:
@@ -909,7 +913,7 @@ def _run_on_schedule(
             except IndexError:
                 break
             flow_state = prefect.engine.state.Scheduled(
-                start_time=next_run_time, result=cached_tasks
+                start_time=next_run_time, result={}
             )
         return flow_state
 
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,6 +573,15 @@ 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()}


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
@@ -1651,6 +1651,169 @@ def store_y(y):
 
         assert storage == dict(y=[[1, 1, 1], [1, 1, 1], [3, 3, 3]])
 
+    def test_flow_dot_run_handles_cached_states_across_runs(self):
+        class MockSchedule(prefect.schedules.Schedule):
+            call_count = 0
+
+            def next(self, n):
+                if self.call_count < 3:
+                    self.call_count += 1
+                    return [pendulum.now("utc")]
+                else:
+                    return []
+
+        class StatefulTask(Task):
+            def __init__(self, maxit=False, **kwargs):
+                self.maxit = maxit
+                super().__init__(**kwargs)
+
+            call_count = 0
+
+            def run(self):
+                # returns 1 on the first run, 2 on the second run, and 1 on the third
+                self.call_count += 1
+                return self.call_count if self.call_count < 3 else 1
+
+        @task(cache_for=datetime.timedelta(minutes=10), cache_validator=all_inputs)
+        def return_x(x):
+            return round(random.random(), 4)
+
+        storage = {"output": []}
+
+        @task(trigger=prefect.triggers.always_run)
+        def store_output(y):
+            storage["output"].append(y)
+
+        t = StatefulTask()
+        schedule = MockSchedule()
+        with Flow(name="test", schedule=schedule) as f:
+            res = store_output(return_x(x=t))
+
+        f.run()
+
+        first_run = storage["output"][0]
+        second_run = storage["output"][1]
+        third_run = storage["output"][2]
+
+        ## first run: nothing interesting
+        assert first_run > 0
+
+        ## second run: all tasks succeed, no cache used
+        assert first_run != second_run
+
+        ## third run: all tasks succeed, caching from previous runs used
+        assert third_run == first_run
+
+    def test_flow_dot_run_handles_mapped_cached_states_across_runs(self):
+        class MockSchedule(prefect.schedules.Schedule):
+            call_count = 0
+
+            def next(self, n):
+                if self.call_count < 3:
+                    self.call_count += 1
+                    return [pendulum.now("utc")]
+                else:
+                    return []
+
+        class StatefulTask(Task):
+            def __init__(self, maxit=False, **kwargs):
+                self.maxit = maxit
+                super().__init__(**kwargs)
+
+            call_count = 0
+            return_vals = {1: [1], 2: [2, 2], 3: [1, 2, 3]}
+
+            def run(self):
+                # returns [1] on the first run, [2, 2] on the second run, and [1, 2, 3] on the third
+                self.call_count += 1
+                return self.return_vals[self.call_count]
+
+        @task(cache_for=datetime.timedelta(minutes=10), cache_validator=all_inputs)
+        def return_x(x):
+            return round(random.random(), 4)
+
+        storage = {"output": []}
+
+        @task(trigger=prefect.triggers.always_run)
+        def store_output(y):
+            storage["output"].append(y)
+
+        t = StatefulTask()
+        schedule = MockSchedule()
+        with Flow(name="test", schedule=schedule) as f:
+            res = store_output(return_x.map(x=t))
+
+        f.run()
+
+        first_run = storage["output"][0]
+        second_run = storage["output"][1]
+        third_run = storage["output"][2]
+
+        ## first run: nothing interesting
+        assert first_run[0] > 0
+
+        ## second run: all tasks succeed, no cache used
+        assert first_run[0] not in second_run
+
+        ## third run: all tasks succeed, caching from previous runs used
+        assert third_run[0] == first_run[0]
+        assert third_run[1] == second_run[0]
+        assert third_run[2] not in first_run
+        assert third_run[2] not in second_run
+
+    def test_flow_dot_run_handles_mapped_cached_states_with_differing_lengths(self):
+        class MockSchedule(prefect.schedules.Schedule):
+            call_count = 0
+
+            def next(self, n):
+                if self.call_count < 3:
+                    self.call_count += 1
+                    return [pendulum.now("utc")]
+                else:
+                    return []
+
+        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
+                # returns [2] on the first run, [2, 2] on the second run, and [3, 3, 3] on the third
+                return [max(self.call_count, 2)] * self.call_count
+
+        @task(cache_for=datetime.timedelta(minutes=10), cache_validator=all_inputs)
+        def return_x(x):
+            return 1 / (x - 1) + round(random.random(), 4)
+
+        storage = {"output": []}
+
+        @task(trigger=prefect.triggers.always_run)
+        def store_output(y):
+            storage["output"].append(y)
+
+        t = StatefulTask()
+        schedule = MockSchedule()
+        with Flow(name="test", schedule=schedule) as f:
+            res = store_output(return_x.map(x=t))
+
+        f.run()
+
+        first_run = storage["output"][0]
+        second_run = storage["output"][1]
+        third_run = storage["output"][2]
+
+        ## first run: nothing interesting
+        assert first_run[0] > 0
+
+        ## second run: all tasks succeed and use cache
+        assert second_run == [first_run[0], first_run[0]]
+
+        ## third run: all tasks succeed, no caching used
+        assert all(x != first_run[0] for x in third_run)
+
     def test_flow_dot_run_handles_mapped_cached_states_with_non_cached(self):
         class MockSchedule(prefect.schedules.Schedule):
             call_count = 0
@@ -1676,12 +1839,9 @@ def run(self):
                 else:
                     return [self.call_count + i for i in range(3)]
 
-        @task(
-            cache_for=datetime.timedelta(minutes=1),
-            cache_validator=partial_inputs_only(validate_on=["x"]),
-        )
+        @task(cache_for=datetime.timedelta(minutes=10), cache_validator=all_inputs)
         def return_x(x, y):
-            return 1 / (y - 1)
+            return 1 / (y - 1) + round(random.random(), 4)
 
         storage = {"y": []}
 
@@ -1702,14 +1862,13 @@ def store_y(y):
 
         ## first run: one child fails, the other two succeed
         assert isinstance(first_run[0], ZeroDivisionError)
-        assert first_run[1:] == [1.0, 0.5]
 
-        ## second run: all tasks succeed, the latter two use cached state
-        assert second_run[0] == 1.0
-        assert second_run[1:] == [1.0, 0.5]
+        ## second run: all tasks succeed, the first two use cached state
+        assert second_run[:2] == first_run[1:]
+        assert second_run[-1] not in first_run
 
         ## third run: all tasks succeed, no caching used
-        assert third_run == [1 / 2, 1 / 3, 1 / 4]
+        assert all(x not in first_run + second_run for x in third_run)
 
     def test_scheduled_runs_handle_mapped_retries(self):
         class StatefulTask(Task):