Shuffle Service with Scheduler Logic (#6007)

.pre-commit-config.yaml

@@ -49,3 +49,4 @@ repos:
           - dask
           - tornado
           - zict
+          - pyarrow

continuous_integration/environment-3.10.yaml

@@ -24,6 +24,7 @@ dependencies:
   - pre-commit
   - prometheus_client
   - psutil
+  - pyarrow=7
   - pytest
   - pytest-cov
   - pytest-faulthandler

continuous_integration/environment-3.8.yaml

@@ -25,6 +25,7 @@ dependencies:
   - pre-commit
   - prometheus_client
   - psutil
+  - pyarrow=7
   - pynvml  # Only tested here
   - pytest
   - pytest-cov

continuous_integration/environment-3.9.yaml

@@ -26,6 +26,8 @@ dependencies:
   - pre-commit
   - prometheus_client
   - psutil
+  - pyarrow=7
+  - pynvml  # Only tested here
   - pytest
   - pytest-cov
   - pytest-faulthandler

distributed/dashboard/components/scheduler.py

@@ -1032,7 +1032,7 @@ class SystemTimeseries(DashboardComponent):
     """
 
     @log_errors
-    def __init__(self, scheduler, **kwargs):
+    def __init__(self, scheduler, follow_interval=20000, **kwargs):
         self.scheduler = scheduler
         self.source = ColumnDataSource(
             {
@@ -1048,7 +1048,9 @@ def __init__(self, scheduler, **kwargs):
 
         update(self.source, self.get_data())
 
-        x_range = DataRange1d(follow="end", follow_interval=20000, range_padding=0)
+        x_range = DataRange1d(
+            follow="end", follow_interval=follow_interval, range_padding=0
+        )
         tools = "reset, xpan, xwheel_zoom"
 
         self.bandwidth = figure(
@@ -3465,6 +3467,261 @@ def update(self):
         self.source.data.update(data)
 
 
+class Shuffling(DashboardComponent):
+    """Occupancy (in time) per worker"""
+
+    def __init__(self, scheduler, **kwargs):
+        with log_errors():
+            self.scheduler = scheduler
+            self.source = ColumnDataSource(
+                {
+                    "worker": [],
+                    "y": [],
+                    "comm_memory": [],
+                    "comm_memory_limit": [],
+                    "comm_buckets": [],
+                    "comm_active": [],
+                    "comm_avg_duration": [],
+                    "comm_avg_size": [],
+                    "comm_read": [],
+                    "comm_written": [],
+                    "comm_color": [],
+                    "disk_memory": [],
+                    "disk_memory_limit": [],
+                    "disk_buckets": [],
+                    "disk_active": [],
+                    "disk_avg_duration": [],
+                    "disk_avg_size": [],
+                    "disk_read": [],
+                    "disk_written": [],
+                    "disk_color": [],
+                }
+            )
+            self.totals_source = ColumnDataSource(
+                {
+                    "x": ["Network Send", "Network Receive", "Disk Write", "Disk Read"],
+                    "values": [0, 0, 0, 0],
+                }
+            )
+
+            self.comm_memory = figure(
+                title="Comms Buffer",
+                tools="",
+                toolbar_location="above",
+                x_range=Range1d(0, 100_000_000),
+                **kwargs,
+            )
+            self.comm_memory.hbar(
+                source=self.source,
+                right="comm_memory",
+                y="y",
+                height=0.9,
+                color="comm_color",
+            )
+            hover = HoverTool(
+                tooltips=[
+                    ("Memory Used", "@comm_memory{0.00 b}"),
+                    ("Average Write", "@comm_avg_size{0.00 b}"),
+                    ("# Buckets", "@comm_buckets"),
+                    ("Average Duration", "@comm_avg_duration"),
+                ],
+                formatters={"@comm_avg_duration": "datetime"},
+                mode="hline",
+            )
+            self.comm_memory.add_tools(hover)
+            self.comm_memory.x_range.start = 0
+            self.comm_memory.x_range.end = 1
+            self.comm_memory.xaxis[0].formatter = NumeralTickFormatter(format="0.0 b")
+
+            self.disk_memory = figure(
+                title="Disk Buffer",
+                tools="",
+                toolbar_location="above",
+                x_range=Range1d(0, 100_000_000),
+                **kwargs,
+            )
+            self.disk_memory.yaxis.visible = False
+
+            self.disk_memory.hbar(
+                source=self.source,
+                right="disk_memory",
+                y="y",
+                height=0.9,
+                color="disk_color",
+            )
+
+            hover = HoverTool(
+                tooltips=[
+                    ("Memory Used", "@disk_memory{0.00 b}"),
+                    ("Average Write", "@disk_avg_size{0.00 b}"),
+                    ("# Buckets", "@disk_buckets"),
+                    ("Average Duration", "@disk_avg_duration"),
+                ],
+                formatters={"@disk_avg_duration": "datetime"},
+                mode="hline",
+            )
+            self.disk_memory.add_tools(hover)
+            self.disk_memory.xaxis[0].formatter = NumeralTickFormatter(format="0.0 b")
+
+            self.totals = figure(
+                title="Total movement",
+                tools="",
+                toolbar_location="above",
+                **kwargs,
+            )
+            titles = ["Network Send", "Network Receive", "Disk Write", "Disk Read"]
+            self.totals = figure(
+                x_range=titles,
+                title="Totals",
+                toolbar_location=None,
+                tools="",
+                **kwargs,
+            )
+
+            self.totals.vbar(
+                x="x",
+                top="values",
+                width=0.9,
+                source=self.totals_source,
+            )
+
+            self.totals.xgrid.grid_line_color = None
+            self.totals.y_range.start = 0
+            self.totals.yaxis[0].formatter = NumeralTickFormatter(format="0.0 b")
+
+            hover = HoverTool(
+                tooltips=[("Total", "@values{0.00b}")],
+                mode="vline",
+            )
+            self.totals.add_tools(hover)
+
+            self.root = row(self.comm_memory, self.disk_memory)
+
+    @without_property_validation
+    def update(self):
+        with log_errors():
+            input = self.scheduler.extensions["shuffle"].heartbeats
+            if not input:
+                return
+
+            input = list(input.values())[-1]  # TODO: multiple concurrent shuffles
+
+            data = {
+                "worker": [],
+                "y": [],
+                "comm_memory": [],
+                "comm_memory_limit": [],
+                "comm_buckets": [],
+                "comm_active": [],
+                "comm_avg_duration": [],
+                "comm_avg_size": [],
+                "comm_read": [],
+                "comm_written": [],
+                "comm_color": [],
+                "disk_memory": [],
+                "disk_memory_limit": [],
+                "disk_buckets": [],
+                "disk_active": [],
+                "disk_avg_duration": [],
+                "disk_avg_size": [],
+                "disk_read": [],
+                "disk_written": [],
+                "disk_color": [],
+            }
+            now = time()
+
+            for i, (worker, d) in enumerate(input.items()):
+                data["y"].append(i)
+                data["worker"].append(worker)
+                data["comm_memory"].append(d["comms"]["memory"])
+                data["comm_memory_limit"].append(d["comms"]["memory_limit"])
+                data["comm_buckets"].append(d["comms"]["buckets"])
+                data["comm_active"].append(d["comms"]["active"])
+                data["comm_avg_duration"].append(
+                    d["comms"]["diagnostics"].get("avg_duration", 0)
+                )
+                data["comm_avg_size"].append(
+                    d["comms"]["diagnostics"].get("avg_size", 0)
+                )
+                data["comm_read"].append(d["comms"]["read"])
+                data["comm_written"].append(d["comms"]["written"])
+                try:
+                    if self.scheduler.workers[worker].last_seen < now - 5:
+                        data["comm_color"].append("gray")
+                    elif d["comms"]["active"]:
+                        data["comm_color"].append("green")
+                    elif d["comms"]["memory"] > d["comms"]["memory_limit"]:
+                        data["comm_color"].append("red")
+                    else:
+                        data["comm_color"].append("blue")
+                except KeyError:
+                    data["comm_color"].append("black")
+
+                data["disk_memory"].append(d["disk"]["memory"])
+                data["disk_memory_limit"].append(d["disk"]["memory_limit"])
+                data["disk_buckets"].append(d["disk"]["buckets"])
+                data["disk_active"].append(d["disk"]["active"])
+                data["disk_avg_duration"].append(
+                    d["disk"]["diagnostics"].get("avg_duration", 0)
+                )
+                data["disk_avg_size"].append(
+                    d["disk"]["diagnostics"].get("avg_size", 0)
+                )
+                data["disk_read"].append(d["disk"]["read"])
+                data["disk_written"].append(d["disk"]["written"])
+                try:
+                    if self.scheduler.workers[worker].last_seen < now - 5:
+                        data["disk_color"].append("gray")
+                    elif d["disk"]["active"]:
+                        data["disk_color"].append("green")
+                    elif d["disk"]["memory"] > d["disk"]["memory_limit"]:
+                        data["disk_color"].append("red")
+                    else:
+                        data["disk_color"].append("blue")
+                except KeyError:
+                    data["disk_color"].append("black")
+
+            """
+            singletons = {
+                "comm_avg_duration": [
+                    sum(data["comm_avg_duration"]) / len(data["comm_avg_duration"])
+                ],
+                "comm_avg_size": [
+                    sum(data["comm_avg_size"]) / len(data["comm_avg_size"])
+                ],
+                "disk_avg_duration": [
+                    sum(data["disk_avg_duration"]) / len(data["disk_avg_duration"])
+                ],
+                "disk_avg_size": [
+                    sum(data["disk_avg_size"]) / len(data["disk_avg_size"])
+                ],
+            }
+            singletons["comm_avg_bandwidth"] = [
+                singletons["comm_avg_size"][0] / singletons["comm_avg_duration"][0]
+            ]
+            singletons["disk_avg_bandwidth"] = [
+                singletons["disk_avg_size"][0] / singletons["disk_avg_duration"][0]
+            ]
+            singletons["y"] = [data["y"][-1] / 2]
+            """
+
+            totals = {
+                "x": ["Network Send", "Network Receive", "Disk Write", "Disk Read"],
+                "values": [
+                    sum(data["comm_written"]),
+                    sum(data["comm_read"]),
+                    sum(data["disk_written"]),
+                    sum(data["disk_read"]),
+                ],
+            }
+            update(self.totals_source, totals)
+
+            update(self.source, data)
+            limit = max(data["comm_memory_limit"] + data["disk_memory_limit"]) * 1.2
+            self.comm_memory.x_range.end = limit
+            self.disk_memory.x_range.end = limit
+
+
 class SchedulerLogs:
     def __init__(self, scheduler, start=None):
         logs = scheduler.get_logs(start=start, timestamps=True)
@@ -3509,6 +3766,41 @@ def systemmonitor_doc(scheduler, extra, doc):
     doc.theme = BOKEH_THEME
 
 
+@log_errors
+def shuffling_doc(scheduler, extra, doc):
+    doc.title = "Dask: Shuffling"
+
+    shuffling = Shuffling(scheduler, width=400, height=400)
+    workers_memory = WorkersMemory(scheduler, width=400, height=400)
+    timeseries = SystemTimeseries(
+        scheduler, width=1600, height=200, follow_interval=3000
+    )
+    event_loop = EventLoop(scheduler, width=200, height=400)
+
+    add_periodic_callback(doc, shuffling, 200)
+    add_periodic_callback(doc, workers_memory, 200)
+    add_periodic_callback(doc, timeseries, 500)
+    add_periodic_callback(doc, event_loop, 500)
+
+    timeseries.bandwidth.y_range = timeseries.disk.y_range
+
+    doc.add_root(
+        column(
+            row(
+                workers_memory.root,
+                shuffling.comm_memory,
+                shuffling.disk_memory,
+                shuffling.totals,
+                event_loop.root,
+            ),
+            row(column(timeseries.bandwidth, timeseries.disk)),
+        )
+    )
+    doc.template = env.get_template("simple.html")
+    doc.template_variables.update(extra)
+    doc.theme = BOKEH_THEME
+
+
 @log_errors
 def stealing_doc(scheduler, extra, doc):
     occupancy = Occupancy(scheduler)

distributed/dashboard/scheduler.py

@@ -37,6 +37,7 @@
     individual_profile_server_doc,
     profile_doc,
     profile_server_doc,
+    shuffling_doc,
     status_doc,
     stealing_doc,
     systemmonitor_doc,
@@ -49,6 +50,7 @@
 
 applications = {
     "/system": systemmonitor_doc,
+    "/shuffle": shuffling_doc,
     "/stealing": stealing_doc,
     "/workers": workers_doc,
     "/events": events_doc,

distributed/dashboard/tests/test_scheduler_bokeh.py

@@ -30,6 +30,7 @@
     Occupancy,
     ProcessingHistogram,
     ProfileServer,
+    Shuffling,
     StealingEvents,
     StealingTimeSeries,
     SystemMonitor,
@@ -1007,6 +1008,21 @@ async def test_prefix_bokeh(s, a, b):
     assert bokeh_app.prefix == f"/{prefix}"
 
 
+@gen_cluster(client=True, worker_kwargs={"dashboard": True})
+async def test_shuffling(c, s, a, b):
+    dd = pytest.importorskip("dask.dataframe")
+    ss = Shuffling(s)
+
+    df = dask.datasets.timeseries()
+    df2 = dd.shuffle.shuffle(df, "x", shuffle="p2p").persist()
+
+    start = time()
+    while not ss.source.data["disk_read"]:
+        ss.update()
+        await asyncio.sleep(0.1)
+        assert time() < start + 5
+
+
 @gen_cluster(client=True, nthreads=[], scheduler_kwargs={"dashboard": True})
 async def test_hardware(c, s):
     plot = Hardware(s)