Withhold root tasks [no co assignment]

distributed/collections.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import heapq
+import itertools
 import weakref
 from collections import OrderedDict, UserDict
 from collections.abc import Callable, Hashable, Iterator
@@ -72,6 +73,9 @@ def _unpickle(
     def __contains__(self, value: object) -> bool:
         return value in self._data
 
+    def __bool__(self) -> bool:
+        return bool(self._data)
+
     def __len__(self) -> int:
         return len(self._data)
 
@@ -109,6 +113,20 @@ def pop(self) -> T:
                 self._data.discard(value)
                 return value
 
+    def topk(self, k: int) -> Iterator[T]:
+        "Iterator over the largest K elements. This is O(1) for k == 1, O(n*logn) otherwise."
+        if k <= 0:
+            return  # empty iterator
+        if k == 1:
+            yield self.peek()
+        else:
+            # NOTE: we could pop K items off the queue, then push them back.
+            # But copying the list K times is probably slower than just sorting it
+            # with fast C code.
+            # If we had a `heappop` that sliced the list instead of popping from it,
+            # we could implement an optimized version for small `k`s.
+            yield from itertools.islice(self.sorted(), k)
+
     def peekright(self) -> T:
         """Return one of the largest elements (not necessarily the largest!) without
         removing it. It's guaranteed that ``self.peekright() >= self.peek()``.
@@ -145,7 +163,8 @@ def sorted(self) -> Iterator[T]:
         elements in order, from smallest to largest according to the key and insertion
         order.
         """
-        for _, _, vref in sorted(self._heap):
+        self._heap.sort()  # A sorted list maintains the heap invariant
+        for _, _, vref in self._heap:
             value = vref()
             if value in self._data:
                 yield value

distributed/dashboard/components/scheduler.py

@@ -2097,8 +2097,8 @@ def __init__(self, scheduler, **kwargs):
 
         node_colors = factor_cmap(
             "state",
-            factors=["waiting", "processing", "memory", "released", "erred"],
-            palette=["gray", "green", "red", "blue", "black"],
+            factors=["waiting", "queued", "processing", "memory", "released", "erred"],
+            palette=["gray", "yellow", "green", "red", "blue", "black"],
         )
 
         self.root = figure(title="Task Graph", **kwargs)
@@ -2986,7 +2986,7 @@ def __init__(self, scheduler, **kwargs):
         self.scheduler = scheduler
 
         data = progress_quads(
-            dict(all={}, memory={}, erred={}, released={}, processing={})
+            dict(all={}, memory={}, erred={}, released={}, processing={}, queued={})
         )
         self.source = ColumnDataSource(data=data)
 
@@ -3058,6 +3058,18 @@ def __init__(self, scheduler, **kwargs):
             fill_alpha=0.35,
             line_alpha=0,
         )
