Add cuDF spilling statistics to RMM/GPU memory plot

continuous_integration/gpuci/build.sh

@@ -56,3 +56,7 @@ conda list --show-channel-urls
 
 rapids-logger "Python py.test for distributed"
 py.test distributed -v -m gpu --runslow --junitxml="$WORKSPACE/junit-distributed.xml"
+
+# cuDF spill stats monitoring must be enabled for this test
+CUDF_SPILL=on CUDF_SPILL_STATS=1 DASK_DISTRIBUTED__DIAGNOSTICS__CUDF=1 \
+    py.test distributed/diagnostics/tests/test_cudf_diagnostics.py -v -m gpu --runslow --junitxml="$WORKSPACE/junit-distributed.xml"

distributed/dashboard/components/rmm.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-import math
-from textwrap import dedent
+from collections.abc import Iterable
+from typing import TypeVar
 
 from bokeh.core.properties import without_property_validation
 from bokeh.models import (
@@ -10,6 +10,7 @@
     HoverTool,
     NumeralTickFormatter,
     OpenURL,
+    Range1d,
     TapTool,
 )
 from bokeh.plotting import figure
@@ -18,12 +19,19 @@
 from dask.utils import format_bytes
 
 from distributed.dashboard.components import DashboardComponent, add_periodic_callback
-from distributed.dashboard.components.scheduler import BOKEH_THEME, TICKS_1024
+from distributed.dashboard.components.scheduler import (
+    BOKEH_THEME,
+    TICKS_1024,
+    XLABEL_ORIENTATION,
+    MemoryColor,
+)
 from distributed.dashboard.utils import update
 from distributed.utils import log_errors
 
+T = TypeVar("T")
+
 
-class RMMMemoryUsage(DashboardComponent):
+class RMMMemoryUsage(DashboardComponent, MemoryColor):
     """
     GPU memory usage plot that includes information about memory
     managed by RMM. If an RMM pool is being used, shows the amount of
@@ -32,168 +40,166 @@
 
     @log_errors
     def __init__(self, scheduler, width=600, **kwargs):
+        DashboardComponent.__init__(self)
+        MemoryColor.__init__(self, neutral_color="#76B900")
+
         self.last = 0
         self.scheduler = scheduler
         self.source = ColumnDataSource(
             {
-                "rmm-used": [1, 2],
-                "rmm-used-half": [0.5, 1],
-                "rmm-total": [2, 4],
-                "rmm-total-half": [1, 2],
-                "external-used": [2, 1],
-                "external-used-x": [3, 4.5],
-                "worker": ["a", "b"],
-                "gpu-index": [0, 0],
-                "y": [1, 2],
-                "escaped_worker": ["a", "b"],
-                "rmm_memory_text": [
-                    "RMM memory used: 1B/1B\nTotal GPU memory used: 1B/2B",
-                    "RMM memory used: 1B/1B\nTotal GPU memory used: 1B/2B",
-                ],
+                "width": [],
+                "x": [],
+                "y": [],
+                "color": [],
+                "alpha": [],
+                "worker": [],
+                "escaped_worker": [],
+                "rmm_used": [],
+                "rmm_total": [],
+                "gpu_used": [],
+                "gpu_total": [],
+                "spilled": [],
             }
         )
 
-        memory = figure(
-            title="RMM Memory",
+        self.root = figure(
+            title="RMM memory used",
             tools="",
             width=int(width / 2),
-            name="rmm_memory_histogram",
+            name="rmm_memory",
             **kwargs,
         )
-
-        rect = memory.rect(
-            source=self.source,
-            x="rmm-used-half",
-            y="y",
-            width="rmm-used",
-            height=1,
-            color="#76B900",
-            alpha=1.0,
-        )
-        rect.nonselection_glyph = None
-
-        rect = memory.rect(
+        rect = self.root.rect(
             source=self.source,
-            x="rmm-total-half",
+            x="x",
             y="y",
-            width="rmm-total",
-            height=1,
-            color="#76B900",
-            alpha=0.75,
+            width="width",
+            height=0.9,
+            color="color",
+            fill_alpha="alpha",
+            line_width=0,
         )
         rect.nonselection_glyph = None
 
-        rect = memory.rect(
-            source=self.source,
-            x="external-used-x",
-            y="y",
-            width="external-used",
-            height=1,
-            color="#76B900",
-            alpha=0.5,
+        self.root.axis[0].ticker = BasicTicker(**TICKS_1024)
+        self.root.xaxis[0].formatter = NumeralTickFormatter(format="0.0 b")
+        self.root.xaxis.major_label_orientation = XLABEL_ORIENTATION
+        self.root.xaxis.minor_tick_line_alpha = 0
+        self.root.x_range = Range1d(start=0)
+        self.root.yaxis.visible = False
+        self.root.ygrid.visible = False
+        self.root.toolbar_location = None
+
+        tap = TapTool(callback=OpenURL(url="./info/worker/@escaped_worker.html"))
+        self.root.add_tools(tap)
+
+        hover = HoverTool(
+            point_policy="follow_mouse",
+            tooltips="""
+            <div>
+                <span style="font-size: 12px; font-weight: bold;">Worker:</span>&nbsp;
+                <span style="font-size: 10px; font-family: Monaco, monospace;">@worker</span>
+            </div>
+            <div>
+                <span style="font-size: 12px; font-weight: bold;">RMM memory used:</span>&nbsp;
+                <span style="font-size: 10px; font-family: Monaco, monospace;">@rmm_used{0.00 b} / @rmm_total{0.00 b}</span>
+            </div>
+            <div>
+                <span style="font-size: 12px; font-weight: bold;">GPU memory used:</span>&nbsp;
+                <span style="font-size: 10px; font-family: Monaco, monospace;">@gpu_used{0.00 b} / @gpu_total{0.00 b}</span>
+            </div>
+            <div>
+                <span style="font-size: 12px; font-weight: bold;">Spilled to CPU:</span>&nbsp;
+                <span style="font-size: 10px; font-family: Monaco, monospace;">@spilled{0.00 b}</span>
+            </div>
+            """,
         )
-        rect.nonselection_glyph = None
-
-        memory.axis[0].ticker = BasicTicker(**TICKS_1024)
-        memory.xaxis[0].formatter = NumeralTickFormatter(format="0.0 b")
-        memory.xaxis.major_label_orientation = -math.pi / 12
-        memory.x_range.start = 0
-
-        for fig in [memory]:
-            fig.xaxis.minor_tick_line_alpha = 0
-            fig.yaxis.visible = False
-            fig.ygrid.visible = False
-
-            tap = TapTool(callback=OpenURL(url="./info/worker/@escaped_worker.html"))
-            fig.add_tools(tap)
-
-            fig.toolbar_location = None
-            fig.yaxis.visible = False
-
-        hover = HoverTool()
-        hover.tooltips = "@worker : @rmm_memory_text"
-        hover.point_policy = "follow_mouse"
-        memory.add_tools(hover)
-
-        self.memory_figure = memory
+        self.root.add_tools(hover)
 
     @without_property_validation
     @log_errors
     def update(self):
+        def quadlist(i: Iterable[T]) -> list[T]:
+            out = []
+            for ii in i:
+                out += [ii, ii, ii, ii]
+            return out
+
         workers = list(self.scheduler.workers.values())
-        rmm_total = []
+
+        width = []
+        x = []
+        color = []
+        max_limit = 0
         rmm_used = []
-        external_used = []
-        gpu_index = []
-        y = []
-        worker = []
-        external_used_x = []
-        memory_max = 0
+        rmm_total = []
+        gpu_used = []
         gpu_total = []
-        rmm_memory_text = []
+        spilled = []
 
-        for idx, ws in enumerate(workers):
+        for ws in workers:
             try:
                 rmm_metrics = ws.metrics["rmm"]
                 gpu_metrics = ws.metrics["gpu"]
                 gpu_info = ws.extra["gpu"]
             except KeyError:
-                continue
-            rmm_total_worker = rmm_metrics["rmm-total"]  # RMM memory only
-            rmm_used_worker = rmm_metrics["rmm-used"]
-            gpu_total_worker = gpu_info["memory-total"]  # All GPU memory
-            gpu_used_worker = gpu_metrics["memory-used"]
+                rmm_metrics = {"rmm-used": 0, "rmm-total": 0}
+                gpu_metrics = {"memory-used": 0}
+                gpu_info = {"memory-total": 0}
+
+            try:
+                cudf_metrics = ws.metrics["cudf"]
+            except KeyError:
+                cudf_metrics = {"cudf-spilled": 0}
 
-            external_used_worker = gpu_used_worker - rmm_total_worker
+            rmm_used_worker = rmm_metrics["rmm-used"]  # RMM memory only
+            rmm_total_worker = rmm_metrics["rmm-total"]
+            gpu_used_worker = gpu_metrics["memory-used"]  # All GPU memory
+            gpu_total_worker = gpu_info["memory-total"]
+            spilled_worker = cudf_metrics["cudf-spilled"] or 0  # memory spilled to host
 
-            rmm_total.append(rmm_total_worker)
+            max_limit = max(
+                max_limit, gpu_total_worker, gpu_used_worker + spilled_worker
+            )
+            color_i = self._memory_color(gpu_used_worker, gpu_total_worker, ws.status)
+
+            width += [
+                rmm_used_worker,
+                rmm_total_worker - rmm_used_worker,
+                gpu_used_worker - rmm_total_worker,
+                spilled_worker,
+            ]
+            x += [sum(width[-4:i]) + width[i] / 2 for i in range(-4, 0)]
+            color += [color_i, color_i, color_i, "grey"]
+
+            # memory info
             rmm_used.append(rmm_used_worker)
+            rmm_total.append(rmm_total_worker)
+            gpu_used.append(gpu_used_worker)
             gpu_total.append(gpu_total_worker)
-            external_used.append(external_used_worker)
-            external_used_x.append(rmm_total_worker + external_used_worker / 2)
-            worker.append(ws.address)
-            gpu_index.append(idx)
-            y.append(idx)
-
-            memory_max = max(memory_max, gpu_total_worker)
-
-            rmm_memory_text.append(
-                "RMM memory used: {}/{}\nTotal GPU memory used: {}/{}".format(
-                    format_bytes(rmm_used_worker),
-                    format_bytes(rmm_total_worker),
-                    format_bytes(gpu_used_worker),
-                    format_bytes(gpu_total_worker),
-                )
-            )
+            spilled.append(spilled_worker)
 
-        self.memory_figure.title.text = dedent(
-            """\
-            RMM Utilization: {} / {}
-            GPU Memory: {} / {}
-            """.format(
-                format_bytes(sum(rmm_used)),
-                format_bytes(sum(rmm_total)),
-                format_bytes(sum([*rmm_total, *external_used])),
-                format_bytes(sum(gpu_total)),
-            )
-        )
+        title = f"RMM memory used: {format_bytes(sum(rmm_used))} / {format_bytes(sum(rmm_total))}\nGPU memory used: {format_bytes(sum(gpu_used))} / {format_bytes(sum(gpu_total))}"
+        if sum(spilled):
+            title += f" + {format_bytes(sum(spilled))} spilled to CPU"
+        self.root.title.text = title
 
         result = {
-            "rmm-total": rmm_total,
-            "rmm-used": rmm_used,
-            "external-used": external_used,
-            "rmm-total-half": [m // 2 for m in rmm_total],
-            "rmm-used-half": [m // 2 for m in rmm_used],
-            "external-used-x": external_used_x,
-            "worker": worker,
-            "gpu-index": gpu_index,
-            "y": y,
-            "escaped_worker": [escape.url_escape(w) for w in worker],
-            "rmm_memory_text": rmm_memory_text,
+            "width": width,
+            "x": x,
+            "y": quadlist(range(len(workers))),
+            "color": color,
+            "alpha": [1, 0.7, 0.4, 1] * len(workers),
+            "worker": quadlist(ws.address for ws in workers),
+            "escaped_worker": quadlist(escape.url_escape(ws.address) for ws in workers),
+            "rmm_used": quadlist(rmm_used),
+            "rmm_total": quadlist(rmm_total),
+            "gpu_used": quadlist(gpu_used),
+            "gpu_total": quadlist(gpu_total),
+            "spilled": quadlist(spilled),
         }
 
-        self.memory_figure.x_range.end = memory_max
-
+        self.root.x_range.end = max_limit
         update(self.source, result)
 
 
@@ -202,5 +208,5 @@
     rmm_load = RMMMemoryUsage(scheduler, sizing_mode="stretch_both")
     rmm_load.update()
     add_periodic_callback(doc, rmm_load, 100)
-    doc.add_root(rmm_load.memory_figure)
+    doc.add_root(rmm_load.root)
     doc.theme = BOKEH_THEME

distributed/dashboard/components/scheduler.py

@@ -276,25 +276,32 @@ class MemoryColor:
     orange: float
     red: float
 
-    def __init__(self):
+    def __init__(
+        self, neutral_color="blue", target_color="orange", terminated_color="red"
+    ):
+        self.neutral_color = neutral_color
+        self.target_color = target_color
+        self.terminated_color = terminated_color
+
         target = dask.config.get("distributed.worker.memory.target")
         spill = dask.config.get("distributed.worker.memory.spill")
         terminate = dask.config.get("distributed.worker.memory.terminate")
+
         # These values can be False. It's also common to configure them to impossibly
         # high values to achieve the same effect.
         self.orange = min(target or math.inf, spill or math.inf)
         self.red = min(terminate or math.inf, 1.0)
 
     def _memory_color(self, current: int, limit: int, status: Status) -> str:
         if status != Status.running:
-            return "red"
+            return self.terminated_color
         if not limit:
-            return "blue"
+            return self.neutral_color
         if current >= limit * self.red:
-            return "red"
+            return self.terminated_color
         if current >= limit * self.orange:
-            return "orange"
-        return "blue"
+            return self.target_color
+        return self.neutral_color
 
 
 class ClusterMemory(DashboardComponent, MemoryColor):