Refactor dashboard code

distributed/client.py

@@ -3861,7 +3861,7 @@ async def _get_task_stream(
             from .diagnostics.task_stream import rectangles
 
             rects = rectangles(msgs)
-            from .dashboard.components import task_stream_figure
+            from .dashboard.components.scheduler import task_stream_figure
 
             source, figure = task_stream_figure(sizing_mode="stretch_both")
             source.data.update(rects)
@@ -4424,15 +4424,15 @@ class get_task_stream(object):
 
     To share this file with others you may wish to upload and serve it online.
     A common way to do this is to upload the file as a gist, and then serve it
-    on https://rawgit.com ::
+    on https://raw.githack.com ::
 
        $ pip install gist
        $ gist task-stream.html
        https://gist.github.com/8a5b3c74b10b413f612bb5e250856ceb
 
     You can then navigate to that site, click the "Raw" button to the right of
     the ``task-stream.html`` file, and then provide that URL to
-    https://rawgit.com .  This process should provide a sharable link that
+    https://raw.githack.com .  This process should provide a sharable link that
     others can use to see your task stream plot.
 
     See Also

distributed/dashboard/components/__init__.py

@@ -0,0 +1,93 @@
+import asyncio
+from bisect import bisect
+from operator import add
+from time import time
+import weakref
+
+from bokeh.layouts import row, column
+from bokeh.models import (
+    ColumnDataSource,
+    Plot,
+    DataRange1d,
+    LinearAxis,
+    HoverTool,
+    BoxZoomTool,
+    ResetTool,
+    PanTool,
+    WheelZoomTool,
+    Range1d,
+    Quad,
+    TapTool,
+    OpenURL,
+    Button,
+    Select,
+)
+from bokeh.palettes import Spectral9
+from bokeh.plotting import figure
+import dask
+from tornado import gen
+import toolz
+
+from distributed.dashboard.utils import without_property_validation, BOKEH_VERSION
+from distributed.diagnostics.progress_stream import nbytes_bar
+from distributed import profile
+from distributed.utils import log_errors, parse_timedelta
+
+if dask.config.get("distributed.dashboard.export-tool"):
+    from distributed.dashboard.export_tool import ExportTool
+else:
+    ExportTool = None
+
+
+profile_interval = dask.config.get("distributed.worker.profile.interval")
+profile_interval = parse_timedelta(profile_interval, default="ms")
+
+
+class DashboardComponent(object):
+    """ Base class for Dask.distributed UI dashboard components.
+
+    This class must have two attributes, ``root`` and ``source``, and one
+    method ``update``:
+
+    *  source: a Bokeh ColumnDataSource
+    *  root: a Bokeh Model
+    *  update: a method that consumes the messages dictionary found in
+               distributed.bokeh.messages
+    """
+
+    def __init__(self):
+        self.source = None
+        self.root = None
+
+    def update(self, messages):
+        """ Reads from bokeh.distributed.messages and updates self.source """
+
+
+def add_periodic_callback(doc, component, interval):
+    """ Add periodic callback to doc in a way that avoids reference cycles
+
+    If we instead use ``doc.add_periodic_callback(component.update, 100)`` then
+    the component stays in memory as a reference cycle because its method is
+    still around.  This way we avoid that and let things clean up a bit more
+    nicely.
+
+    TODO: we still have reference cycles.  Docs seem to be referred to by their
+    add_periodic_callback methods.
+    """
+    ref = weakref.ref(component)
+
+    doc.add_periodic_callback(lambda: update(ref), interval)
+    _attach(doc, component)
+
+
+def update(ref):
+    comp = ref()
+    if comp is not None:
+        comp.update()
+
+
+def _attach(doc, component):
+    if not hasattr(doc, "components"):
+        doc.components = set()
+
+    doc.components.add(component)

distributed/dashboard/nvml.py → distributed/dashboard/components/nvml.py

@@ -1,6 +1,6 @@
 import math
 
-from .components import DashboardComponent, add_periodic_callback
+from distributed.dashboard.components import DashboardComponent, add_periodic_callback
 
 from bokeh.plotting import figure
 from bokeh.models import (
@@ -13,9 +13,17 @@
 )
 from tornado import escape
 from dask.utils import format_bytes
-from ..utils import log_errors
-from .scheduler import update, applications, BOKEH_THEME
-from .utils import without_property_validation
+from distributed.utils import log_errors
+from distributed.dashboard.scheduler import BOKEH_THEME
+from distributed.dashboard.utils import without_property_validation, update
+
+
+try:
+    import pynvml
+
+    pynvml.nvmlInit()
+except Exception:
+    pass
 
 
 class GPUCurrentLoad(DashboardComponent):
@@ -181,14 +189,3 @@ def gpu_utilization_doc(scheduler, extra, doc):
     add_periodic_callback(doc, gpu_load, 100)
     doc.add_root(gpu_load.utilization_figure)
     doc.theme = BOKEH_THEME
-
-
-try:
-    import pynvml
-
-    pynvml.nvmlInit()
-except Exception:
-    pass
-else:
-    applications["/individual-gpu-memory"] = gpu_memory_doc
-    applications["/individual-gpu-utilization"] = gpu_utilization_doc