Shuffle Service with Scheduler Logic

.github/workflows/tests.yaml

@@ -75,6 +75,10 @@ jobs:
         if: ${{ matrix.os == 'windows-latest' && matrix.python-version == '3.9' }}
         run: mamba uninstall ipython
 
+      - name: Install dask branch
+        shell: bash -l {0}
+        run: python -m pip install --no-deps git+https://github.com/mrocklin/dask@p2p-shuffle
+
       - name: Install
         shell: bash -l {0}
         run: |

.pre-commit-config.yaml

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

continuous_integration/environment-3.10.yaml

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

continuous_integration/environment-3.8.yaml

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

continuous_integration/environment-3.9.yaml

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

distributed/dashboard/components/scheduler.py

@@ -1043,7 +1043,7 @@ class SystemTimeseries(DashboardComponent):
     from ws.metrics["val"] for ws in scheduler.workers.values() divided by nuber of workers.
     """
 
-    def __init__(self, scheduler, **kwargs):
+    def __init__(self, scheduler, follow_interval=20000, **kwargs):
         with log_errors():
             self.scheduler = scheduler
             self.source = ColumnDataSource(
@@ -1060,7 +1060,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(
@@ -3487,6 +3489,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)
@@ -3531,6 +3788,41 @@ def systemmonitor_doc(scheduler, extra, doc):
         doc.theme = BOKEH_THEME
 
 
+def shuffling_doc(scheduler, extra, doc):
+    with log_errors():
+        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
+
+
 def stealing_doc(scheduler, extra, doc):
     with log_errors():
         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,