Repository: dask/distributed

Issue:
New Worker not used by previously submitted map
To roughly reproduce, set up a scheduler on one machine with `dscheduler`, start up several workers with `dworker`, and then do

``` python
from distributed import Executor

executor = Executor('larry:8786')
L = executor.map(long_running_function, inputs)
```

where the number of `inputs` is greater than the number of processes available among the workers. 

Adding an additional worker with `dworker` later is recognized by the executor, but none of the pending operations appear to be sent to that new worker. New submissions using e.g. `Executor.map`, however, do make use of the new worker.

Not sure if this is technically a bug, but it would be cool if adding more workers helps move pending operations from a `map` further along.



Discussion:
Agreed.  

The current state is that a task is assigned a worker as soon as all dependencies for that task become available.  It lingers in the stack until it goes to the worker for processing or until the assigned worker dies, at which point it will be appropriately reassigned.

We could consider re-balancing ready tasks on adding a new worker.  We could also consider re-balancing ready tasks whenever a worker becomes idle, which would also help to redistributed straggler tasks.

The only reason we haven't done this yet is that the scheduler doesn't currently hold the data structures necessary to execute this re-balancing operation in linear time relative to the number of moved tasks; so it's potentially costly if we have very many available tasks.  We would need to walk through the tasks on the worker stacks one by one to find good candidates.  It is tricky to add and track additional state to make this more efficient, but doable.

My plan until now has been to handle some other items first.  I'd be happy to re-prioritize though if dynamic scaling is of concrete use.  @dotsdl how meaningfully would this impact your current work?


Resolving patch:
diff --git a/distributed/nanny.py b/distributed/nanny.py
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -11,7 +11,7 @@
 from tornado import gen
 
 from .core import Server, rpc, write
-from .utils import get_ip
+from .utils import get_ip, ignoring
 
 
 logger = logging.getLogger(__name__)
@@ -201,4 +201,5 @@ def start():
             q.put({'port': worker.port, 'dir': worker.local_dir})  # pragma: no cover
 
     loop.add_callback(start)  # pragma: no cover
-    loop.start()  # pragma: no cover
+    with ignoring(KeyboardInterrupt):
+        loop.start()  # pragma: no cover
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -76,8 +76,12 @@ class Scheduler(Server):
         Dictionary like dependencies but excludes keys already computed
     * **waiting_data:** ``{key: {key}}``:
         Dictionary like dependents but excludes keys already computed
+    * **ready:** ``deque(key)``
+        Keys that are ready to run, but not yet assigned to a worker
     * **ncores:** ``{worker: int}``:
         Number of cores owned by each worker
+    * **idle:** ``{worker}``:
+        Set of workers that are not fully utilized
     * **services:** ``{str: port}``:
         Other services running on this scheduler, like HTTP
     * **worker_services:** ``{worker: {str: port}}``:
