Repository: PrefectHQ/prefect

Issue:
Cannot restart flow run after mapped task
An example flow that maps and then fails:
```
from datetime import timedelta
from prefect import *
from prefect.engine.executors import DaskExecutor
from prefect.engine.results import GCSResult
from prefect.environments.storage import GCS

with Flow("TestFlow", result=GCSResult(bucket="model_bigquery_tmp"), storage=GCS(bucket="model_bigquery_tmp"), executor=DaskExecutor("brett-daskscheduler:8786")) as TestFlow:
    @task(max_retries=1, retry_delay=timedelta(seconds=0.1))
    def generate_random_list():
        n = 10
        return list(range(n))

    @task(max_retries=1, retry_delay=timedelta(seconds=0.1))
    def wait(n):
        from time import sleep
        sleep(n)
        return n

    @task(max_retries=1, retry_delay=timedelta(seconds=0.1))
    def fail(values):
        raise ValueError(f"n: {len(values)}")

    values = wait.map(generate_random_list())
    fail(values)
```
On restarting I get the following error:
```
brett_replicahq restarted this flow run
Submitted for execution: Job prefect-job-871b5d1e
Downloading testflow/2020-10-16t14-25-00-009370-00-00 from model_bigquery_tmp
Beginning Flow run for 'TestFlow'
Task 'wait': Starting task run...
Task 'wait': finished task run for task with final state: 'Mapped'
Unexpected error: TypeError("Cannot map over unsubscriptable object of type <class 'NoneType'>: None...")
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/prefect/engine/runner.py", line 48, in inner
    new_state = method(self, state, *args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/prefect/engine/flow_runner.py", line 526, in get_flow_run_state
    executors.prepare_upstream_states_for_mapping(
  File "/usr/local/lib/python3.8/site-packages/prefect/utilities/executors.py", line 372, in prepare_upstream_states_for_mapping
    raise TypeError(
TypeError: Cannot map over unsubscriptable object of type <class 'NoneType'>: None...
```
cc @cicdw seems related to #3322 ?


Discussion:
Hey @bnaul I ran your example and was not able to encounter the error you posted. The first run ran, failed as expected, then when I restarted it successfully ran through the mapped tasks (via their results in GCS) and then failed as expected again. Any more information you could share about your environment / how you are running this?
Very strange...I was running in K8s but now I'm just running locally and see the same behavior, here's everything else:
```
from prefect.environments import LocalEnvironment

TestFlow.environment = LocalEnvironment(
    executor=prefect.engine.executors.DaskExecutor(address="localhost:8786")
)
TestFlow.register("Model")
TestFlow.run_agent(<token>)
# "Quick Run" in the UI
```

Versions:
```
prefect==0.13.12
dask==2.30.0
distributed=2.30.0
```
@joshmeek update: I think I'm seeing it succeed/fail intermittently actually, maybe there's some kinda race condition......? Lemme try re-running a bunch and see what happens
^ actually @joshmeek it's not random: when I run with `LocalEnvironment(DaskExecutor(...))` like above (with a local dask scheduler+workers running in another tab), I reliably get `TypeError: Cannot map over unsubscriptable object of type <class 'NoneType'>: None...`. When I comment out the `environment` and run with the default executor, everything works as intended. That ring any bells for you?
@bnaul Thanks for the clarification! I am able to reproduce, looking into it now
Update: @bnaul we have identified the issue and are working on a fix, _should_ be in by the next release

