Repository: dask/dask

Issue:
Color palette lookup error in `visualize`
It appears that `diagnostics.profile_visualize.get_colors` assumes that the color palettes in `bokeh.palettes` have entries for each number of colors from `low` to `high`, but this doesn't appear to be the case. Maybe `bokeh.palettes.small_palettes` is actually what's wanted since that omits the 256-color palettes? Changing `getattr(palettes, palette)` to `palettes.small_palettes[palette]` fixes the issue for me.

Full traceback:
```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-10-4df906118825> in <module>()
----> 1 prof.visualize()

/Users/brettnaul/miniconda3/envs/cesium/lib/python3.5/site-packages/dask/diagnostics/profile.py in visualize(self, **kwargs)
     92         """
     93         from .profile_visualize import visualize
---> 94         return visualize(self, **kwargs)
     95 
     96     def clear(self):

/Users/brettnaul/miniconda3/envs/cesium/lib/python3.5/site-packages/dask/diagnostics/profile_visualize.py in visualize(profilers, file_path, show, save, **kwargs)
    174     if not isinstance(profilers, list):
    175         profilers = [profilers]
--> 176     figs = [prof._plot(**kwargs) for prof in profilers]
    177     # Stack the plots
    178     if len(figs) == 1:

/Users/brettnaul/miniconda3/envs/cesium/lib/python3.5/site-packages/dask/diagnostics/profile_visualize.py in <listcomp>(.0)
    174     if not isinstance(profilers, list):
    175         profilers = [profilers]
--> 176     figs = [prof._plot(**kwargs) for prof in profilers]
    177     # Stack the plots
    178     if len(figs) == 1:

/Users/brettnaul/miniconda3/envs/cesium/lib/python3.5/site-packages/dask/diagnostics/profile.py in _plot(self, **kwargs)
     82     def _plot(self, **kwargs):
     83         from .profile_visualize import plot_tasks
---> 84         return plot_tasks(self.results, self._dsk, **kwargs)
     85 
     86     def visualize(self, **kwargs):

/Users/brettnaul/miniconda3/envs/cesium/lib/python3.5/site-packages/dask/diagnostics/profile_visualize.py in plot_tasks(results, dsk, palette, label_size, **kwargs)
    259         data['y'] = [id_lk[i] + 1 for i in ids]
    260         data['function'] = funcs = [pprint_task(i, dsk, label_size) for i in tasks]
--> 261         data['color'] = get_colors(palette, funcs)
    262         data['key'] = [str(i) for i in keys]
    263 

/Users/brettnaul/miniconda3/envs/cesium/lib/python3.5/site-packages/dask/diagnostics/profile_visualize.py in get_colors(palette, funcs)
    137         colors = palette_lookup[low]
    138     else:
--> 139         colors = palette_lookup[n_funcs]
    140     color_lookup = dict(zip(unique_funcs, colors))
    141     return [color_lookup[n] for n in funcs]

KeyError: 32
```


Discussion:

Resolving patch:
diff --git a/dask/diagnostics/profile_visualize.py b/dask/diagnostics/profile_visualize.py
--- a/dask/diagnostics/profile_visualize.py
+++ b/dask/diagnostics/profile_visualize.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
+import random
+from bisect import bisect_left
 from itertools import cycle
 from operator import itemgetter, add
 
@@ -119,7 +121,8 @@ def get_colors(palette, funcs):
     Parameters
     ----------
     palette : string