@@ -163,6 +167,8 @@ def __init__(self, center=None, loop=None,
         self.stacks = dict()
         self.waiting = dict()
         self.waiting_data = dict()
+        self.ready = deque()
+        self.idle = set()
         self.who_has = defaultdict(set)
         self.deleted_keys = defaultdict(set)
         self.who_wants = defaultdict(set)
@@ -188,7 +194,7 @@ def __init__(self, center=None, loop=None,
                                  'client-releases-keys': self.client_releases_keys,
                                  'restart': self.restart}
 
-        self.handlers = {'register-client': self.control_stream,
+        self.handlers = {'register-client': self.add_client,
                          'scatter': self.scatter,
                          'register': self.add_worker,
                          'unregister': self.remove_worker,
@@ -246,7 +252,8 @@ def start(self, port=8786, start_queues=True):
         collections = [self.tasks, self.dependencies, self.dependents,
                 self.waiting, self.waiting_data, self.in_play, self.keyorder,
                 self.nbytes, self.processing, self.restrictions,
-                self.loose_restrictions]
+                self.loose_restrictions, self.ready, self.who_wants,
+                self.wants_what]
         for collection in collections:
             collection.clear()
 
@@ -268,7 +275,6 @@ def start(self, port=8786, start_queues=True):
 
         self.heal_state()
 
-
         if start_queues:
             self.handle_queues(self.scheduler_queues[0], None)
 
@@ -332,7 +338,13 @@ def cleanup(self):
             q.put_nowait({'op': 'close'})
 
     def mark_ready_to_run(self, key):
-        """ Send task to an appropriate worker.  Trigger that worker.
+        """
+        Mark a task as ready to run.
+
+        If the task should be assigned to a worker then make that determination
+        and assign appropriately.  Otherwise place task in the ready queue.
+
+        Trigger appropriate workers if idle.
 
         See Also
         --------
@@ -340,16 +352,32 @@ def mark_ready_to_run(self, key):
         Scheduler.ensure_occupied
         """
         logger.debug("Mark %s ready to run", key)
+        if key in self.exceptions_blame:
+            logger.debug("Don't mark ready, task has already failed")
+            return
+
         if key in self.waiting:
             assert not self.waiting[key]
             del self.waiting[key]
 
-        new_worker = decide_worker(self.dependencies, self.stacks,
-                self.who_has, self.restrictions, self.loose_restrictions,
-                self.nbytes, key)
+        if self.dependencies.get(key, None) or key in self.restrictions:
+            new_worker = decide_worker(self.dependencies, self.stacks,
+                    self.who_has, self.restrictions, self.loose_restrictions,
+                    self.nbytes, key)
+            if not new_worker:
+                raise ValueError("No valid workers found")
+            else:
+                self.stacks[new_worker].append(key)
+                self.ensure_occupied(new_worker)
+        else:
+            self.ready.appendleft(key)
+            self.ensure_idle_ready()
 
-        self.stacks[new_worker].append(key)
-        self.ensure_occupied(new_worker)
+    def ensure_idle_ready(self):
+        """ Run ready tasks on idle workers """
+        while self.idle and self.ready and self.ncores:
+            worker = min(self.idle, key=lambda w: len(self.processing[w]))
+            self.ensure_occupied(worker)
 
     def mark_key_in_memory(self, key, workers=None, type=None):
         """ Mark that a key now lives in distributed memory """
@@ -387,13 +415,21 @@ def mark_key_in_memory(self, key, workers=None, type=None):
         self.report(msg)
 
     def ensure_occupied(self, worker):
-        """ Send tasks to worker while it has tasks and free cores """
+        """ Send tasks to worker while it has tasks and free cores
+
+        These tasks may come from the worker's own stacks or from the global
+        ready deque.
+
+        We update the idle workers set appropriately.
+        """
         logger.debug('Ensure worker is occupied: %s', worker)
         while (self.stacks[worker] and
                self.ncores[worker] > len(self.processing[worker])):
             key = self.stacks[worker].pop()
             if key not in self.tasks:
                 continue
+            if self.who_has.get(key):
+                continue
             self.processing[worker].add(key)
             logger.debug("Send job to worker: %s, %s, %s", worker, key)
             self.worker_queues[worker].put_nowait(
@@ -403,27 +439,26 @@ def ensure_occupied(self, worker):
                      'who_has': {dep: self.who_has[dep] for dep in
                                  self.dependencies[key]}})
 
-    def seed_ready_tasks(self, keys=None):
-        """ Distribute many leaf tasks among workers
+        while (self.ready and
+               self.ncores[worker] > len(self.processing[worker])):
+            key = self.ready.pop()
+            if key not in self.tasks:
+                continue
+            if self.who_has.get(key):
+                continue
+            self.processing[worker].add(key)
+            logger.debug("Send job to worker: %s, %s, %s", worker, key)
+            self.worker_queues[worker].put_nowait(
+                    {'op': 'compute-task',
+                     'key': key,
+                     'task': self.tasks[key],
+                     'who_has': {dep: self.who_has[dep] for dep in
+                                 self.dependencies[key]}})
 
-        Takes an iterable of keys to consider for execution
-
-        See Also
-        --------
-        assign_many_tasks
-        Scheduler.ensure_occupied
-        """
-        if keys is None:
-            keys = self.tasks
-        new_stacks = assign_many_tasks(
-                self.dependencies, self.waiting, self.keyorder, self.who_has,
-                self.stacks, self.restrictions, self.loose_restrictions,
-                self.nbytes,
-                [k for k in keys if k in self.waiting and not self.waiting[k]])
-        logger.debug("Seed ready tasks: %s", new_stacks)
-        for worker, stack in new_stacks.items():
-            if stack:
-                self.ensure_occupied(worker)
+        if self.ncores[worker] > len(self.processing[worker]):
+            self.idle.add(worker)
+        elif worker in self.idle:
+            self.idle.remove(worker)
 
     def update_data(self, who_has=None, nbytes=None, client=None):
         """
@@ -525,8 +560,6 @@ def mark_missing_data(self, missing=None, key=None, worker=None):
             logger.debug('task missing data, %s, %s', key, self.waiting)
             self.ensure_occupied(worker)
 
-        self.seed_ready_tasks()
-
         # self.validate(allow_overlap=True, allow_bad_stacks=True)
 
     def log_state(self, msg=''):
@@ -554,6 +587,8 @@ def remove_worker(self, stream=None, address=None, heal=True):
         del self.stacks[address]
         del self.processing[address]
         del self.worker_services[address]
+        if address in self.idle:
+            self.idle.remove(address)
         if not self.stacks:
             logger.critical("Lost all workers")
         missing_keys = set()
@@ -580,20 +615,20 @@ def add_worker(self, stream=None, address=None, keys=(), ncores=None,
             self.processing[address] = set()
             self.stacks[address] = []
             self.worker_queues[address] = Queue()
+
         for key in keys:
             self.mark_key_in_memory(key, [address])
 
         self._worker_coroutines.append(self.worker(address))
 
+        if self.ncores[address] > len(self.processing[address]):
+            self.idle.add(address)
+
+        self.ensure_occupied(address)
+
         logger.info("Register %s", str(address))
         return b'OK'
 
-    def remove_client(self, client=None):
-        logger.info("Remove client %s", client)
-        self.client_releases_keys(self.wants_what.get(client, ()), client)
-        with ignoring(KeyError):
-            del self.wants_what[client]
-
     def update_graph(self, client=None, tasks=None, keys=None,
                      dependencies=None, restrictions=None,
                      loose_restrictions=None):
@@ -605,10 +640,50 @@ def update_graph(self, client=None, tasks=None, keys=None,
             if tasks[k] is k:
                 del tasks[k]
 
-        update_state(self.tasks, self.dependencies, self.dependents,
-                self.who_wants, self.wants_what, self.who_has, self.in_play,
-                self.waiting, self.waiting_data, tasks, keys, dependencies,
-                client)
+        if not isinstance(keys, set):
+            keys = set(keys)
+
+        for key, task in tasks.items():
+            if key not in self.tasks:
+                self.tasks[key] = task
+
+        for key in tasks:  # add dependencies/dependents
+            if key in self.dependencies:
+                continue
+
+            task = self.tasks[key]
+            self.dependencies[key] = set(dependencies.get(key, ()))
+
+            for dep in self.dependencies[key]:
+                if dep not in self.dependents:
+                    self.dependents[dep] = set()
+                self.dependents[dep].add(key)
+
+            if key not in self.dependents:
+                self.dependents[key] = set()
+
+        for k in keys:
+            self.who_wants[k].add(client)
+            self.wants_what[client].add(k)
+
+        ready_to_run = set()
+
+        exterior = keys_outside_frontier(self.dependencies, keys, self.in_play)
+        self.in_play |= exterior
+        for key in exterior:
+            deps = self.dependencies[key]
+            wait_keys = {d for d in deps if not self.who_has.get(d)}
+            if wait_keys:
+                self.waiting[key] = wait_keys
+            else:
+                ready_to_run.add(key)
+            for dep in deps:
+                if dep not in self.waiting_data:
+                    self.waiting_data[dep] = set()
+                self.waiting_data[dep].add(key)
+
+            if key not in self.waiting_data:
+                self.waiting_data[key] = set()
 
         if restrictions:
             restrictions = {k: set(map(ensure_ip, v))
@@ -630,9 +705,8 @@ def update_graph(self, client=None, tasks=None, keys=None,
                 if dep in self.exceptions_blame:
                     self.mark_failed(key, self.exceptions_blame[dep])
 
-        self.seed_ready_tasks(tasks)
         for key in keys:
-            if self.who_has[key]:
+            if self.who_has.get(key):
                 self.mark_key_in_memory(key)
 
         for plugin in self.plugins[:]:
@@ -641,6 +715,11 @@ def update_graph(self, client=None, tasks=None, keys=None,
             except Exception as e:
                 logger.exception(e)
 
+        for key in ready_to_run:
+            self.mark_ready_to_run(key)
+
+        self.ensure_idle_ready()
+
     def client_releases_keys(self, keys=None, client=None):
         for k in list(keys):
             with ignoring(KeyError):
@@ -726,26 +805,33 @@ def heal_state(self):
         logger.debug("Heal state")
         self.log_state("Before Heal")
         state = heal(self.dependencies, self.dependents, self.who_has,
-                self.stacks, self.processing, self.waiting, self.waiting_data)
+                self.stacks, self.processing, self.waiting, self.waiting_data,
+                self.ready)
         released = state['released']
         self.in_play.clear(); self.in_play.update(state['in_play'])
-        add_keys = {k for k, v in self.waiting.items() if not v}
+
         for key in set(self.who_wants) & released:
             self.report({'op': 'lost-key', 'key': key})
-        if self.stacks:
-            for key in add_keys:
-                self.mark_ready_to_run(key)
+
+        for key in state['ready']:
+            self.mark_ready_to_run(key)
 
         self.delete_data(keys=set(self.who_has) & released - set(self.who_wants))
         self.in_play.update(self.who_has)
+
+        for worker in self.ncores:
+            self.ensure_occupied(worker)
+
         self.log_state("After Heal")
 
     def my_heal_missing_data(self, missing):
         """ Recover from lost data """
         logger.debug("Heal from missing data")
-        return heal_missing_data(self.tasks, self.dependencies, self.dependents,
-                self.who_has, self.in_play, self.waiting,
+        ready_to_run = heal_missing_data(self.tasks, self.dependencies,
+                self.dependents, self.who_has, self.in_play, self.waiting,
                 self.waiting_data, missing)
+        for key in ready_to_run:
+            self.mark_ready_to_run(key)
 
     def report(self, msg):
         """ Publish updates to all listening Queues and Streams """
@@ -779,7 +865,7 @@ def handle_queues(self, scheduler_queue, report_queue):
         return future
 
     @gen.coroutine
-    def control_stream(self, stream, address=None, client=None):
+    def add_client(self, stream, address=None, client=None):
         """ Listen to messages from an IOStream """
         logger.info("Connection to %s, %s", type(self).__name__, client)
         self.streams[client] = stream
@@ -793,6 +879,12 @@ def control_stream(self, stream, address=None, client=None):
             logger.info("Close connection to %s, %s", type(self).__name__,
                         client)
 
+    def remove_client(self, client=None):
+        logger.info("Remove client %s", client)
+        self.client_releases_keys(self.wants_what.get(client, ()), client)
+        with ignoring(KeyError):
+            del self.wants_what[client]
+
     @gen.coroutine
     def handle_messages(self, in_queue, report, client=None):
         """ Master coroutine.  Handles inbound messages.
@@ -902,6 +994,7 @@ def worker_core(self, ident, i):
                     serialized = False
                 else:
                     serialized = True
+
                 response, content = yield worker.compute(who_has=who_has,
                                                          key=key,
                                                          report=self.center
@@ -1029,6 +1122,7 @@ def restart(self):
         assert all(resp == b'OK' for resp in resps)
         if self.center:
             yield self.sync_center()
+
         self.start()
 
         logger.debug("All workers reporting in")
@@ -1043,7 +1137,7 @@ def restart(self):
     def validate(self, allow_overlap=False, allow_bad_stacks=False):
         released = set(self.tasks) - self.in_play
         validate_state(self.dependencies, self.dependents, self.waiting,
-                self.waiting_data, self.who_has, self.stacks,
+                self.waiting_data, self.ready, self.who_has, self.stacks,
                 self.processing, None, released, self.in_play, self.who_wants,
                 self.wants_what,
                 allow_overlap=allow_overlap, allow_bad_stacks=allow_bad_stacks)
@@ -1159,7 +1253,7 @@ def decide_worker(dependencies, stacks, who_has, restrictions,
                 else:
                     raise ValueError("Task has no valid workers", key, r)
     if not workers or not stacks:
-        raise ValueError("No workers found")
+        return None
 
     commbytes = {w: sum(nbytes[k] for k in dependencies[key]
                                    if w not in who_has[k])
@@ -1172,62 +1266,7 @@ def decide_worker(dependencies, stacks, who_has, restrictions,
     return worker
 
 
-def update_state(tasks, dependencies, dependents, who_wants, wants_what,
-                 who_has, in_play,
-                 waiting, waiting_data, new_tasks, new_keys, new_dependencies,
-                 client):
-    """ Update state given new dask graph and output keys
-
-    This should operate in linear time relative to the size of edges of the
-    added graph.  It assumes that the current runtime state is valid.
-    """
-    tasks.update(new_tasks)
-    if not isinstance(new_keys, set):
-        new_keys = set(new_keys)
-
-    for key in new_tasks:  # add dependencies/dependents
-        if key in dependencies:
-            continue
-
-        task = new_tasks[key]
-        deps = new_dependencies.get(key, ())
-        dependencies[key] = set(deps)
-
-        for dep in deps:
-            if dep not in dependents:
-                dependents[dep] = set()
-            dependents[dep].add(key)
-
-        if key not in dependents:
-            dependents[key] = set()
-
-    for k in new_keys:
-        who_wants[k].add(client)
-        wants_what[client].add(k)
-
-    exterior = keys_outside_frontier(dependencies, new_keys, in_play)
-    in_play |= exterior
-    for key in exterior:
-        deps = dependencies[key]
-        waiting[key] = {d for d in deps if not (d in who_has and who_has[d])}
-        for dep in deps:
-            if dep not in waiting_data:
-                waiting_data[dep] = set()
-            waiting_data[dep].add(key)
-
-        if key not in waiting_data:
-            waiting_data[key] = set()
-
-    return {'tasks': tasks,
-            'dependencies': dependencies,
-            'dependents': dependents,
-            'who_wants': who_wants,
-            'wants_what': wants_what,
-            'waiting': waiting,
-            'waiting_data': waiting_data}
-
-
-def validate_state(dependencies, dependents, waiting, waiting_data,
+def validate_state(dependencies, dependents, waiting, waiting_data, ready,
         who_has, stacks, processing, finished_results, released, in_play,
         who_wants, wants_what, allow_overlap=False, allow_bad_stacks=False,
         **kwargs):
@@ -1249,6 +1288,7 @@ def validate_state(dependencies, dependents, waiting, waiting_data,
     def check_key(key):
         """ Validate a single key, recurse downards """
         val = sum([key in waiting,
+                   key in ready,
                    key in in_stacks,
                    key in in_processing,
                    who_has.get(key) is not None,
@@ -1272,6 +1312,7 @@ def check_key(key):
                 raise ValueError("Key in stacks/processing without all deps",
                                  key)
             assert not waiting.get(key)
+            assert key not in ready
 
         if finished_results is not None:
             if key in finished_results:
@@ -1286,6 +1327,13 @@ def check_key(key):
             for client in s:
                 assert key in wants_what[client]
 
+        if key in waiting:
+            assert waiting[key], 'waiting empty'
+
+
+        if key in ready:
+            assert key not in waiting
+
         return True
 
     assert all(map(check_key, keys))
@@ -1334,7 +1382,7 @@ def keys_outside_frontier(dependencies, keys, frontier):
 
 
 def heal(dependencies, dependents, who_has, stacks, processing, waiting,
-        waiting_data, **kwargs):
+        waiting_data, ready, **kwargs):
     """ Make a possibly broken runtime state consistent
 
     Sometimes we lose intermediate values.  In these cases it can be tricky to
@@ -1352,6 +1400,8 @@ def heal(dependencies, dependents, who_has, stacks, processing, waiting,
 
     waiting_data.clear()
     waiting.clear()
+    ready.clear()
+    ready = set()
     finished_results = set()
 
     def remove_from_stacks(key):
@@ -1378,8 +1428,11 @@ def make_accessible(key):
         for dep in dependencies[key]:
             make_accessible(dep)  # recurse
 
-        waiting[key] = {dep for dep in dependencies[key]
-                            if not who_has.get(dep)}
+        s = {dep for dep in dependencies[key] if not who_has.get(dep)}
+        if s:
+            waiting[key] = s
+        else:
+            ready.add(key)
 
     for key in outputs:
         make_accessible(key)
@@ -1406,11 +1459,14 @@ def unrunnable(key):
         if key in waiting:
             assert not waiting[key]
             del waiting[key]
+        if key in ready:
+            ready.remove(key)
 
     finished_results = {key for key in outputs if who_has.get(key)}
 
     in_play = ({k for k, v in who_has.items() if v}
             | set(waiting)
+            | ready
             | (set.union(*processing.values()) if processing else set())
             | set(concat(stacks.values())))
 
@@ -1419,7 +1475,9 @@ def unrunnable(key):
              'waiting': waiting, 'waiting_data': waiting_data,
              'who_has': who_has, 'processing': processing, 'stacks': stacks,
              'finished_results': finished_results, 'released': new_released,
-             'in_play': in_play, 'who_wants': {}, 'wants_what': {}}
+             'in_play': in_play, 'who_wants': {}, 'wants_what': {},
+             'ready': ready}
+
     validate_state(**output)
     return output
 
@@ -1427,58 +1485,6 @@ def unrunnable(key):
 _round_robin = [0]
 
 
-def assign_many_tasks(dependencies, waiting, keyorder, who_has, stacks,
-        restrictions, loose_restrictions, nbytes, keys):
-    """ Assign many new ready tasks to workers
-
-    Often at the beginning of computation we have to assign many new leaves to
-    several workers.  This initial seeding of work can have dramatic effects on
-    total runtime.
-
-    This takes some typical state variables (dependencies, waiting, keyorder,
-    who_has, stacks) as well as a list of keys to assign to stacks.
-
-    This mutates waiting and stacks in place and returns a dictionary,
-    new_stacks, that serves as a diff between the old and new stacks.  These
-    new tasks have yet to be put on worker queues.
-    """
-    leaves = list()  # ready tasks without data dependencies
-    ready = list()   # ready tasks with data dependencies
-    new_stacks = defaultdict(list)
-
-    for k in keys:
-        assert not waiting.pop(k)
-        if not dependencies[k] and k not in restrictions:
-            leaves.append(k)
-        else:
-            ready.append(k)
-
-    if not stacks:
-        raise ValueError("No workers found")
-
-    leaves = sorted(leaves, key=keyorder.get)
-
-    workers = list(stacks)
-
-    k = _round_robin[0] % len(workers)
-    workers = workers[k:] + workers[:k]
-    _round_robin[0] += 1
-
-    k = int(ceil(len(leaves) / len(workers)))
-    for i, worker in enumerate(workers):
-        keys = leaves[i*k: (i + 1)*k][::-1]
-        new_stacks[worker].extend(keys)
-        stacks[worker].extend(keys)
-
-    for key in ready:
-        worker = decide_worker(dependencies, stacks, who_has, restrictions,
-                loose_restrictions, nbytes, key)
-        new_stacks[worker].append(key)
-        stacks[worker].append(key)
-
-    return new_stacks
-
-
 def heal_missing_data(tasks, dependencies, dependents,
                       who_has, in_play, waiting, waiting_data, missing):
     """ Return to healthy state after discovering missing data
@@ -1487,6 +1493,9 @@ def heal_missing_data(tasks, dependencies, dependents,
     evaluate those keys.
     """
     logger.debug("Healing missing: %s", missing)
+
+    ready_to_run = set()
+
     for key in missing:
         if key in in_play:
             in_play.remove(key)
@@ -1497,8 +1506,13 @@ def ensure_key(key):
         for dep in dependencies[key]:
             ensure_key(dep)
             waiting_data[dep].add(key)
-        waiting[key] = {dep for dep in dependencies[key] if not who_has.get(dep)}
-        logger.debug("Added key to waiting: %s", key)
+        s = {dep for dep in dependencies[key] if not who_has.get(dep)}
+        if s:
+            waiting[key] = s
+            logger.debug("Added key to waiting: %s", key)
+        else:
+            ready_to_run.add(key)
+            logger.debug("Added key to ready-to-run: %s", key)
         waiting_data[key] = {dep for dep in dependents[key] if dep in in_play
                                                     and not who_has.get(dep)}
         in_play.add(key)
@@ -1508,6 +1522,8 @@ def ensure_key(key):
 
     assert set(missing).issubset(in_play)
 
+    return ready_to_run
+
 
 def _maybe_complex(task):
     """ Possibly contains a nested task """
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -190,12 +190,12 @@ def compute(self, stream, function=None, key=None, args=(), kwargs={},
             task=None, needed=[], who_has=None, report=True, serialized=False):
         """ Execute function """
         self.active.add(key)
-        if needed:
-            local_data = {k: self.data[k] for k in needed if k in self.data}
-            needed = [n for n in needed if n not in self.data]
-        elif who_has:
+        if who_has:
             local_data = {k: self.data[k] for k in who_has if k in self.data}
-            who_has = {k: v for k, v in who_has.items() if k not in self.data}
+            who_has2 = {k: v for k, v in who_has.items() if k not in self.data}
+        elif needed:
+            local_data = {k: self.data[k] for k in needed if k in self.data}
+            needed2 = [n for n in needed if n not in self.data]
         else:
             local_data = {}
 
@@ -204,13 +204,13 @@ def compute(self, stream, function=None, key=None, args=(), kwargs={},
             try:
                 if who_has:
                     logger.info("gather %d keys from peers: %s",
-                                len(who_has), str(who_has))
-                    other = yield gather_from_workers(who_has)
+                                len(who_has2), str(who_has2))
+                    other = yield gather_from_workers(who_has2)
                 elif needed:
                     logger.info("gather %d keys from peers: %s",
-                                len(needed), str(needed))
-                    other = yield _gather(self.center, needed=needed)
-                    other = dict(zip(needed, other))
+                                len(needed2), str(needed2))
+                    other = yield _gather(self.center, needed=needed2)
+                    other = dict(zip(needed2, other))
                 else:
                     raise ValueError()
             except KeyError as e:


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
@@ -1103,17 +1103,21 @@ def test_directed_scatter_sync(loop):
 
 def test_iterator_scatter(loop):
     with cluster() as (s, [a, b]):
-        with Executor(('127.0.0.1', s['port']), loop=loop) as ee:
+        with Executor(('127.0.0.1', s['port']), loop=loop) as e:
 
-            aa = ee.scatter([1,2,3])
-            assert [1,2,3] == ee.gather(aa)
+            aa = e.scatter([1,2,3])
+            assert [1,2,3] == e.gather(aa)
 
             g = (i for i in range(10))
-            futures = ee.scatter(g)
+            futures = e.scatter(g)
             assert isinstance(futures, Iterator)
 
             a = next(futures)
-            assert ee.gather(a) == 0
+            assert e.gather(a) == 0
+
+            futures = list(futures)
+            assert len(futures) == 9
+            assert e.gather(futures) == [1, 2, 3, 4, 5, 6, 7, 8, 9]
 
 
 def test_queue_scatter(loop):
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -1,4 +1,4 @@
-from collections import defaultdict
+from collections import defaultdict, deque
 from copy import deepcopy
 from operator import add
 from time import time
@@ -8,20 +8,25 @@
 from tornado.queues import Queue
 from tornado.iostream import StreamClosedError
 from tornado.gen import TimeoutError
+from tornado import gen
+
 import pytest
 
 from distributed import Center, Nanny, Worker
 from distributed.core import connect, read, write, rpc, loads
 from distributed.client import WrappedKey
-from distributed.scheduler import (validate_state, heal, update_state,
-        decide_worker, assign_many_tasks, heal_missing_data, Scheduler,
+from distributed.scheduler import (validate_state, heal,
+        decide_worker, heal_missing_data, Scheduler,
         _maybe_complex, dumps_function, dumps_task, apply)
-from distributed.utils_test import inc, ignoring, dec
+from distributed.utils_test import (inc, ignoring, dec, gen_cluster, gen_test,
+        loop)
+from distributed.utils import All
 
 
 alice = 'alice'
 bob = 'bob'
 
+
 def test_heal():
     dsk = {'x': 1, 'y': (inc, 'x')}
     dependencies = {'x': set(), 'y': {'x'}}
@@ -31,14 +36,14 @@ def test_heal():
     stacks = {'alice': [], 'bob': []}
     processing = {'alice': set(), 'bob': set()}
 
-    waiting = {'x': set(), 'y': {'x'}}
+    waiting = {'y': {'x'}}
+    ready = {'x'}
     waiting_data = {'x': {'y'}, 'y': set()}
     finished_results = set()
 
     local = {k: v for k, v in locals().items() if '@' not in k}
 
-    output = heal(dependencies, dependents, who_has, stacks, processing,
-            {}, {})
+    output = heal(**local)
 
     assert output['dependencies'] == dependencies
     assert output['dependents'] == dependents
@@ -52,7 +57,7 @@ def test_heal():
     state = {'who_has': dict(),
              'stacks': {'alice': ['x'], 'bob': []},
              'processing': {'alice': set(), 'bob': set()},
-             'waiting': {}, 'waiting_data': {}}
+             'waiting': {}, 'waiting_data': {}, 'ready': set()}
 
     heal(dependencies, dependents, **state)
 
@@ -71,10 +76,11 @@ def test_heal_2():
     state = {'who_has': {'y': {alice}, 'a': {alice}},  # missing 'b'
              'stacks': {'alice': ['z'], 'bob': []},
              'processing': {'alice': set(), 'bob': set(['c'])},
-             'waiting': {}, 'waiting_data': {}}
+             'waiting': {}, 'waiting_data': {}, 'ready': set()}
 
     output = heal(dependencies, dependents, **state)
-    assert output['waiting'] == {'b': set(), 'c': {'b'}, 'result': {'c', 'z'}}
+    assert output['waiting'] == {'c': {'b'}, 'result': {'c', 'z'}}
+    assert output['ready'] == {'b'}
     assert output['waiting_data'] == {'a': {'b'}, 'b': {'c'}, 'c': {'result'},
                                       'y': {'z'}, 'z': {'result'},
                                       'result': set()}
@@ -92,7 +98,7 @@ def test_heal_restarts_leaf_tasks():
     state = {'who_has': dict(),  # missing 'b'
              'stacks': {'alice': ['a'], 'bob': ['x']},
              'processing': {'alice': set(), 'bob': set()},
-             'waiting': {}, 'waiting_data': {}}
+             'waiting': {}, 'waiting_data': {}, 'ready': set()}
 
     del state['stacks']['bob']
     del state['processing']['bob']
@@ -109,7 +115,7 @@ def test_heal_culls():
     state = {'who_has': {'c': {alice}, 'y': {alice}},
              'stacks': {'alice': ['a'], 'bob': []},
              'processing': {'alice': set(), 'bob': set('y')},
-             'waiting': {}, 'waiting_data': {}}
+             'waiting': {}, 'waiting_data': {}, 'ready': set()}
 
     output = heal(dependencies, dependents, **state)
     assert 'a' not in output['stacks']['alice']
@@ -117,162 +123,156 @@ def test_heal_culls():
     assert output['finished_results'] == {'c'}
     assert 'y' not in output['processing']['bob']
 
-    assert output['waiting']['z'] == set()
+    assert output['ready'] == {'z'}
 
 
-def test_update_state():
-    dsk = {'x': 1, 'y': (inc, 'x')}
-    dependencies = {'x': set(), 'y': {'x'}}
-    dependents = {'x': {'y'}, 'y': set()}
+@gen_cluster()
+def test_ready_add_worker(s, a, b):
+    s.add_client(client='client')
+    s.add_worker(address='alice')
 
-    waiting = {'y': set()}
-    waiting_data = {'x': {'y'}, 'y': set()}
+    s.update_graph(tasks={'x-%d' % i: dumps_task((inc, i)) for i in range(20)},
+                   keys=['x-%d' % i for i in range(20)],
+                   client='client',
+                   dependencies={'x-%d' % i: set() for i in range(20)})
 
-    who_wants = defaultdict(set, y={'client'})
-    wants_what = defaultdict(set, client={'y'})
-    who_has = {'x': {'alice'}}
-    processing = set()
-    released = set()
-    in_play = {'x', 'y'}
+def test_update_state(loop):
+    s = Scheduler()
+    s.start()
+    s.add_worker(address='alice', ncores=1)
+    s.update_graph(tasks={'x': 1, 'y': (inc, 'x')},
+                   keys=['y'],
+                   dependencies={'y': 'x', 'x': set()},
+                   client='client')
 
-    new_dsk = {'a': 1, 'z': (add, 'y', 'a')}
-    new_dependencies = {'z': {'y', 'a'}}
-    new_keys = {'z'}
+    s.mark_task_finished('x', 'alice', nbytes=10, type=int)
 
-    e_dsk = {'x': 1, 'y': (inc, 'x'), 'a': 1, 'z': (add, 'y', 'a')}
-    e_dependencies = {'x': set(), 'a': set(), 'y': {'x'}, 'z': {'a', 'y'}}
-    e_dependents = {'z': set(), 'y': {'z'}, 'a': {'z'}, 'x': {'y'}}
+    assert s.processing['alice'] == {'y'}
+    assert not s.ready
+    assert s.who_wants == {'y': {'client'}}
+    assert s.wants_what == {'client': {'y'}}
 
-    e_waiting = {'y': set(), 'a': set(), 'z': {'a', 'y'}}
-    e_waiting_data = {'x': {'y'}, 'y': {'z'}, 'a': {'z'}, 'z': set()}
+    s.update_graph(tasks={'a': 1, 'z': (add, 'y', 'a')},
+                   keys=['z'],
+                   dependencies={'z': {'y', 'a'}},
+                   client='client')
 
-    e_who_wants = {'z': {'client'}, 'y': {'client'}}
-    e_wants_what = {'client': {'y', 'z'}}
 
-    update_state(dsk, dependencies, dependents, who_wants, wants_what,
-                 who_has, in_play,
-                 waiting, waiting_data, new_dsk, new_keys, new_dependencies,
-                 'client')
+    assert s.tasks == {'x': 1, 'y': (inc, 'x'), 'a': 1, 'z': (add, 'y', 'a')}
+    assert s.dependencies == {'x': set(), 'a': set(), 'y': {'x'}, 'z': {'a', 'y'}}
+    assert s.dependents == {'z': set(), 'y': {'z'}, 'a': {'z'}, 'x': {'y'}}
 
-    assert dsk == e_dsk
-    assert dependencies == e_dependencies
-    assert dependents == e_dependents
-    assert waiting == e_waiting
-    assert waiting_data == e_waiting_data
-    assert who_wants == e_who_wants
-    assert wants_what == e_wants_what
-    assert in_play == {'x', 'y', 'a', 'z'}
+    assert s.waiting == {'z': {'a', 'y'}}
+    assert s.waiting_data == {'x': {'y'}, 'y': {'z'}, 'a': {'z'}, 'z': set()}
 
+    assert s.who_wants == {'z': {'client'}, 'y': {'client'}}
+    assert s.wants_what == {'client': {'y', 'z'}}
 
-def test_update_state_with_processing():
-    dsk = {'x': 1, 'y': (inc, 'x'), 'z': (inc, 'y')}
-    dependencies, dependents = get_deps(dsk)
+    assert list(s.ready) == ['a']
+    assert s.in_play == {'a', 'x', 'y', 'z'}
 
-    waiting = {'z': {'y'}}
-    waiting_data = {'x': {'y'}, 'y': {'z'}, 'z': set()}
+    s.stop()
 
-    who_wants = defaultdict(set, z={'client'})
-    wants_what = defaultdict(set, client={'z'})
-    who_has = {'x': {'alice'}}
-    processing = {'y'}
-    released = set()
-    in_play = {'z', 'x', 'y'}
 
-    new_dsk = {'a': (inc, 'x'), 'b': (add, 'a', 'y'), 'c': (inc, 'z')}
-    new_dependencies = {'a': {'x'}, 'b': {'a', 'y'}, 'c': {'z'}}
-    new_keys = {'b', 'c'}
+def test_update_state_with_processing(loop):
+    s = Scheduler()
+    s.start()
+    s.add_worker(address='alice', ncores=1)
+    s.update_graph(tasks={'x': 1, 'y': (inc, 'x'), 'z': (inc, 'y')},
+                   keys=['z'],
+                   dependencies={'y': {'x'}, 'x': set(), 'z': {'y'}},
+                   client='client')
 
-    e_waiting = {'z': {'y'}, 'a': set(), 'b': {'a', 'y'}, 'c': {'z'}}
-    e_waiting_data = {'x': {'y', 'a'}, 'y': {'z', 'b'}, 'z': {'c'},
-                      'a': {'b'}, 'b': set(), 'c': set()}
+    s.mark_task_finished('x', 'alice', nbytes=10, type=int)
 
-    e_who_wants = {'b': {'client'}, 'c': {'client'}, 'z': {'client'}}
-    e_wants_what = {'client': {'b', 'c', 'z'}}
+    assert s.waiting == {'z': {'y'}}
+    assert s.waiting_data == {'x': {'y'}, 'y': {'z'}, 'z': set()}
+    assert list(s.ready) == []
 
-    update_state(dsk, dependencies, dependents, who_wants, wants_what,
-                 who_has, in_play,
-                 waiting, waiting_data, new_dsk, new_keys, new_dependencies,
-                 'client')
+    assert s.who_wants == {'z': {'client'}}
+    assert s.wants_what == {'client': {'z'}}
 
-    assert waiting == e_waiting
-    assert waiting_data == e_waiting_data
-    assert who_wants == e_who_wants
-    assert wants_what == e_wants_what
-    assert in_play == {'x', 'y', 'z', 'a', 'b', 'c'}
+    assert s.who_has == {'x': {'alice'}}
+    assert s.in_play == {'z', 'x', 'y'}
 
+    s.update_graph(tasks={'a': (inc, 'x'), 'b': (add,'a','y'), 'c': (inc, 'z')},
+                   keys=['b', 'c'],
+                   dependencies={'a': {'x'}, 'b': {'a', 'y'}, 'c': {'z'}},
+                   client='client')
 
-def test_update_state_respects_data_in_memory():
-    dsk = {'x': 1, 'y': (inc, 'x')}
-    dependencies, dependents = get_deps(dsk)
+    assert s.waiting == {'z': {'y'}, 'b': {'a', 'y'}, 'c': {'z'}}
+    assert s.stacks['alice'] == ['a']
+    assert not s.ready
+    assert s.waiting_data == {'x': {'y', 'a'}, 'y': {'z', 'b'}, 'z': {'c'},
+                              'a': {'b'}, 'b': set(), 'c': set()}
 
-    waiting = dict()
-    waiting_data = {'y': set()}
+    assert s.who_wants == {'b': {'client'}, 'c': {'client'}, 'z': {'client'}}
+    assert s.wants_what == {'client': {'b', 'c', 'z'}}
+    assert s.in_play == {'x', 'y', 'z', 'a', 'b', 'c'}
 
-    who_wants = defaultdict(set, y={'client'})
-    wants_what = defaultdict(set, client={'y'})
-    who_has = {'y': {'alice'}}
-    processing = set()
-    released = {'x'}
-    in_play = {'y'}
+    s.stop()
 
-    new_dsk = {'x': 1, 'y': (inc, 'x'), 'z': (add, 'y', 'x')}
-    new_dependencies = {'y': {'x'}, 'z': {'y', 'x'}}
-    new_keys = {'z'}
 
-    e_dsk = new_dsk.copy()
-    e_waiting = {'z': {'x'}, 'x': set()}
-    e_waiting_data = {'x': {'z'}, 'y': {'z'}, 'z': set()}
-    e_who_wants = {'y': {'client'}, 'z': {'client'}}
-    e_wants_what = {'client': {'y', 'z'}}
+def test_update_state_respects_data_in_memory(loop):
+    s = Scheduler()
+    s.start()
+    s.add_worker(address='alice', ncores=1)
+    s.update_graph(tasks={'x': 1, 'y': (inc, 'x')},
+                   keys=['y'],
+                   dependencies={'y': {'x'}, 'x': set()},
+                   client='client')
 
-    update_state(dsk, dependencies, dependents, who_wants, wants_what,
-                 who_has, in_play,
-                 waiting, waiting_data, new_dsk, new_keys, new_dependencies,
-                 'client')
+    s.mark_task_finished('x', 'alice', nbytes=10, type=int)
+    s.mark_task_finished('y', 'alice', nbytes=10, type=int)
 
-    assert dsk == e_dsk
-    assert waiting == e_waiting
-    assert waiting_data == e_waiting_data
-    assert who_wants == e_who_wants
-    assert wants_what == e_wants_what
-    assert in_play == {'x', 'y', 'z'}
+    assert s.who_has == {'y': {'alice'}}
 
+    s.update_graph(tasks={'x': 1, 'y': (inc, 'x'), 'z': (add, 'y', 'x')},
+                   keys=['z'],
+                   dependencies={'y': {'x'}, 'z': {'y', 'x'}},
+                   client='client')
 
-def test_update_state_supports_recomputing_released_results():
-    dsk = {'x': 1, 'y': (inc, 'x'), 'z': (inc, 'x')}
-    dependencies, dependents = get_deps(dsk)
+    assert s.waiting == {'z': {'x'}}
+    assert s.processing['alice'] == {'x'}  # x was released, need to recompute
+    assert s.waiting_data == {'x': {'z'}, 'y': {'z'}, 'z': set()}
+    assert s.who_wants == {'y': {'client'}, 'z': {'client'}}
+    assert s.wants_what == {'client': {'y', 'z'}}
+    assert s.in_play == {'x', 'y', 'z'}
 
-    waiting = dict()
-    waiting_data = {'z': set()}
+    s.stop()
 
-    who_wants = defaultdict(set, z={'client'})
-    wants_what = defaultdict(set, client={'z'})
-    who_has = {'z': {'alice'}}
-    processing = set()
-    released = {'x', 'y'}
-    in_play = {'z'}
 
-    new_dsk = {'x': 1, 'y': (inc, 'x')}
-    new_dependencies = {'y': {'x'}}
-    new_keys = {'y'}
+def test_update_state_supports_recomputing_released_results(loop):
+    s = Scheduler()
+    s.start()
+    s.add_worker(address='alice', ncores=1)
+    s.update_graph(tasks={'x': 1, 'y': (inc, 'x'), 'z': (inc, 'x')},
+                   keys=['z'],
+                   dependencies={'y': {'x'}, 'x': set(), 'z': {'y'}},
+                   client='client')
+
+    s.mark_task_finished('x', 'alice', nbytes=10, type=int)
+    s.mark_task_finished('y', 'alice', nbytes=10, type=int)
+    s.mark_task_finished('z', 'alice', nbytes=10, type=int)
 
-    e_dsk = dsk.copy()
-    e_waiting = {'x': set(), 'y': {'x'}}
-    e_waiting_data = {'x': {'y'}, 'y': set(), 'z': set()}
-    e_who_wants = {'z': {'client'}, 'y': {'client'}}
-    e_wants_what = {'client': {'y', 'z'}}
+    assert not s.waiting
+    assert not s.ready
+    assert s.waiting_data == {'z': set()}
 
-    update_state(dsk, dependencies, dependents, who_wants, wants_what,
-                 who_has, in_play,
-                 waiting, waiting_data, new_dsk, new_keys, new_dependencies,
-                 'client')
+    assert s.who_has == {'z': {'alice'}}
 
-    assert dsk == e_dsk
-    assert waiting == e_waiting
-    assert waiting_data == e_waiting_data
-    assert who_wants == e_who_wants
-    assert wants_what == e_wants_what
-    assert in_play == {'x', 'y', 'z'}
+    s.update_graph(tasks={'x': 1, 'y': (inc, 'x')},
+                   keys=['y'],
+                   dependencies={'y': {'x'}},
+                   client='client')
+
+    assert s.waiting == {'y': {'x'}}
+    assert s.waiting_data == {'x': {'y'}, 'y': set(), 'z': set()}
+    assert s.who_wants == {'z': {'client'}, 'y': {'client'}}
+    assert s.wants_what == {'client': {'y', 'z'}}
+    assert s.processing['alice'] == {'x'}
+
+    s.stop()
 
 
 def test_decide_worker_with_many_independent_leaves():
@@ -345,14 +345,14 @@ def test_decide_worker_with_loose_restrictions():
 
 
 def test_decide_worker_without_stacks():
-    with pytest.raises(ValueError):
-        result = decide_worker({'x': []}, [], {}, {}, set(), {}, 'x')
+    assert not decide_worker({'x': []}, [], {}, {}, set(), {}, 'x')
 
 
 def test_validate_state():
     dsk = {'x': 1, 'y': (inc, 'x')}
     dependencies = {'x': set(), 'y': {'x'}}
-    waiting = {'y': {'x'}, 'x': set()}
+    waiting = {'y': {'x'}}
+    ready = deque(['x'])
     dependents = {'x': {'y'}, 'y': set()}
     waiting_data = {'x': {'y'}}
     who_has = dict()
@@ -370,18 +370,23 @@ def test_validate_state():
     with pytest.raises(Exception):
         validate_state(**locals())
 
-    del waiting['x']
+    ready.remove('x')
     with pytest.raises(Exception):
         validate_state(**locals())
 
     waiting['y'].remove('x')
+    with pytest.raises(Exception):
+        validate_state(**locals())
+
+    del waiting['y']
+    ready.appendleft('y')
     validate_state(**locals())
 
     stacks['alice'].append('y')
     with pytest.raises(Exception):
         validate_state(**locals())
 
-    waiting.pop('y')
+    ready.remove('y')
     validate_state(**locals())
 
     stacks['alice'].pop()
@@ -403,50 +408,6 @@ def test_validate_state():
     validate_state(**locals())
 
 
-def test_assign_many_tasks():
-    alice, bob, charlie = ('alice', 8000), ('bob', 8000), ('charlie', 8000)
-    dependencies = {'y': {'x'}, 'b': {'a'}, 'x': set(), 'a': set()}
-    waiting = {'y': set(), 'b': {'a'}, 'a': set()}
-    keyorder = {c: 1 for c in 'abxy'}
-    who_has = {'x': {alice}}
-    stacks = {alice: [], bob: []}
-    nbytes = {}
-    restrictions = {}
-    keys = ['y', 'a']
-
-    new_stacks = assign_many_tasks(dependencies, waiting, keyorder, who_has,
-                                   stacks, restrictions, set(), nbytes, ['y', 'a'])
-
-    assert 'y' in stacks[alice]
-    assert 'a' in stacks[alice] + stacks[bob]
-    assert 'a' not in waiting
-    assert 'y' not in waiting
-
-    assert set(concat(new_stacks.values())) == set(concat(stacks.values()))
-
-
-def test_assign_many_tasks_with_restrictions():
-    alice, bob, charlie = ('alice', 8000), ('bob', 8000), ('charlie', 8000)
-    dependencies = {'y': {'x'}, 'b': {'a'}, 'x': set(), 'a': set()}
-    waiting = {'y': set(), 'b': {'a'}, 'a': set()}
-    keyorder = {c: 1 for c in 'abxy'}
-    who_has = {'x': {alice}}
-    stacks = {alice: [], bob: []}
-    nbytes = {'x': 0}
-    restrictions = {'y': {'bob'}, 'a': {'alice'}}
-    keys = ['y', 'a']
-
-    new_stacks = assign_many_tasks(dependencies, waiting, keyorder, who_has,
-                                   stacks, restrictions, set(), nbytes, ['y', 'a'])
-
-    assert 'y' in stacks[bob]
-    assert 'a' in stacks[alice]
-    assert 'a' not in waiting
-    assert 'y' not in waiting
-
-    assert set(concat(new_stacks.values())) == set(concat(stacks.values()))
-
-
 def test_fill_missing_data():
     dsk = {'x': 1, 'y': (inc, 'x'), 'z': (inc, 'y')}
     dependencies, dependents = get_deps(dsk)
@@ -461,7 +422,8 @@ def test_fill_missing_data():
     released = set()
     in_play = {'z'}
 
-    e_waiting = {'x': set(), 'y': {'x'}, 'z': {'y'}}
+    e_waiting = {'y': {'x'}, 'z': {'y'}}
+    e_ready = {'x'}
     e_waiting_data = {'x': {'y'}, 'y': {'z'}, 'z': set()}
     e_in_play = {'x', 'y', 'z'}
 
@@ -469,18 +431,15 @@ def test_fill_missing_data():
     del who_has['z']
     in_play.remove('z')
 
-    heal_missing_data(dsk, dependencies, dependents, who_has, in_play, waiting,
+    ready = heal_missing_data(dsk, dependencies, dependents, who_has, in_play, waiting,
             waiting_data, lost)
 
     assert waiting == e_waiting
     assert waiting_data == e_waiting_data
+    assert ready == e_ready
     assert in_play == e_in_play
 
 
-from distributed.utils_test import gen_cluster, gen_test
-from distributed.utils import All
-from tornado import gen
-
 def div(x, y):
     return x / y
 
@@ -623,12 +582,13 @@ def test_remove_worker_from_scheduler(s, a, b):
     dsk = {('x', i): (inc, i) for i in range(10)}
     s.update_graph(tasks=valmap(dumps_task, dsk), keys=list(dsk),
                    dependencies={k: set() for k in dsk})
-    assert s.stacks[a.address]
+    assert s.ready
+    assert not any(stack for stack in s.stacks.values())
 
     assert a.address in s.worker_queues
     s.remove_worker(address=a.address)
     assert a.address not in s.ncores
-    assert len(s.stacks[b.address]) + len(s.processing[b.address]) == \
+    assert len(s.ready) + len(s.processing[b.address]) == \
             len(dsk)  # b owns everything
     assert all(k in s.in_play for k in dsk)
     s.validate()
@@ -724,7 +684,8 @@ def test_scheduler_as_center():
         yield gen.sleep(0.01)
     assert 'a' in a.data or 'a' in b.data or 'a' in c.data
 
-    yield [w._close() for w in [a, b, c]]
+    with ignoring(StreamClosedError):
+        yield [w._close() for w in [a, b, c]]
 
     assert s.ncores == {}
     assert s.who_has == {}
@@ -848,3 +809,73 @@ def test_dumps_task():
     assert loads(d['function'])(1, 2) == 3
     assert loads(d['args']) == (1,)
     assert loads(d['kwargs']) == {'y': 10}
+
+
+@gen_cluster()
+def test_ready_remove_worker(s, a, b):
+    s.add_client(client='client')
+    s.update_graph(tasks={'x-%d' % i: dumps_task((inc, i)) for i in range(20)},
+                   keys=['x-%d' % i for i in range(20)],
+                   client='client',
+                   dependencies={'x-%d' % i: set() for i in range(20)})
+
+    assert all(len(s.processing[w]) == s.ncores[w]
+                for w in s.ncores)
+    assert not any(stack for stack in s.stacks.values())
+    assert len(s.ready) + sum(map(len, s.processing.values())) == 20
+
+    s.remove_worker(address=a.address)
+
+    for collection in [s.ncores, s.stacks, s.processing]:
+        assert set(collection) == {b.address}
+    assert all(len(s.processing[w]) == s.ncores[w]
+                for w in s.ncores)
+    assert set(s.processing) == {b.address}
+    assert not any(stack for stack in s.stacks.values())
+    assert len(s.ready) + sum(map(len, s.processing.values())) == 20
+
+
+@gen_cluster(Worker=Nanny)
+def test_restart(s, a, b):
+    s.add_client(client='client')
+    s.update_graph(tasks={'x-%d' % i: dumps_task((inc, i)) for i in range(20)},
+                   keys=['x-%d' % i for i in range(20)],
+                   client='client',
+                   dependencies={'x-%d' % i: set() for i in range(20)})
+
+    assert len(s.ready) + sum(map(len, s.processing.values())) == 20
+    assert s.ready
+
+    yield s.restart()
+
+    for c in [s.stacks, s.processing, s.ncores]:
+        assert len(c) == 2
+
+    for c in [s.stacks, s.processing]:
+        assert not any(v for v in c.values())
+
+    assert not s.ready
+    assert not s.tasks
+    assert not s.dependencies
+
+
+@gen_cluster()
+def test_ready_add_worker(s, a, b):
+    s.add_client(client='client')
+    s.update_graph(tasks={'x-%d' % i: dumps_task((inc, i)) for i in range(20)},
+                   keys=['x-%d' % i for i in range(20)],
+                   client='client',
+                   dependencies={'x-%d' % i: set() for i in range(20)})
+
+    assert all(len(s.processing[w]) == s.ncores[w]
+                for w in s.ncores)
+    assert len(s.ready) + sum(map(len, s.processing.values())) == 20
+
+    w = Worker(s.ip, s.port, ncores=3, ip='127.0.0.1')
+    w.listen(0)
+    s.add_worker(address=w.address, ncores=w.ncores)
+
+    assert w.address in s.ncores
+    assert all(len(s.processing[w]) == s.ncores[w]
+                for w in s.ncores)
+    assert len(s.ready) + sum(map(len, s.processing.values())) == 20