Repository: PrefectHQ/prefect

Issue:
Mapped task names are not grouped on dask / bokeh
## Description
*A clear description of the bug*
See #2813 for original context.

From the 0.12.0 release task names weren't being inferred / passed properly to the dask scheduler. That was fixed for single tasks after #2814 was merged. However after just testing out current git head it seems like mapped tasks are copied over _with_ their index in the name which results in `bokeh` looking like:

```
mapped_task [0] 1/1
mapped_task [1] 1/1
mapped_task[2] 1/1
...
```

## Expected Behavior
*What did you expect to happen instead?*

Previous behavior in 0.11.0 where mapped tasks were all named the same thing and so `bokeh` had a single progress bar for them all.

## Reproduction
*A minimal example that exhibits the behavior.*

```python
from distributed import Client
from prefect import Flow, task
from prefect.engine.executors import DaskExecutor

@task
def add_one(x):
    return x + 1

with Flow("test") as flow:
    values = add_one.map(range(10))

cluster = LocalCluster()
client = Client(cluster)

print(f"Dashboard available at: {client.dashboard_link}")

flow.run(executor=DaskExecutor(address=cluster.scheduler_address))
```

## Environment
*Any additional information about your environment*

`prefect`: current git head
`dask[bag]==2.18.1`
`distributed==2.18.0`
`bokeh==2.0.2`

*Optionally run `prefect diagnostics` from the command line and paste the information here*



Discussion:

Resolving patch:
diff --git a/src/prefect/engine/executors/dask.py b/src/prefect/engine/executors/dask.py
--- a/src/prefect/engine/executors/dask.py
+++ b/src/prefect/engine/executors/dask.py
@@ -2,7 +2,7 @@
 import uuid
 import warnings
 from contextlib import contextmanager
-from typing import Any, Callable, Iterator, TYPE_CHECKING, Union
+from typing import Any, Callable, Iterator, TYPE_CHECKING, Union, Optional
 
 from prefect import context
 from prefect.engine.executors.base import Executor
@@ -28,6 +28,18 @@
 }
 
 