Resolving patch:
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
@@ -430,8 +430,18 @@ def check_task_ready_to_map(
             - ENDRUN: either way, we dont continue past this point
         """
         if state.is_mapped():
+            # this indicates we are executing a re-run of a mapped pipeline;
+            # in this case, we populate both `map_states` and `cached_inputs`
+            # to ensure the flow runner can properly regenerate the child tasks,
+            # regardless of whether we mapped over an exchanged piece of data
+            # or a non-data-exchanging upstream dependency
             if len(state.map_states) == 0 and state.n_map_states > 0:  # type: ignore
                 state.map_states = [None] * state.n_map_states  # type: ignore
+            state.cached_inputs = {
+                edge.key: state._result  # type: ignore
+                for edge, state in upstream_states.items()
+                if edge.key
+            }
             raise ENDRUN(state)
 
         # we can't map if there are no success states with iterables upstream
diff --git a/src/prefect/utilities/executors.py b/src/prefect/utilities/executors.py
--- a/src/prefect/utilities/executors.py
+++ b/src/prefect/utilities/executors.py
@@ -506,21 +506,14 @@ def prepare_upstream_states_for_mapping(
                         not state.is_mapped()
                         or upstream_state._result != prefect.engine.result.NoResult
                     ):
-                        # this line should never be hit due to a check
-                        # in the TaskRunner when evaluating the mapped parent
                         if not hasattr(upstream_state.result, "__getitem__"):
-                            raise TypeError(
-                                (
-                                    "Cannot map over unsubscriptable object of type {t}: {val}..."
-                                ).format(
-                                    t=type(upstream_state.result),
-                                    val=repr(upstream_state.result)[:10],
-                                )
-                            )
-                        upstream_result = upstream_state._result.from_value(  # type: ignore
-                            upstream_state.result[i]
-                        )
+                            value = None
+                        else:
+                            value = upstream_state.result[i]
+                        upstream_result = upstream_state._result.from_value(value)  # type: ignore
                         states[edge].result = upstream_result
+                        if state.map_states and i >= len(state.map_states):  # type: ignore
+                            raise IndexError()
                     elif state.is_mapped():
                         if i >= len(state.map_states):  # type: ignore
                             raise IndexError()


Test patch:
diff --git a/tests/core/test_task_map.py b/tests/core/test_task_map.py
--- a/tests/core/test_task_map.py
+++ b/tests/core/test_task_map.py
@@ -1154,3 +1154,34 @@ def test_flatmap_one_unmappable_input(self, executor):
             state.result[z].message
             == "At least one upstream state has an unmappable result."
         )
+
+
+def test_mapped_retries_regenerate_child_pipelines():
+    """
+    This test sets up a situation analogous to one found in Cloud: if a reduce task fails, and a user
+    retries it in the future, we want to make sure that the mapped children pipelines are correctly
+    regenerated.  When run against Cloud, these child tasks will correctly query for their states and
+    the run will proceed with the correct data.
+
+    This test mimics this scenario by running this flow with a provided set of states that only contain
+    metadata about the runs with no actual data to reference.  The child runs should still be produced
+    based only on the n_map_states attribute of the parent.
+    """
+    idt = IdTask()
+    ll = ListTask()
+    with Flow("test") as flow:
+        mapped = idt.map(ll)
+        reduced = idt(mapped)
+
+    flow_state = flow.run()
+    assert flow_state.is_successful()
+    assert flow_state.result[mapped].is_mapped()
+    assert flow_state.result[reduced].is_successful()
+    assert flow_state.result[reduced].result == [1, 2, 3]
+
+    second_pass_states = {mapped: Mapped(n_map_states=3), ll: Success(result=Result())}
+
+    new_state = flow.run(task_states=second_pass_states)
+    assert new_state.is_successful()
+    assert new_state.result[mapped].is_mapped()
+    assert new_state.result[reduced].is_successful()
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
@@ -1891,6 +1891,23 @@ def test_run_mapped_returns_mapped(self, state):
             )
         assert exc.value.state.is_mapped()
 
+    @pytest.mark.parametrize("state", [Pending(), Mapped(), Scheduled()])
+    def test_run_mapped_returns_cached_inputs_if_rerun(self, state):
+        """
+        This is important to communicate result information back to the
+        FlowRunner for regenerating the mapped children.
+        """
+        result = LocalResult(value="y")
+        edge = Edge(Task(), Task(), key="x")
+        with pytest.raises(ENDRUN) as exc:
+            TaskRunner(task=Task()).check_task_ready_to_map(
+                state=state, upstream_states={edge: Success(result=result)}
+            )
+        if state.is_mapped():
+            assert exc.value.state.cached_inputs == dict(x=result)
+        else:
+            assert exc.value.state.cached_inputs == dict()
+
     def test_run_mapped_returns_failed_if_no_success_upstream(self):
         with pytest.raises(ENDRUN) as exc:
             TaskRunner(task=Task()).check_task_ready_to_map(