Repository: PrefectHQ/prefect

Issue:
BokehRunner not compatible with maps
This is a tricky one; I think ideally, since the `BokehRunner` isn't a production tool anyway, we can introspect the flow state and expand out any `mapped` tasks into the appropriate number of individual tasks for the visualization.


Discussion:
IIRC, BokehRunner first runs the flow anyway, then generates it's walkthrough. So that seems totally fine to me. 

Resolving patch:
diff --git a/src/prefect/utilities/bokeh_runner.py b/src/prefect/utilities/bokeh_runner.py
--- a/src/prefect/utilities/bokeh_runner.py
+++ b/src/prefect/utilities/bokeh_runner.py
@@ -3,6 +3,7 @@
 import os
 import subprocess
 import tempfile
+from collections import defaultdict
 from typing import Any, Dict, Iterable, Iterator, List, Union
 
 import cloudpickle
@@ -56,9 +57,11 @@ def run(
             context=context,
             task_contexts=task_contexts,
         )
+        self.reset_flow(self.flow_state)
         self.title = title or "Prefect Flow Interactive Demonstration: {}".format(
             self.flow.name
         )
+
         if viz:
             with tempfile.NamedTemporaryFile() as tmp:
                 cloudpickle.dump(self, tmp)
@@ -71,6 +74,55 @@ def run(
                 )
         return self.flow_state
 
+    def reset_flow(self, state):
+        map_counts = defaultdict(lambda: 0)
+        mapped_tasks = {}
+        edges = []
+
+        for task in self.flow.sorted_tasks():
+            old_edges = self.flow.edges_to(task)
+            mapped_edges = [e for e in old_edges if e.mapped]
+            unmapped_edges = [e for e in old_edges if not e.mapped]
+
+            # update unmapped edges
+            for edge in unmapped_edges:
+                if edge.upstream_task in mapped_tasks:
+                    new_edges = [
+                        prefect.core.edge.Edge(upstream_task=t, downstream_task=task)
+                        for t in mapped_tasks[edge.upstream_task]
+                    ]
+                    edges.extend(new_edges)
+                else:
+                    edges.append(edge)
+
+            # update mapped edges
+            for edge in mapped_edges:
+                # create copies corresponding to each mapped result
+                new_tasks = [task.copy() for _ in state.result[task]]
+                mapped_tasks[task] = new_tasks
+                for t in new_tasks:
+                    t.name = task.name + "<mapped-{}>".format(map_counts[task.name])
+                    map_counts[task.name] += 1
+
+                if edge.upstream_task in mapped_tasks:
+                    upstream = [t for t in mapped_tasks[edge.upstream_task]]
+                else:
+                    upstream = [edge.upstream_task for _ in range(len(new_tasks))]
+
+                new_edges = [
+                    prefect.core.edge.Edge(
+                        upstream_task=up, downstream_task=task, key=edge.key
+                    )
+                    for up, task in zip(upstream, new_tasks)
+                ]
+                edges.extend(new_edges)
+                states = state.result.pop(task)
+                for t, s in zip(new_tasks, states):
+                    state.result[t] = s
+
+        self.flow = prefect.Flow(edges=edges, name=self.flow.name)
+        self.flow_state = state
+
     def compute_depths(self):
         flow = self.flow
         depths = {task: 0 for task in flow.tasks}


Test patch: