Merge fine metrics classes, all plots update on bytes toggle [skip ci]

Author: milesgranger

distributed/dashboard/components/scheduler.py

@@ -3377,87 +3377,16 @@ def update(self):
         )
 
 
-class _PerfMetricsSendData(DashboardComponent):
-    @log_errors
-    def __init__(self, scheduler, **kwargs):
-        self.scheduler = scheduler
-        self.data = defaultdict(list)
-        self.source = ColumnDataSource(data=dict())
-        self.substantial_change = False
-        self.init_root()
-
-    def init_root(self):
-        self.fig = figure()
-        self.root = row([self.fig])
-
-    @without_property_validation
-    @log_errors
-    def update(self):
-        items = self.scheduler.cumulative_worker_metrics.items()
-        items = (
-            (k, v)
-            for k, v in items
-            if isinstance(k, tuple) and k[0] == "get-data" and k[-1] == "seconds"
-        )
-        for (_type, operation, _freq), val in items:
-            if operation not in self.data["operation"]:
-                self.substantial_change = True
-                self.data["operation"].append(operation)
-
-            idx = self.data["operation"].index(operation)
-            while idx >= len(self.data["value"]):
-                self.data["value"].append(float("NaN"))
-            self.data["value"][idx] = val
-
-        self.data["text"] = [format_time(n) for n in self.data["value"]]
-        self.data["angle"] = [
-            value / sum(self.data["value"]) * 2 * math.pi
-            for value in self.data["value"]
-        ]
-        self.data["color"] = small_palettes["YlGnBu"].get(
-            len(self.data["operation"]), []
-        )
-
-        if self.data:
-            self.source.data = dict(self.data)
-            fig = figure(
-                height=500,
-                title="Send data, by activity",
-                tools="hover",
-                tooltips="@{operation}: @text",
-                x_range=(-0.5, 1.0),
-            )
-            fig.wedge(
-                x=0,
-                y=1,
-                radius=0.4,
-                start_angle=cumsum("angle", include_zero=True),
-                end_angle=cumsum("angle"),
-                line_color="white",
-                fill_color="color",
-                legend_field="operation",
-                source=self.source,
-            )
-            fig.axis.axis_label = None
-            fig.axis.visible = False
-            fig.grid.grid_line_color = None
-
-            if self.substantial_change:
-                self.root.children[0] = fig
-                self.substantial_change = False
-            else:
-                self.fig = fig
-
-
-class _PerfMetricsExecutionByPrefixAndActivity(DashboardComponent):
+class FinePerformanceMetrics(DashboardComponent):
     """
-    Stacked bar-chart displaying breakdown of function execution times
-    between de/serialization, thread-non/cpu, etc.
+    The main overview of the Fine Performance Metrics page.
     """
 
     @log_errors
     def __init__(self, scheduler, **kwargs):
         self.scheduler = scheduler
+        self.senddata = defaultdict(list)
+        self.sendsource = ColumnDataSource(data=dict())
         self.data = defaultdict(list)
         self.source = ColumnDataSource(data=dict())
         self.piechart_source = ColumnDataSource(data=dict())
@@ -3480,10 +3409,12 @@ def handle_toggle(_toggle):
         self.toggle.on_click(handle_toggle)
         self.barchart = figure()
         self.piechart = figure()
+        self.senddata_piechart = figure()
         self.root = column(
             self.function_selector,
             self.toggle,
             row([self.piechart, self.barchart]),
+            row([self.senddata_piechart]),
             sizing_mode="scale_width",
         )
 
@@ -3493,35 +3424,53 @@ def update(self):
         items = (
             (k, v)
             for k, v in self.scheduler.cumulative_worker_metrics.items()
-            if isinstance(k, tuple)
-            and k[0] == "execute"
-            and k[-1] in ("bytes", "seconds")
-        )
-        for (_type, function_name, operation, freq), val in items:
-            if operation not in self.operations:
-                self.substantial_change = True
-                self.operations.append(operation)
-
-            if function_name not in self.data["functions"]:
-                self.substantial_change = True
-                self.function_selector.options.append(function_name)
-                self.data["functions"].append(function_name)
-                self.data["timestamp"].append(datetime.utcnow())
-            idx = self.data["functions"].index(function_name)
-
-            # Some function/operation combos missing, so need to keep columns aligned
-            for op in self.operations:
-                while len(self.data[f"{op}_value"]) != len(self.data["functions"]):
-                    self.data[f"{op}_value"].append(0)
-                    self.data[f"{op}_bytes"].append(0)
-                    self.data[f"{op}_text"].append("")
-
-            if freq == "seconds":
-                self.data[f"{operation}_text"][idx] = format_time(val)
-                self.data[f"{operation}_value"][idx] = val
-            elif freq == "bytes":
-                self.data[f"{operation}_text"][idx] = format_bytes(val)
-                self.data[f"{operation}_bytes"][idx] = val
+            if isinstance(k, tuple) and k[-1] in ("bytes", "seconds")
+        )
+        for (type, *parts), val in items:
+            if type == "get-data":
+                operation, freq = parts
+                if operation not in self.senddata["operation"]:
+                    self.substantial_change = True
+                    self.senddata["operation"].append(operation)
+
+                idx = self.senddata["operation"].index(operation)
+                while idx >= len(self.senddata[f"{operation}_value"]):
+                    self.senddata[f"{operation}_bytes"].append(0)
+                    self.senddata[f"{operation}_value"].append(0)
+                    self.senddata[f"{operation}_text"].append(0)
+                if freq == "bytes":
+                    self.senddata[f"{operation}_text"][idx] = format_bytes(val)
+                    self.senddata[f"{operation}_bytes"][idx] = val
+                elif freq == "seconds":
+                    self.senddata[f"{operation}_text"][idx] = format_time(val)
+                    self.senddata[f"{operation}_value"][idx] = val
+
+            elif type == "execute":
+                function_name, operation, freq = parts
+                if operation not in self.operations:
+                    self.substantial_change = True
+                    self.operations.append(operation)
+
+                if function_name not in self.data["functions"]:
+                    self.substantial_change = True
+                    self.function_selector.options.append(function_name)
+                    self.data["functions"].append(function_name)
+                    self.data["timestamp"].append(datetime.utcnow())
+                idx = self.data["functions"].index(function_name)
+
+                # Some function/operation combos missing, so need to keep columns aligned
+                for op in self.operations:
+                    while len(self.data[f"{op}_value"]) != len(self.data["functions"]):
+                        self.data[f"{op}_value"].append(0)
+                        self.data[f"{op}_bytes"].append(0)
+                        self.data[f"{op}_text"].append("")
+
+                if freq == "seconds":
+                    self.data[f"{operation}_text"][idx] = format_time(val)
+                    self.data[f"{operation}_value"][idx] = val
+                elif freq == "bytes":
+                    self.data[f"{operation}_text"][idx] = format_bytes(val)
+                    self.data[f"{operation}_bytes"][idx] = val
 
         data = self.data.copy()
 
@@ -3623,6 +3572,52 @@ def update(self):
             self.substantial_change = True
         self.source.data = dict(data)
 
+        senddata = dict()
+        senddata["operation"] = self.senddata["operation"]
+        senddata["value"] = [
+            (sum(self.senddata[f"{op}_{'bytes' if show_bytes else 'value'}"]))
+            for op in self.senddata["operation"]
+        ]
+        senddata["text"] = [
+            format_bytes(n) if show_bytes else format_time(n) for n in senddata["value"]
+        ]
+        senddata["angle"] = [
+            (
+                (
+                    sum(self.senddata[f"{op}_{'bytes' if show_bytes else 'value'}"])
+                    / sum(senddata["value"])
+                )
+                if sum(senddata["value"])
+                else 0.0
+            )
+            * 2
+            * math.pi
+            for op in senddata["operation"]
+        ]
+        senddata["color"] = small_palettes["YlGnBu"].get(len(senddata["operation"]), [])
+        self.sendsource.data = senddata
+        senddata_piechart = figure(
+            height=500,
+            title="Send data, by activity",
+            tools="hover",
+            tooltips="@{operation}: @text",
+            x_range=(-0.5, 1.0),
+        )
+        senddata_piechart_renderers = senddata_piechart.wedge(
+            x=0,
+            y=1,
+            radius=0.4,
+            start_angle=cumsum("angle", include_zero=True),
+            end_angle=cumsum("angle"),
+            line_color="white",
+            fill_color="color",
+            legend_field="operation",
+            source=self.sendsource,
+        )
+        senddata_piechart.axis.axis_label = None
+        senddata_piechart.axis.visible = False
+        senddata_piechart.grid.grid_line_color = None
+
         # Show total number of functions to choose from
         self.function_selector.title = (
             f"Filter by function ({len(self.function_selector.options)}):"
@@ -3631,39 +3626,16 @@ def update(self):
         # replacing the child causes small blips if done every iteration vs updating renderers
         # but it's needed when new functions and/or operations show up to rerender plot
         if self.substantial_change:
-            self.root.children[-1].children[0] = barchart
-            self.root.children[-1].children[1] = piechart
+            self.root.children[-2].children[0] = barchart
+            self.root.children[-2].children[1] = piechart
+            if self.senddata:
+                self.root.children[-1].children[0] = senddata_piechart
             self.substantial_change = False
         else:
             self.barchart.renderers = renderers
             self.piechart.renderers = [piechart_renderers]
-
-
-class FinePerformanceMetrics(DashboardComponent):
-    """
-    The main overview of the Fine Performance Metrics page.
-    """
-
-    @log_errors
-    def __init__(self, scheduler, **kwargs):
-        self.scheduler = scheduler
-
-        self.metrics_by_prefix_and_activity = _PerfMetricsExecutionByPrefixAndActivity(
-            scheduler, **kwargs
-        )
-        self.metrics_send_data = _PerfMetricsSendData(scheduler, **kwargs)
-        self.root = column(
-            [
-                self.metrics_by_prefix_and_activity.root,
-                self.metrics_send_data.root,
-            ]
-        )
-
-    @without_property_validation
-    @log_errors
-    def update(self):
-        self.metrics_by_prefix_and_activity.update()
-        self.metrics_send_data.update()
+            if self.senddata:
+                self.senddata_piechart.renderers = [senddata_piechart_renderers]
 
 
 class Contention(DashboardComponent):