Repository: dask/distributed

Issue:
KeyError in Scheduler.update_graph
Hi, there seems to be an issue with the cleanup logic added to `Scheduler.update_graph` in #1840.

For instance, when I execute

```python
import time

def foo():
    return 'foo'

def bar(x):
    return x + 'bar'

def baz(x):
    time.sleep(100)
    return x + 'baz'

if __name__ == '__main__':
    from distributed import Client, LocalCluster
    cluster = LocalCluster()
    c = Client(cluster)
    x = c.submit(foo)
    y = c.submit(bar, x)
    y.result()
    z = c.submit(baz, y)
    del y
    time.sleep(1)
    f = c.submit(bar, x)
```

I get with distributed 1.22.1 the following error:

```
Traceback (most recent call last):
  File "/home/stephan/.virtualenvs/pymor/lib/python3.5/site-packages/distributed/core.py", line 320, in handle_comm
    result = yield result
  File "/home/stephan/.virtualenvs/pymor/lib/python3.5/site-packages/tornado/gen.py", line 1133, in run
    value = future.result()
  File "/usr/lib/python3.5/asyncio/futures.py", line 293, in result
    raise self._exception
  File "/home/stephan/.virtualenvs/pymor/lib/python3.5/site-packages/tornado/gen.py", line 1147, in run
    yielded = self.gen.send(value)
  File "/home/stephan/.virtualenvs/pymor/lib/python3.5/site-packages/distributed/scheduler.py", line 1924, in add_client
    yield self.handle_stream(comm=comm, extra={'client': client})
  File "/home/stephan/.virtualenvs/pymor/lib/python3.5/site-packages/tornado/gen.py", line 1133, in run
    value = future.result()
  File "/usr/lib/python3.5/asyncio/futures.py", line 293, in result
    raise self._exception
  File "/home/stephan/.virtualenvs/pymor/lib/python3.5/site-packages/tornado/gen.py", line 1147, in run
    yielded = self.gen.send(value)
  File "/home/stephan/.virtualenvs/pymor/lib/python3.5/site-packages/distributed/core.py", line 374, in handle_stream
    handler(**merge(extra, msg))
  File "/home/stephan/.virtualenvs/pymor/lib/python3.5/site-packages/distributed/scheduler.py", line 1357, in update_graph
    for dep in dependencies[key]:
KeyError: 'foo-b55e006af80e0cfdc1a300901102086d'
```

The problem seems to be that the code assumes that each dependency is again a key in the `dependencies` dict, which is not the case in this example.


Discussion:
Thank you for the excellent reproducible minimal example @sdrave .  This is great.

Is this something that you'd like to take a shot at resolving?  If not I can probably look at it in a day or two.

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -1369,7 +1369,11 @@ def update_graph(self, client=None, tasks=None, keys=None,
             while stack:  # remove unnecessary dependencies
                 key = stack.pop()
                 ts = self.tasks[key]
-                for dep in dependencies[key]:
+                try:
+                    deps = dependencies[key]
+                except KeyError:
+                    deps = self.dependencies[key]
+                for dep in deps:
                     if all(d in done for d in dependents[dep]):
                         if dep in self.tasks:
                             done.add(dep)
@@ -1377,7 +1381,7 @@ def update_graph(self, client=None, tasks=None, keys=None,
 
             for d in done:
                 tasks.pop(d, None)
-                del dependencies[d]
+                dependencies.pop(d, None)
 
         # Get or create task states
         stack = list(keys)


Test patch:
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
@@ -1362,3 +1362,25 @@ def test_resources_reset_after_cancelled_task(c, s, w):
     assert w.available_resources == {'A': 1}
 
     yield c.submit(inc, 1, resources={'A': 1})
+
+
+@gen_cluster(client=True)
+def test_gh2187(c, s, a, b):
+    def foo():
+        return 'foo'
+
+    def bar(x):
+        return x + 'bar'
+
+    def baz(x):
+        sleep(0.1)
+        return x + 'baz'
+
+    x = c.submit(foo, key='x')
+    y = c.submit(bar, x, key='y')
+    yield y
+    z = c.submit(baz, y, key='z')
+    del y
+    yield gen.sleep(0.1)
+    f = c.submit(bar, x, key='y')
+    yield f