Repository: dask/distributed

Issue:
Thoughts on Work Stealing 
I've been thinking about work stealing recently.  I thought I'd put down some thoughts.
### Situation

Recall that the scheduler sends a worker about as many tasks as it has cores.  Excess tasks that are known to be particularly well suited to the worker, either because it has most of the necessary data or because of a user-specified restriction, are kept in the `stacks`, a list associated with each worker.  Additionally, there is a common pool of tasks from which any worker may draw.

The need for work stealing arises when both of the following conditions hold:
1.  Some workers are _saturated_, meaning that those workers are busy processing as many tasks as they have cores, and they also have some excess tasks in their stack.
2.  Some workers are _idle_, meaning that they are not processing sufficient tasks, and so necessarily their stack and the common ready pool are both empty.

In these cases we may choose to pull tasks from the stack of a _saturated_ worker onto the stack of an _idle_ worker.  The idle worker _steals_ from the saturated worker.

Recall though that in some sense this is sub-optimal; there there was an original reason why the stolen task was assigned to the saturated worker.  We're going to end up moving excess data or violating a loose user-provided restriction in order to achieve load balancing.
### Objectives

We want to find "good" tasks to steal among the stacks of saturated workers.  A task can be good in a few ways:
1.  **Computation/communication time ratio**:  The ratio of its computation time to its data transfer time is high.  We would like tasks that take minutes to compute who only have a few bytes of data that can be transferred in a millisecond.
2.  **Steal from the back of the stack**: Workers put preferred tasks at the front of their stacks, leaving stale unwanted tasks at the back.  This practice aids low-memory-footprint operation.  The tasks at the front of the stack are fresh for the worker, and something related to the worker's current computations.  The tasks at the back of the stack are of lower priority for that worker, and so are more polite to steal.
3.  **Steal from the rich**.  Some workers have longer stacks than others.  We prefer to steal from these overburdened workers.
### Performance

So we could formulate an objective function and search over the stacks for tasks that satisfy this objective function.  However, the set of tasks in the stacks may be quite large (in the millions) and we may want to steal frequently.  As with all scheduling operations we would like to keep the administrative costs of work-stealing low, preferably even as the number of tasks increases.  This constant-time restriction stops us from considering traditional sorting techniques, like keeping around a heap.  This is especially important because work stealing often never comes up.  In the common case we would like to have low cost.
### Three approaches

Lets consider a few approaches that optimize for different things.  None of these will be ideal but they may bring us closer to a solution.
### Keep a sorted collection of time ratios

We want to steal tasks that are expensive to compute but cheap to move.  So as we move any task into the stacks we compute this ratio (using estimated costs) and put it into a heap.  If the task leaves the stacks then we remove it from the heap.  Each operation takes `log(n)` time for `n` the number of tasks in the stacks.  When we want to steal we just take the top element of the heap, a fast operation.

This solution optimizes computation/communication time ratio.  It does not select tasks low in the worker's stack, nor does it steal from the rich.  However, it doesn't do a particularly bad job at these either.  All things equal we're likely to pick a task from the middle of a stack and we're more likely to steal from the rich than from the poor.  In practice work stealing often arises in heavily imbalanced situations, where a small number of workers hold all of the pending tasks.
### Keep a pair of sets

In order to avoid the logarithmic time cost of the heap, we instead keep around two sets, "good to steal" and "not good to steal", and put tasks into each bin according to whether or not their time ratio is over a certain threshold.  This is faster, but results in less optimal task selection.  Like the sorted collection it doesn't take from the bottom of the stack nor does it steal from the rich.

We can extend this approach to "Keeping an ordered collection of k sets", say one for each order of magnitude of time ratio.  We can still add and remove tasks at constant time but now have more granularity about how good a task is to steal.  This also lets us respond to desperate times by stealing worse tasks when the task distribution is heavily lopsided.  This solution starts to look similar to non-comparative integer sorting techniques like Radix sort.
### Pull from the end

To optimize for "steal from the back" and "steal from the rich" we might find the worker with the largest stack (at a linear search through the number of workers, typically less than 1000) and then take the task from the end of that worker's stack.  This does not optimize for the time ratio but we could add in a quick check to only steal if it's reasonable and just not steal otherwise.



Discussion:
Quick question: in the current implementation, if a task is stolen, are its dependents re-evaluated? Such a cascade could factor into this discussion.
Also, this all assumes that we can measure the likely processing time of some task, which we can only do if we've seen something similar before.

The dependents are still `'waiting'` (because one of their dependencies (this key) is not yet in memory), so their ideal worker hasn't yet been assigned.  It will be assigned when all dependent tasks have been completed, including this one.

> Also, this all assumes that we can measure the likely processing time of some task, which we can only do if we've seen something similar before.

Correct, we maintain an exponential moving average of task duration in `self.task_durations`


Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -4,6 +4,7 @@
 from datetime import datetime, timedelta
 import logging
 import math
+from math import log
 import os
 import pickle
 import random
@@ -137,7 +138,8 @@ class Scheduler(Server):
         Which clients want each key.  The active targets of computation.
     * **nbytes:** ``{key: int}``:
         Number of bytes for a key as reported by workers holding that key.
-
+    * **stealable:** ``[[key]]``
+        A list of stacks of stealable keys, ordered by stealability
     * **ncores:** ``{worker: int}``:
         Number of cores owned by each worker
     * **idle:** ``{worker}``:
@@ -221,6 +223,8 @@ def __init__(self, center=None, loop=None,
         self.exceptions = dict()
         self.tracebacks = dict()
         self.exceptions_blame = dict()
+        self.stealable = [set() for i in range(12)]
+        self.stealable_unknown_durations = defaultdict(set)
 
         # Worker state
         self.ncores = dict()
@@ -1649,6 +1653,7 @@ def transition_waiting_ready(self, key):
                 else:
                     self.stacks[new_worker].append(key)
                     self.maybe_idle.add(new_worker)
+                    self.put_key_in_stealable(key)
                     self.task_state[key] = 'stacks'
             else:
                 self.ready.appendleft(key)
@@ -1678,6 +1683,7 @@ def transition_ready_processing(self, key, worker=None, latency=5e-3):
             self.rprocessing[key].add(worker)
             self.occupancy[worker] += duration
             self.task_state[key] = 'processing'
+            self.remove_key_from_stealable(key)
 
             logger.debug("Send job to worker: %s, %s", worker, key)
 
@@ -1728,6 +1734,11 @@ def transition_processing_memory(self, key, nbytes=None, type=None,
                                   + 0.5 * new_duration)
 
                 self.task_duration[ks] = avg_duration
+                if ks in self.stealable_unknown_durations:
+                    for k in self.stealable_unknown_durations.pop(ks, ()):
+                        if self.task_state[k] == 'stacks':
+                            self.put_key_in_stealable(k)
+
                 info['last-task'] = compute_stop
 
             ############################
@@ -2259,7 +2270,7 @@ def keyfunc(w):
                         for worker, count in zip(workers2, counts):
                             self.ensure_occupied_queue(worker, count=count)
 
-            if self.idle and self.saturated:
+            if self.idle and any(self.stealable):
                 thieves = self.work_steal()
                 for worker in thieves:
                     self.ensure_occupied_stacks(worker)
@@ -2285,7 +2296,7 @@ def ensure_occupied_stacks(self, worker):
                 self.occupancy[worker] < latency * self.ncores[worker])):
             key = stack.pop()
 
-            if self.task_state[key] == 'stacks':
+            if self.task_state.get(key) == 'stacks':
                 r = self.transition(key, 'processing',
                                     worker=worker, latency=latency)
 
@@ -2298,6 +2309,19 @@ def ensure_occupied_stacks(self, worker):
                 self.saturated.remove(worker)
             self._check_idle(worker)
 
+    def put_key_in_stealable(self, key):
+        ratio, loc = self.steal_time_ratio(key)
+        if ratio is not None:
+            self.stealable[loc].add(key)
+
+    def remove_key_from_stealable(self, key):
+        ratio, loc = self.steal_time_ratio(key)
+        if ratio is not None:
+            try:
+                self.stealable[loc].remove(key)
+            except:
+                pass
+
     def ensure_occupied_queue(self, worker, count):
         """
         Send at most count tasks from the ready queue to the specified worker
@@ -2318,7 +2342,7 @@ def ensure_occupied_queue(self, worker, count):
 
         self._check_idle(worker)
 
-    def work_steal(self, bandwidth=None):
+    def work_steal(self):
         """ Steal tasks from saturated workers to idle workers
 
         This moves tasks from the bottom of the stacks of over-occupied workers
@@ -2328,87 +2352,75 @@ def work_steal(self, bandwidth=None):
         --------
         Scheduler.ensure_occupied
         """
-        bandwidth = bandwidth if bandwidth is not None else BANDWIDTH
-        if not self.idle or not self.saturated:
-            return
-
-        thieves = set()  # Output list
-
-        idle = iter(self.idle)  # we will walk down these two sequences
-        remove_idle = set()
-        saturated = iter(self.saturated)
-        remove_saturated = set()
-
-        thief = next(idle)
-        victim = next(saturated)
-        try:
-            while True:
-                if victim == thief:
-                    raise ValueError()
-                thieves.add(thief)  # add to output
-                n = (self.ncores[thief]  # number of tasks to consume
-                  - len(self.processing[thief])
-                  - len(self.stacks[thief]))
-                stack = self.stacks[victim]
-                while n > 0 and stack:
-                    key = stack.popleft()
-                    if key not in self.tasks:
-                        continue
-                    if self.should_steal(key):
-                        self.stacks[thief].append(key)
-                        n -= 1
-                    else:
-                        stack.appendleft(key)  # replace task in victim's stack
-                        remove_saturated.add(victim)
-                        victim = next(saturated)
-                        break
-
-                if not stack:
-                    remove_saturated.add(victim)
-                    victim = next(saturated)
-
-                if n <= 0:
-                    remove_idle.add(thief)
-                    thief = next(idle)
-        except StopIteration:
-            pass
-        for worker in remove_saturated:
-            self.saturated.remove(worker)
-        for worker in remove_idle:
-            self.idle.remove(worker)
-        logger.debug('Stolen tasks for %d workers', len(thieves))
-        return thieves
-
-    def should_steal(self, key, bandwidth=None):
-        """ Is a key good for stealing from its chosen worker?
+        with log_errors(pdb=True):
+            thieves = set()
+            for level, stealable in enumerate(self.stealable[:-1]):
+                if not stealable:
+                    continue
+                if len(self.idle) == len(self.ncores):  # no stacks
+                    stealable.clear()
+                    continue
+                # Enough idleness to continue?
+                ratio = 2 ** (level - 3)
+                n_saturated = len(self.ncores) - len(self.idle)
+                duration_if_hold = len(stealable) / n_saturated
+                duration_if_steal = ratio
+                if level > 1 and duration_if_hold < duration_if_steal:
+                    break
+                while stealable and self.idle:
+                    for w in list(self.idle):
+                        try:
+                            key = stealable.pop()
+                        except:
+                            break
+                        else:
+                            if self.task_state.get(key, 'stacks'):
+                                self.stacks[w].append(key)
+                                thieves.add(w)
+                                if (self.ncores[w] <=
+                                    len(self.processing[w]) + len(self.stacks[w])):
+                                    self.idle.remove(w)
+
+                if stealable:
+                    break
+            logger.debug('Stolen tasks for %d workers', len(thieves))
+            return thieves
 
-        It must have the following attributes
+    def steal_time_ratio(self, key, bandwidth=BANDWIDTH):
+        """ The compute to communication time ratio of a key
 
-        1.  Not have too many dependencies
-        2.  Not be restricted to run on that worker
-        3   Take less time to transfer than to compute
+        Returns
+        -------
 
-        See also
-        --------
-        Scheduler.work_steal
+        ratio: The compute/communication time ratio of the task
+        loc: The self.stealable bin into which this key should go
         """
-        bandwidth = bandwidth if bandwidth is not None else BANDWIDTH
-        if len(self.dependencies[key]) > 10:
-            return False
         if key in self.restrictions and key not in self.loose_restrictions:
-            return False
+            return None, None  # don't steal
 
         nbytes = sum(self.nbytes.get(k, 1000) for k in self.dependencies[key])
         transfer_time = nbytes / bandwidth
         try:
             compute_time = self.task_duration[key_split(key)]
-            return transfer_time < compute_time
         except KeyError:
-            return False
+            self.stealable_unknown_durations[key_split(key)].add(key)
+            return None, None
+        else:
+            if transfer_time:
+                ratio = compute_time / transfer_time
+            else:
+                ratio = 10000
+            if ratio > 8:
+                loc = 0
+            elif ratio < 2**-8:
+                loc = -1
+            else:
+                loc = int(-round(log(ratio) / log(2), 0) + 3)
+            return ratio, loc
 
     def issaturated(self, worker, latency=5e-3):
         """
-        Determin if a worker has enough work to avoid being idle
+        Determine if a worker has enough work to avoid being idle
 
         A worker is saturated if the following criteria are met
 


Test patch:
diff --git a/distributed/tests/test_executor.py b/distributed/tests/test_executor.py
--- a/distributed/tests/test_executor.py
+++ b/distributed/tests/test_executor.py
@@ -2690,18 +2690,6 @@ def test_workers_register_indirect_data(e, s, a, b):
     s.validate_state()
 
 
-@pytest.mark.skipif(not sys.platform.startswith('linux'),
-                    reason="Need 127.0.0.2 to mean localhost")
-@gen_cluster(executor=True, ncores=[('127.0.0.1', 2), ('127.0.0.2', 2)],
-        timeout=20)
-def test_work_stealing(e, s, a, b):
-    [x] = yield e._scatter([1])
-    futures = e.map(slowadd, range(50), [x] * 50)
-    yield _wait(futures)
-    assert len(a.data) > 10
-    assert len(b.data) > 10
-
-
 @gen_cluster(executor=True)
 def test_submit_on_cancelled_future(e, s, a, b):
     x = e.submit(inc, 1)
@@ -2945,78 +2933,6 @@ def test_scheduler_saturates_cores_random(e, s, a, b):
             yield gen.sleep(0.01)
 
 
-@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
-def test_dont_steal_expensive_data_fast_computation(e, s, a, b):
-    np = pytest.importorskip('numpy')
-    x = e.submit(np.arange, 1000000, workers=a.address)
-    yield _wait([x])
-    future = e.submit(np.sum, [1], workers=a.address)  # learn that sum is fast
-    yield _wait([future])
-
-    cheap = [e.submit(np.sum, x, pure=False, workers=a.address,
-                      allow_other_workers=True) for i in range(10)]
-    yield _wait(cheap)
-    assert len(b.data) == 0
-    assert len(a.data) == 12
-
-
-@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
-def test_steal_cheap_data_slow_computation(e, s, a, b):
-    x = e.submit(slowinc, 100, delay=0.1)  # learn that slowinc is slow
-    yield _wait([x])
-
-    futures = e.map(slowinc, range(10), delay=0.01, workers=a.address,
-                    allow_other_workers=True)
-    yield _wait(futures)
-    assert abs(len(a.data) - len(b.data)) < 3
-
-
-@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
-def test_steal_expensive_data_slow_computation(e, s, a, b):
-    np = pytest.importorskip('numpy')
-
-    x = e.submit(slowinc, 100, delay=0.1, workers=a.address)
-    yield _wait([x])  # learn that slowinc is slow
-
-    x = e.submit(np.arange, 1000000, workers=a.address)  # put expensive data
-    yield _wait([x])
-
-    slow = [e.submit(slowinc, x, delay=0.1, pure=False) for i in range(4)]
-    yield _wait([slow])
-
-    assert b.data  # not empty
-
-
-@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 10)
-def test_worksteal_many_thieves(e, s, *workers):
-    x = e.submit(slowinc, -1, delay=0.1)
-    yield x._result()
-
-    xs = e.map(slowinc, [x] * 100, pure=False, delay=0.01)
-
-    yield _wait(xs)
-
-    for w, keys in s.has_what.items():
-        assert 2 < len(keys) < 50
-
-
-@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
-def test_dont_steal_unknown_functions(e, s, a, b):
-    futures = e.map(inc, [1, 2], workers=a.address, allow_other_workers=True)
-    yield _wait(futures)
-    assert len(a.data) == 2
-    assert len(b.data) == 0
-
-
-@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
-def test_eventually_steal_unknown_functions(e, s, a, b):
-    futures = e.map(slowinc, range(10), delay=0.1,  workers=a.address,
-                    allow_other_workers=True)
-    yield _wait(futures)
-    assert len(a.data) >= 3
-    assert len(b.data) >= 3
-
-
 @gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 4, timeout=None)
 def test_cancel_stress(e, s, *workers):
     da = pytest.importorskip('dask.array')
diff --git a/distributed/tests/test_steal.py b/distributed/tests/test_steal.py
new file mode 100644
--- /dev/null
+++ b/distributed/tests/test_steal.py
@@ -0,0 +1,131 @@
+from __future__ import print_function, division, absolute_import
+
+import sys
+
+import pytest
+from toolz import sliding_window
+from tornado import gen
+
+import dask
+from dask import delayed
+from distributed.executor import Executor, _wait, wait
+from distributed.utils_test import (cluster, slowinc, slowadd, randominc,
+        _test_scheduler, loop, inc, dec, div, throws, gen_cluster, gen_test,
+        double, deep)
+
+import pytest
+
+
+@pytest.mark.skipif(not sys.platform.startswith('linux'),
+                    reason="Need 127.0.0.2 to mean localhost")
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 2), ('127.0.0.2', 2)],
+        timeout=20)
+def test_work_stealing(e, s, a, b):
+    [x] = yield e._scatter([1])
+    futures = e.map(slowadd, range(50), [x] * 50)
+    yield _wait(futures)
+    assert len(a.data) > 10
+    assert len(b.data) > 10
+
+
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
+def test_dont_steal_expensive_data_fast_computation(e, s, a, b):
+    np = pytest.importorskip('numpy')
+    x = e.submit(np.arange, 1000000, workers=a.address)
+    yield _wait([x])
+    future = e.submit(np.sum, [1], workers=a.address)  # learn that sum is fast
+    yield _wait([future])
+
+    cheap = [e.submit(np.sum, x, pure=False, workers=a.address,
+                      allow_other_workers=True) for i in range(10)]
+    yield _wait(cheap)
+    assert len(b.data) == 0
+    assert len(a.data) == 12
+
+
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
+def test_steal_cheap_data_slow_computation(e, s, a, b):
+    x = e.submit(slowinc, 100, delay=0.1)  # learn that slowinc is slow
+    yield _wait([x])
+
+    futures = e.map(slowinc, range(10), delay=0.01, workers=a.address,
+                    allow_other_workers=True)
+    yield _wait(futures)
+    assert abs(len(a.data) - len(b.data)) < 3
+
+
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
+def test_steal_expensive_data_slow_computation(e, s, a, b):
+    np = pytest.importorskip('numpy')
+
+    x = e.submit(slowinc, 100, delay=0.1, workers=a.address)
+    yield _wait([x])  # learn that slowinc is slow
+
+    x = e.submit(np.arange, 1000000, workers=a.address)  # put expensive data
+    yield _wait([x])
+
+    slow = [e.submit(slowinc, x, delay=0.1, pure=False) for i in range(4)]
+    yield _wait([slow])
+
+    assert b.data  # not empty
+
+
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 10)
+def test_worksteal_many_thieves(e, s, *workers):
+    x = e.submit(slowinc, -1, delay=0.1)
+    yield x._result()
+
+    xs = e.map(slowinc, [x] * 100, pure=False, delay=0.01)
+
+    yield _wait(xs)
+
+    for w, keys in s.has_what.items():
+        assert 2 < len(keys) < 50
+
+
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
+def test_dont_steal_unknown_functions(e, s, a, b):
+    futures = e.map(inc, [1, 2], workers=a.address, allow_other_workers=True)
+    yield _wait(futures)
+    assert len(a.data) == 2
+    assert len(b.data) == 0
+
+
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 2)
+def test_eventually_steal_unknown_functions(e, s, a, b):
+    futures = e.map(slowinc, range(10), delay=0.1,  workers=a.address,
+                    allow_other_workers=True)
+    yield _wait(futures)
+    assert len(a.data) >= 3
+    assert len(b.data) >= 3
+
+
+@pytest.mark.xfail(reason='')
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 3)
+def test_steal_related_tasks(e, s, a, b, c):
+    futures = e.map(slowinc, range(20), delay=0.05, workers=a.address,
+                    allow_other_workers=True)
+
+    yield _wait(futures)
+
+    nearby = 0
+    for f1, f2 in sliding_window(2, futures):
+        if s.who_has[f1.key] == s.who_has[f2.key]:
+            nearby += 1
+
+    assert nearby > 10
+
+
+@gen_cluster(executor=True, ncores=[('127.0.0.1', 1)] * 10, timeout=1000)
+def test_dont_steal_fast_tasks(e, s, *workers):
+    np = pytest.importorskip('numpy')
+    x = e.submit(np.random.random, 10000000, workers=workers[0].address)
+
+    def do_nothing(x, y=None):
+        pass
+
+    futures = e.map(do_nothing, range(1000), y=x)
+
+    yield _wait(futures)
+
+    assert len(s.has_what[workers[0].address]) == 1001