+def _make_task_key(
+    task_name: str = "", task_index: int = None, **kwargs: Any
+) -> Optional[str]:
+    """A helper for generating a dask task key from fields set in `extra_context`"""
+    if task_name:
+        suffix = uuid.uuid4().hex
+        if task_index is not None:
+            return f"{task_name}-{task_index}-{suffix}"
+        return f"{task_name}-{suffix}"
+    return None
+
+
 class DaskExecutor(Executor):
     """
     An executor that runs all functions using the `dask.distributed` scheduler.
@@ -210,16 +222,15 @@ def _prep_dask_kwargs(self, extra_context: dict = None) -> dict:
         if extra_context is None:
             extra_context = {}
 
-        task_name = extra_context.get("task_name", "")
-        task_tags = extra_context.get("task_tags", [])
-
         dask_kwargs = {"pure": False}  # type: dict
 
         # set a key for the dask scheduler UI
-        if task_name:
-            dask_kwargs.update(key=f"{task_name}-{uuid.uuid4().hex}")
+        key = _make_task_key(**extra_context)
+        if key is not None:
+            dask_kwargs["key"] = key
 
         # infer from context if dask resources are being utilized
+        task_tags = extra_context.get("task_tags", [])
         dask_resource_tags = [
             tag for tag in task_tags if tag.lower().startswith("dask-resource")
         ]
@@ -347,9 +358,9 @@ def submit(
         import dask
 
         extra_kwargs = {}
-        task_name = (extra_context or {}).get("task_name", "")
-        if task_name:
-            extra_kwargs["dask_key_name"] = f"{task_name}-{uuid.uuid4().hex}"
+        key = _make_task_key(**(extra_context or {}))
+        if key is not None:
+            extra_kwargs["dask_key_name"] = key
         return dask.delayed(fn, pure=False)(*args, **kwargs, **extra_kwargs)
 
     def wait(self, futures: Any) -> Any:
diff --git a/src/prefect/engine/flow_runner.py b/src/prefect/engine/flow_runner.py
--- a/src/prefect/engine/flow_runner.py
+++ b/src/prefect/engine/flow_runner.py
@@ -6,7 +6,6 @@
     NamedTuple,
     Optional,
     Set,
-    Union,
 )
 
 import pendulum
@@ -20,7 +19,6 @@
     Failed,
     Mapped,
     Pending,
-    Retrying,
     Running,
     Scheduled,
     State,
@@ -399,6 +397,13 @@ def get_flow_run_state(
         if set(return_tasks).difference(self.flow.tasks):
             raise ValueError("Some tasks in return_tasks were not found in the flow.")
 
+        def extra_context(task: Task, task_index: int = None) -> dict:
+            return {
+                "task_name": task.name,
+                "task_tags": task.tags,
+                "task_index": task_index,
+            }
+
         # -- process each task in order
 
         with executor.start():
@@ -457,8 +462,6 @@ def get_flow_run_state(
                         result=ConstantResult(value=val),
                     )
 
-                extra_context = dict(task_name=task.name, task_tags=task.tags)
-
                 # handle mapped tasks
                 if any([edge.mapped for edge in upstream_states.keys()]):
 
@@ -484,7 +487,7 @@ def get_flow_run_state(
                             task_runner_state_handlers=task_runner_state_handlers,
                             upstream_mapped_states=upstream_mapped_states,
                             is_mapped_parent=True,
-                            extra_context=extra_context,
+                            extra_context=extra_context(task),
                         )
                     )
 
@@ -513,7 +516,6 @@ def get_flow_run_state(
                             current_state = task_state
 
                         ## this is where each child is submitted for actual work
-                        extra_context.update(task_name=f"{task.name}[{idx}]")
                         submitted_states.append(
                             executor.submit(
                                 run_task,
@@ -529,7 +531,7 @@ def get_flow_run_state(
                                 task_runner_cls=self.task_runner_cls,
                                 task_runner_state_handlers=task_runner_state_handlers,
                                 upstream_mapped_states=upstream_mapped_states,
-                                extra_context=extra_context,
+                                extra_context=extra_context(task, task_index=idx),
                             )
                         )
                     if isinstance(task_states.get(task), Mapped):
@@ -546,7 +548,7 @@ def get_flow_run_state(
                         task_runner_cls=self.task_runner_cls,
                         task_runner_state_handlers=task_runner_state_handlers,
                         upstream_mapped_states=upstream_mapped_states,
-                        extra_context=extra_context,
+                        extra_context=extra_context(task),
                     )
 
             # ---------------------------------------------


Test patch:
diff --git a/tests/engine/executors/test_executors.py b/tests/engine/executors/test_executors.py
--- a/tests/engine/executors/test_executors.py
+++ b/tests/engine/executors/test_executors.py
@@ -85,6 +85,13 @@ def test_submit_sets_task_name(self):
             assert f.key.startswith("inc-")
             assert res == 2
 
+            f = e.submit(
+                lambda x: x + 1, 1, extra_context={"task_name": "inc", "task_index": 1}
+            )
+            (res,) = e.wait([f])
+            assert f.key.startswith("inc-1-")
+            assert res == 2
+
     def test_only_compute_once(self):
         e = LocalDaskExecutor()
         count = 0
@@ -317,9 +324,14 @@ def test_prep_dask_kwargs(self):
         kwargs = executor._prep_dask_kwargs(
             dict(task_name="FISH!", task_tags=["dask-resource:GPU=1"])
         )
-        assert kwargs["key"].startswith("FISH!")
+        assert kwargs["key"].startswith("FISH!-")
         assert kwargs["resources"] == {"GPU": 1.0}
 
+        kwargs = executor._prep_dask_kwargs(
+            dict(task_name="FISH!", task_tags=["dask-resource:GPU=1"], task_index=1)
+        )
+        assert kwargs["key"].startswith("FISH!-1-")
+
     def test_submit_sets_task_name(self, mthread):
         with mthread.start():
             fut = mthread.submit(lambda x: x + 1, 1, extra_context={"task_name": "inc"})
@@ -327,6 +339,13 @@ def test_submit_sets_task_name(self, mthread):
             assert fut.key.startswith("inc-")
             assert res == 2
 
+            fut = mthread.submit(
+                lambda x: x + 1, 1, extra_context={"task_name": "inc", "task_index": 1}
+            )
+            (res,) = mthread.wait([fut])
+            assert fut.key.startswith("inc-1-")
+            assert res == 2
+
     @pytest.mark.parametrize("executor", ["mproc", "mthread"], indirect=True)
     def test_is_pickleable(self, executor):
         post = cloudpickle.loads(cloudpickle.dumps(executor))
diff --git a/tests/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
--- a/tests/engine/test_flow_runner.py
+++ b/tests/engine/test_flow_runner.py
@@ -3,6 +3,7 @@
 import queue
 import random
 import time
+from contextlib import contextmanager
 from unittest.mock import MagicMock
 
 import pendulum
@@ -1548,3 +1549,41 @@ def add(x):
 
     ## to be safe, ensure '5' isn't in the logs
     assert len([log.message for log in caplog.records if "99" in log.message]) == 0
+
+
+def test_dask_executor_with_flow_runner_sets_task_keys(mthread):
+    """Integration test that ensures the flow runner forwards the proper
+    information to the DaskExecutor so that key names are set based on
+    the task name"""
+    key_names = set()
+
+    class MyExecutor(Executor):
+        @contextmanager
+        def start(self):
+            with mthread.start():
+                yield
+
+        def submit(self, *args, **kwargs):
+            fut = mthread.submit(*args, **kwargs)
+            key_names.add(fut.key.split("-")[0])
+            return fut
+
+        def wait(self, x):
+            return mthread.wait(x)
+
+    @prefect.task
+    def inc(x):
+        return x + 1
+
+    @prefect.task
+    def do_sum(x):
+        return sum(x)
+
+    with Flow("test") as flow:
+        a = inc(1)
+        b = inc.map(range(3))
+        c = do_sum(b)
+
+    flow.run(executor=MyExecutor())
+
+    assert key_names == {"inc", "do_sum"}