+        self.root.quad(
+            source=self.source,
+            top="top",
+            bottom="bottom",
+            left="processing-loc",
+            right="queued-loc",
+            fill_color="gray",
+            hatch_pattern="/",
+            hatch_color="white",
+            fill_alpha=0.35,
+            line_alpha=0,
+        )
         self.root.text(
             source=self.source,
             text="show-name",
@@ -3093,6 +3105,14 @@ def __init__(self, scheduler, **kwargs):
                     <span style="font-size: 14px; font-weight: bold;">All:</span>&nbsp;
                     <span style="font-size: 10px; font-family: Monaco, monospace;">@all</span>
                 </div>
+                <div>
+                    <span style="font-size: 14px; font-weight: bold;">Queued:</span>&nbsp;
+                    <span style="font-size: 10px; font-family: Monaco, monospace;">@queued</span>
+                </div>
+                <div>
+                    <span style="font-size: 14px; font-weight: bold;">Processing:</span>&nbsp;
+                    <span style="font-size: 10px; font-family: Monaco, monospace;">@processing</span>
+                </div>
                 <div>
                     <span style="font-size: 14px; font-weight: bold;">Memory:</span>&nbsp;
                     <span style="font-size: 10px; font-family: Monaco, monospace;">@memory</span>
@@ -3101,10 +3121,6 @@ def __init__(self, scheduler, **kwargs):
                     <span style="font-size: 14px; font-weight: bold;">Erred:</span>&nbsp;
                     <span style="font-size: 10px; font-family: Monaco, monospace;">@erred</span>
                 </div>
-                <div>
-                    <span style="font-size: 14px; font-weight: bold;">Ready:</span>&nbsp;
-                    <span style="font-size: 10px; font-family: Monaco, monospace;">@processing</span>
-                </div>
                 """,
         )
         self.root.add_tools(hover)
@@ -3118,6 +3134,7 @@ def update(self):
             "released": {},
             "processing": {},
             "waiting": {},
+            "queued": {},
         }
 
         for tp in self.scheduler.task_prefixes.values():
@@ -3128,6 +3145,7 @@ def update(self):
                 state["released"][tp.name] = active_states["released"]
                 state["processing"][tp.name] = active_states["processing"]
                 state["waiting"][tp.name] = active_states["waiting"]
+                state["queued"][tp.name] = active_states["queued"]
 
         state["all"] = {k: sum(v[k] for v in state.values()) for k in state["memory"]}
 
@@ -3140,16 +3158,18 @@ def update(self):
 
         totals = {
             k: sum(state[k].values())
-            for k in ["all", "memory", "erred", "released", "waiting"]
+            for k in ["all", "memory", "erred", "released", "waiting", "queued"]
         }
         totals["processing"] = totals["all"] - sum(
             v for k, v in totals.items() if k != "all"
         )
 
         self.root.title.text = (
             "Progress -- total: %(all)s, "
-            "in-memory: %(memory)s, processing: %(processing)s, "
             "waiting: %(waiting)s, "
+            "queued: %(queued)s, "
+            "processing: %(processing)s, "
+            "in-memory: %(memory)s, "
             "erred: %(erred)s" % totals
         )
 

distributed/diagnostics/progress_stream.py

@@ -17,7 +17,7 @@ def counts(scheduler, allprogress):
         {"all": valmap(len, allprogress.all), "nbytes": allprogress.nbytes},
         {
             state: valmap(len, allprogress.state[state])
-            for state in ["memory", "erred", "released", "processing"]
+            for state in ["memory", "erred", "released", "processing", "queued"]
         },
     )
 
@@ -66,23 +66,29 @@ def progress_quads(msg, nrows=8, ncols=3):
     ...        'memory': {'inc': 2, 'dec': 0, 'add': 1},
     ...        'erred': {'inc': 0, 'dec': 1, 'add': 0},
     ...        'released': {'inc': 1, 'dec': 0, 'add': 1},
-    ...        'processing': {'inc': 1, 'dec': 0, 'add': 2}}
+    ...        'processing': {'inc': 1, 'dec': 0, 'add': 2},
+    ...        'queued': {'inc': 1, 'dec': 0, 'add': 2}}
 
     >>> progress_quads(msg, nrows=2)  # doctest: +SKIP
-    {'name': ['inc', 'add', 'dec'],
-     'left': [0, 0, 1],
-     'right': [0.9, 0.9, 1.9],
-     'top': [0, -1, 0],
-     'bottom': [-.8, -1.8, -.8],
-     'released': [1, 1, 0],
-     'memory': [2, 1, 0],
-     'erred': [0, 0, 1],
-     'processing': [1, 0, 2],
-     'done': ['3 / 5', '2 / 4', '1 / 1'],
-     'released-loc': [.2/.9, .25 / 0.9, 1],
-     'memory-loc': [3 / 5 / .9, .5 / 0.9, 1],
-     'erred-loc': [3 / 5 / .9, .5 / 0.9, 1.9],
-     'processing-loc': [4 / 5, 1 / 1, 1]}}
+    {'all': [5, 4, 1],
+    'memory': [2, 1, 0],
+    'erred': [0, 0, 1],
+    'released': [1, 1, 0],
+    'processing': [1, 2, 0],
+    'queued': [1, 2, 0],
+    'name': ['inc', 'add', 'dec'],
+    'show-name': ['inc', 'add', 'dec'],
+    'left': [0, 0, 1],
+    'right': [0.9, 0.9, 1.9],
+    'top': [0, -1, 0],
+    'bottom': [-0.8, -1.8, -0.8],
+    'color': ['#45BF6F', '#2E6C8E', '#440154'],
+    'released-loc': [0.18, 0.225, 1.0],
+    'memory-loc': [0.54, 0.45, 1.0],
+    'erred-loc': [0.54, 0.45, 1.9],
+    'processing-loc': [0.72, 0.9, 1.9],
+    'queued-loc': [0.9, 1.35, 1.9],
+    'done': ['3 / 5', '2 / 4', '1 / 1']}
     """
     width = 0.9
     names = sorted(msg["all"], key=msg["all"].get, reverse=True)
@@ -102,19 +108,28 @@ def progress_quads(msg, nrows=8, ncols=3):
     d["memory-loc"] = []
     d["erred-loc"] = []
     d["processing-loc"] = []
+    d["queued-loc"] = []
     d["done"] = []
-    for r, m, e, p, a, l in zip(
-        d["released"], d["memory"], d["erred"], d["processing"], d["all"], d["left"]
+    for r, m, e, p, q, a, l in zip(
+        d["released"],
+        d["memory"],
+        d["erred"],
+        d["processing"],
+        d["queued"],
+        d["all"],
+        d["left"],
     ):
         rl = width * r / a + l
         ml = width * (r + m) / a + l
         el = width * (r + m + e) / a + l
         pl = width * (p + r + m + e) / a + l
+        ql = width * (p + r + m + e + q) / a + l
         done = "%d / %d" % (r + m + e, a)
         d["released-loc"].append(rl)
         d["memory-loc"].append(ml)
         d["erred-loc"].append(el)
         d["processing-loc"].append(pl)
+        d["queued-loc"].append(ql)
         d["done"].append(done)
 
     return d

distributed/diagnostics/tests/test_progress_stream.py

@@ -18,6 +18,7 @@ def test_progress_quads():
         "erred": {"inc": 0, "dec": 1, "add": 0},
         "released": {"inc": 1, "dec": 0, "add": 1},
         "processing": {"inc": 1, "dec": 0, "add": 2},
+        "queued": {"inc": 1, "dec": 0, "add": 2},
     }
 
     d = progress_quads(msg, nrows=2)
@@ -35,11 +36,13 @@ def test_progress_quads():
         "memory": [2, 1, 0],
         "erred": [0, 0, 1],
         "processing": [1, 2, 0],
+        "queued": [1, 2, 0],
         "done": ["3 / 5", "2 / 4", "1 / 1"],
         "released-loc": [0.9 * 1 / 5, 0.25 * 0.9, 1.0],
         "memory-loc": [0.9 * 3 / 5, 0.5 * 0.9, 1.0],
         "erred-loc": [0.9 * 3 / 5, 0.5 * 0.9, 1.9],
         "processing-loc": [0.9 * 4 / 5, 1 * 0.9, 1 * 0.9 + 1],
+        "queued-loc": [1 * 0.9, 1.5 * 0.9, 1 * 0.9 + 1],
     }
     assert d == expected
 
