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:

Patch:
diff --git a/changes/issue3519.yaml b/changes/issue3519.yaml
new file mode 100644
index 0000000000..d912d92da7
--- /dev/null
+++ b/changes/issue3519.yaml
@@ -0,0 +1,2 @@
+fix:
+  - "Fix issue with retrying mapped pipelines on dask - [#3519](https://github.com/PrefectHQ/prefect/issues/3519)"
diff --git a/changes/pr3630.yaml b/changes/pr3630.yaml
deleted file mode 100644
index 6e12358357..0000000000
--- a/changes/pr3630.yaml
+++ /dev/null
@@ -1,5 +0,0 @@
-enhancement:
-  - "`Client.create_project` and `prefect create project` will skip creating the project if the project already exists. - [#3630](https://github.com/PrefectHQ/prefect/pull/3630)"
-
-contributor:
-  - "[Amanda Wee](https://github.com/amanda-wee)"
diff --git a/src/prefect/cli/create.py b/src/prefect/cli/create.py
index 53209d0814..568586ccd1 100644
--- a/src/prefect/cli/create.py
+++ b/src/prefect/cli/create.py
@@ -33,8 +33,7 @@ def create():
 @click.option("--description", "-d", help="Project description to create", hidden=True)
 def project(name, description):
     """
-    Create projects with the Prefect API that organize flows. Does nothing if
-    the project already exists.
+    Create projects with the Prefect API that organize flows.
 
     \b
     Arguments:
diff --git a/src/prefect/client/client.py b/src/prefect/client/client.py
index aa04ff14b8..9c6bfdffcf 100644
--- a/src/prefect/client/client.py
+++ b/src/prefect/client/client.py
@@ -919,14 +919,14 @@ class Client:
 
     def create_project(self, project_name: str, project_description: str = None) -> str:
         """
-        Create a new project if a project with the name provided does not already exist
+        Create a new Project
 
         Args:
-            - project_name (str): the project that should be created
+            - project_name (str): the project that should contain this flow
             - project_description (str, optional): the project description
 
         Returns:
-            - str: the ID of the newly-created or pre-existing project
+            - str: the ID of the newly-created project
 
         Raises:
             - ClientError: if the project creation failed
@@ -937,29 +937,16 @@ class Client:
             }
         }
 
-        try:
-            res = self.graphql(
-                project_mutation,
-                variables=dict(
-                    input=dict(
-                        name=project_name,
-                        description=project_description,
-                        tenant_id=self.active_tenant_id,
-                    )
-                ),
-            )  # type: Any
-        except ClientError as exc:
-            if "'Uniqueness violation.'" in str(exc):
-                project_query = {
-                    "query": {
-                        with_args(
-                            "project", {"where": {"name": {"_eq": project_name}}}
-                        ): {"id": True}
-                    }
-                }
-                res = self.graphql(project_query)
-                return res.data.project[0].id
-            raise
+        res = self.graphql(
+            project_mutation,
+            variables=dict(
+                input=dict(
+                    name=project_name,
+                    description=project_description,
+                    tenant_id=self.active_tenant_id,
+                )
+            ),
+        )  # type: Any
 
         return res.data.create_project.id
 
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
index 056df601db..501a00cbec 100644
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -430,8 +430,18 @@ class TaskRunner(Runner):
             - 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
index 9dbff70d26..2001cbee4e 100644
--- 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/cli/test_create.py b/tests/cli/test_create.py
index 9dd09be604..f9aa3d6c12 100644
--- a/tests/cli/test_create.py
+++ b/tests/cli/test_create.py
@@ -49,22 +49,3 @@ def test_create_project_description(patch_post, cloud_api):
     result = runner.invoke(create, ["project", "test", "-d", "description"])
     assert result.exit_code == 0
     assert "test created" in result.output
-
-
-def test_create_project_that_already_exists(patch_posts):
-    patch_posts(
-        [
-            {
-                "errors": [
-                    {"message": "Uniqueness violation.", "path": ["create_project"]}
-                ],
-                "data": {"create_project": None},
-            },
-            {"data": {"project": [{"id": "proj-id"}]}},
-        ]
-    )
-
-    runner = CliRunner()
-    result = runner.invoke(create, ["project", "test"])
-    assert result.exit_code == 0
-    assert "test created" in result.output
diff --git a/tests/client/test_client.py b/tests/client/test_client.py
index 854bb6b833..4329bb6a9d 100644
--- a/tests/client/test_client.py
+++ b/tests/client/test_client.py
@@ -593,29 +593,6 @@ def test_client_register_with_bad_proj_name(patch_post, monkeypatch, cloud_api):
     assert "prefect create project 'my-default-project'" in str(exc.value)
 
 
-def test_client_create_project_that_already_exists(patch_posts, monkeypatch):
-    patch_posts(
-        [
-            {
-                "errors": [
-                    {"message": "Uniqueness violation.", "path": ["create_project"]}
-                ],
-                "data": {"create_project": None},
-            },
-            {"data": {"project": [{"id": "proj-id"}]}},
-        ]
-    )
-
-    monkeypatch.setattr(
-        "prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
-    )
-
-    with set_temporary_config({"cloud.auth_token": "secret_token", "backend": "cloud"}):
-        client = Client()
-    project_id = client.create_project(project_name="my-default-project")
-    assert project_id == "proj-id"
-
-
 def test_client_register_with_flow_that_cant_be_deserialized(patch_post, monkeypatch):
     patch_post({"data": {"project": [{"id": "proj-id"}]}})
 
diff --git a/tests/core/test_task_map.py b/tests/core/test_task_map.py
index 7646d486d7..18fbd48d27 100644
--- a/tests/core/test_task_map.py
+++ b/tests/core/test_task_map.py
@@ -1154,3 +1154,34 @@ class TestFlatMap:
             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
index 841623656d..5cc5bd63af 100644
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -1891,6 +1891,23 @@ class TestCheckTaskReadyToMapStep:
             )
         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(
diff --git a/tests/utilities/test_logging.py b/tests/utilities/test_logging.py
index 0ea7b443e8..818e9fa9ba 100644
--- a/tests/utilities/test_logging.py
+++ b/tests/utilities/test_logging.py
@@ -36,7 +36,7 @@ def test_root_logger_datefmt_responds_to_config(caplog, datefmt):
             logger = utilities.logging.configure_logging(testing=True)
             logger.error("badness")
             logs = [r for r in caplog.records if r.levelname == "ERROR"]
-            assert logs[0].asctime == datetime.datetime.now().strftime(datefmt)
+            assert logs[0].asctime == datetime.datetime.utcnow().strftime(datefmt)
     finally:
         # reset root_logger
         logger = utilities.logging.configure_logging(testing=True)