-        Name of the palette. Must be an attribute on bokeh.palettes.
+        Name of the bokeh palette to use, must be a member of
+        bokeh.palettes.all_palettes.
     funcs : iterable
         Iterable of function names
     """
@@ -128,16 +131,16 @@ def get_colors(palette, funcs):
 
     unique_funcs = list(sorted(tz.unique(funcs)))
     n_funcs = len(unique_funcs)
-    palette_lookup = getattr(palettes, palette)
-    keys = list(palette_lookup.keys())
-    low, high = min(keys), max(keys)
-    if n_funcs > high:
-        colors = cycle(palette_lookup[high])
-    elif n_funcs < low:
-        colors = palette_lookup[low]
-    else:
-        colors = palette_lookup[n_funcs]
-    color_lookup = dict(zip(unique_funcs, colors))
+    palette_lookup = palettes.all_palettes[palette]
+    keys = list(sorted(palette_lookup.keys()))
+    index = keys[min(bisect_left(keys, n_funcs), len(keys) - 1)]
+    palette = palette_lookup[index]
+    # Some bokeh palettes repeat colors, we want just the unique set
+    palette = list(tz.unique(palette))
+    if len(palette) > n_funcs:
+        # Consistently shuffle palette - prevents just using low-range
+        random.Random(42).shuffle(palette)
+    color_lookup = dict(zip(unique_funcs, cycle(palette)))
     return [color_lookup[n] for n in funcs]
 
 
@@ -216,8 +219,8 @@ def plot_tasks(results, dsk, palette='Viridis', label_size=60, **kwargs):
     dsk : dict
         The dask graph being profiled.
     palette : string, optional
-        Name of the bokeh palette to use, must be an attribute on
-        bokeh.palettes.
+        Name of the bokeh palette to use, must be a member of
+        bokeh.palettes.all_palettes.
     label_size: int (optional)
         Maximum size of output labels in plot, defaults to 60
     **kwargs
@@ -298,8 +301,8 @@ def plot_resources(results, palette='Viridis', **kwargs):
     results : sequence
         Output of ResourceProfiler.results
     palette : string, optional
-        Name of the bokeh palette to use, must be an attribute on
-        bokeh.palettes.
+        Name of the bokeh palette to use, must be a member of
+        bokeh.palettes.all_palettes.
     **kwargs
         Other keyword arguments, passed to bokeh.figure. These will override
         all defaults set by plot_resources.
@@ -325,7 +328,7 @@ def plot_resources(results, palette='Viridis', **kwargs):
     else:
         t = mem = cpu = []
         p = bp.figure(y_range=(0, 100), x_range=(0, 10), **defaults)
-    colors = getattr(palettes, palette)[6]
+    colors = palettes.all_palettes[palette][6]
     p.line(t, cpu, color=colors[0], line_width=4, legend='% CPU')
     p.yaxis.axis_label = "% CPU"
     p.extra_y_ranges = {'memory': Range1d(start=(min(mem) if mem else 0),
@@ -353,8 +356,8 @@ def plot_cache(results, dsk, start_time, metric_name, palette='Viridis',
     metric_name : string
         Metric used to measure cache size
     palette : string, optional
-        Name of the bokeh palette to use, must be an attribute on
-        bokeh.palettes.
+        Name of the bokeh palette to use, must be a member of
+        bokeh.palettes.all_palettes.
     label_size: int (optional)
         Maximum size of output labels in plot, defaults to 60
     **kwargs


Test patch:
diff --git a/dask/diagnostics/tests/test_profiler.py b/dask/diagnostics/tests/test_profiler.py
--- a/dask/diagnostics/tests/test_profiler.py
+++ b/dask/diagnostics/tests/test_profiler.py
@@ -287,17 +287,26 @@ def test_saves_file():
 @pytest.mark.skipif("not bokeh")
 def test_get_colors():
     from dask.diagnostics.profile_visualize import get_colors
-    from bokeh.palettes import Blues9, Blues5, BrBG3
+    from bokeh.palettes import Blues9, Blues5, Viridis
     from itertools import cycle
     funcs = list(range(11))
     cmap = get_colors('Blues', funcs)
     lk = dict(zip(funcs, cycle(Blues9)))
     assert cmap == [lk[i] for i in funcs]
+
     funcs = list(range(5))
     cmap = get_colors('Blues', funcs)
     lk = dict(zip(funcs, Blues5))
     assert cmap == [lk[i] for i in funcs]
+
     funcs = [0, 1, 0, 1, 0, 1]
     cmap = get_colors('BrBG', funcs)
-    lk = dict(zip([0, 1], BrBG3))
-    assert cmap == [lk[i] for i in funcs]
+    assert len(set(cmap)) == 2
+
+    funcs = list(range(100))
+    cmap = get_colors('Viridis', funcs)
+    assert len(set(cmap)) == 100
+
+    funcs = list(range(300))
+    cmap = get_colors('Viridis', funcs)
+    assert len(set(cmap)) == len(set(Viridis[256]))