@@ -52,6 +55,7 @@ def test_progress_quads_too_many():
         "erred": {k: 0 for k in keys},
         "released": {k: 0 for k in keys},
         "processing": {k: 0 for k in keys},
+        "queued": {k: 0 for k in keys},
     }
 
     d = progress_quads(msg, nrows=6, ncols=3)
@@ -78,6 +82,7 @@ async def test_progress_stream(c, s, a, b):
         "memory": {"div": 9, "inc": 1},
         "released": {"inc": 4},
         "processing": {},
+        "queued": {},
     }
     assert set(nbytes) == set(msg["all"])
     assert all(v > 0 for v in nbytes.values())
@@ -95,6 +100,7 @@ def test_progress_quads_many_functions():
         "erred": {fn: 0 for fn in funcnames},
         "released": {fn: 0 for fn in funcnames},
         "processing": {fn: 0 for fn in funcnames},
+        "queued": {fn: 0 for fn in funcnames},
     }
 
     d = progress_quads(msg, nrows=2)

distributed/distributed-schema.yaml

@@ -117,6 +117,26 @@ properties:
             description: |
               How frequently to balance worker loads
 
+          worker-saturation:
+            type: number
+            description: |
+              Controls how many root tasks are sent to workers (like a `readahead`).
+
+              `max(floor(worker-oversaturation * worker.nthreads), 1)` root tasks are sent to a
+              worker at a time. If `.inf`, all runnable tasks are immediately sent to workers.
+
+              Allowing oversaturation (> 1.0) means a worker will start running a new root task as
+              soon as it completes the previous, even if there is a higher-priority downstream task
+              to run. This reduces worker idleness, by letting workers do something while waiting for
+              further instructions from the scheduler.
+
+              This generally comes  at the expense of increased memory usage. It leads to "wider"
+              (more breadth-first) execution of the graph.
+
+              Compute-bound workloads benefit from oversaturation. Memory-bound workloads should
+              generally leave `worker-saturation` at 1.0, though 1.25-1.5 could slightly improve
+              performance if ample memory is available.
+
           worker-ttl:
             type:
             - string

distributed/distributed.yaml

@@ -22,6 +22,7 @@ distributed:
     events-log-length: 100000
     work-stealing: True     # workers should steal tasks from each other
     work-stealing-interval: 100ms  # Callback time for work stealing
+    worker-saturation: .inf # Send this fraction of nthreads root tasks to workers
     worker-ttl: "5 minutes" # like '60s'. Time to live for workers.  They must heartbeat faster than this
     pickle: True            # Is the scheduler allowed to deserialize arbitrary bytestrings
     preload: []             # Run custom modules with Scheduler

distributed/http/scheduler/tests/test_scheduler_http.py

@@ -137,7 +137,15 @@ async def fetch_metrics():
         ]
         return active_metrics, forgotten_tasks
 
-    expected = {"memory", "released", "processing", "waiting", "no-worker", "erred"}
+    expected = {
+        "memory",
+        "released",
+        "queued",
+        "processing",
+        "waiting",
+        "no-worker",
+        "erred",
+    }
 
     # Ensure that we get full zero metrics for all states even though the
     # scheduler did nothing, yet