Repository: PrefectHQ/prefect

Issue:
Prefect Tasks do not seem to cache when invoked in a map
Hi,

Trying to get the Prefect cache to work for a mapped task. It seems like the tasks are correctly cached, but something about the function signature prevents the cache from correctly being recalled. New to prefect, so definitely may just be operator error...

```
from prefect import Flow, task
from prefect.schedules import IntervalSchedule
import pendulum
from prefect.engine.cache_validators import all_inputs
import time

# Define flow run frequency
refresh_schedule = IntervalSchedule(
    start_date=pendulum.parse("2019-01-01"),
    interval=pendulum.duration(minutes=1)
)



@task
def a():
    return [1,2]

@task(
    cache_for=pendulum.duration(hours=1),
    cache_validator=all_inputs)
def b(a):
    time.sleep(2.5)
    
    

with Flow("test_flow", schedule=refresh_schedule) as test_flow:
    z = a()
    b.map(a=z)


test_flow.run()
```




Discussion:
@jlewis91 thanks for the bug report!  I am able to reproduce locally.  I think I know what's going on..
Great! Looking forward to a fix, love the lib. And thanks for the quick response!! 😄 

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
@@ -883,9 +883,25 @@ def _run_on_schedule(
 
             ## create next scheduled run
             try:
-                cached_tasks = {
-                    t: s for t, s in flow_state.result.items() if s.is_cached()
-                }
+                cached_tasks = {}
+
+                for t, s in flow_state.result.items():
+                    if s.is_cached():
+                        cached_tasks.update({t: 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
+                            for sub_state in s.map_states
+                        ]
+                        cached_tasks.update(
+                            {
+                                t: prefect.engine.state.Mapped(
+                                    map_states=cached_sub_states
+                                )
+                            }
+                        )
                 if self.schedule is not None:
                     next_run_time = self.schedule.next(1)[0]
                 else:


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
@@ -1,5 +1,6 @@
 import datetime
 import logging
+import random
 import sys
 import tempfile
 import uuid
@@ -13,7 +14,7 @@
 from prefect.core.edge import Edge
 from prefect.core.flow import Flow
 from prefect.core.task import Parameter, Task
-from prefect.engine.cache_validators import partial_inputs_only
+from prefect.engine.cache_validators import all_inputs, partial_inputs_only
 from prefect.engine.result_handlers import LocalResultHandler, ResultHandler
 from prefect.engine.signals import PrefectError
 from prefect.engine.state import (
@@ -1604,6 +1605,113 @@ def store_y(y):
 
         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
+
+            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
+                if self.maxit:
+                    return [max(self.call_count, 2)] * 3
+                else:
+                    return [self.call_count] * 3
+
+        @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()
+        schedule = MockSchedule()
+        with Flow(name="test", schedule=schedule) as f:
+            res = store_y(return_x.map(x=t1, y=t2))
+
+        f.run()
+
+        assert storage == dict(y=[[1, 1, 1], [1, 1, 1], [3, 3, 3]])
+
+    def test_flow_dot_run_handles_mapped_cached_states_with_non_cached(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
+                if self.maxit:
+                    return [max(self.call_count, 2)] * 3
+                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"]),
+        )
+        def return_x(x, y):
+            return 1 / (y - 1)
+
+        storage = {"y": []}
+
+        @task(trigger=prefect.triggers.always_run)
+        def store_y(y):
+            storage["y"].append(y)
+
+        t1, t2 = StatefulTask(maxit=True), StatefulTask()
+        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: 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]
+
+        ## third run: all tasks succeed, no caching used
+        assert third_run == [1 / 2, 1 / 3, 1 / 4]
+
     def test_scheduled_runs_handle_mapped_retries(self):
         class StatefulTask(Task):
             call_count = 0