Repository: dask/distributed

Issue:
Task dependency bug in worker
**What happened**:
The following code fails sometimes with an `KeyError` error (3 out of 4 times). `git bisect` says the bug was introduced by https://github.com/dask/distributed/pull/4107.
```
distributed.worker - ERROR - "('split-simple-shuffle-b50994a1a48d47067bc463a19b005e75', 14, 55)"
Traceback (most recent call last):
  File "/home/nfs/mkristensen/repos/distributed/distributed/worker.py", line 1984, in gather_dep
    deps_ts = [self.tasks[key] for key in deps]
  File "/home/nfs/mkristensen/repos/distributed/distributed/worker.py", line 1984, in <listcomp>
    deps_ts = [self.tasks[key] for key in deps]
KeyError: "('split-simple-shuffle-b50994a1a48d47067bc463a19b005e75', 14, 55)"
```

**What you expected to happen**:
For some reason the worker calls `gather_dep()` with a set of `deps` that the task does not depend on. As far as I can see, the client and scheduler maintains task dependencies correctly.

**Minimal Complete Verifiable Example**:
```python

import pandas as pd
import numpy as np

import dask.dataframe as dd
from dask.dataframe.shuffle import shuffle
from distributed import wait
import dask

from distributed import Client, LocalCluster

nparts = 100
max_branch = 100
data_size = nparts * max_branch


def main(client):
    df = pd.DataFrame({"x": np.arange(data_size)})
    ddf = dd.from_pandas(df, npartitions=nparts)
    ddf = ddf.persist(optimize_graph=False)
    wait(ddf)

    with dask.config.set({"optimization.fuse.active": False}):
        s = shuffle(
            ddf, ddf.x, shuffle="tasks", npartitions=nparts, max_branch=max_branch
        )
        s = s.persist()
        wait(s)


if __name__ == "__main__":
    with LocalCluster(scheduler_port=0, asynchronous=False, n_workers=5) as cluster:
        with Client(cluster, asynchronous=False) as client:
            main(client)
```
cc. @gforsyth 



Discussion:
Thanks for the reproducer, @madsbk ! I'll take a look
I'm seeing the same thing in the nightly benchmark we are putting together:
https://github.com/quasiben/dask-scheduler-performance/pull/2
I can confirm the issue.  Here's an inelegant fix while I figure out the correct one:

```
diff --git a/distributed/worker.py b/distributed/worker.py
index b9bc9b633..e77fd984a 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1998,7 +1998,7 @@ class Worker(ServerNode):
 
                 # dep states may have changed before gather_dep runs
                 # if a dep is no longer in-flight then don't fetch it
-                deps_ts = [self.tasks[key] for key in deps]
+                deps_ts = [self.tasks.get(key, TaskState(key)) for key in deps]
                 deps_ts = tuple(ts for ts in deps_ts if ts.state == "flight")
                 deps = [d.key for d in deps_ts]
 
```
Checking in.  Is this the right approach?  @gforsyth what is inelegant about this in particular?  It appears to be blocking some benchmarking work as well.

Also, quick nit, I wonder how expensive it is to construct a TaskState object, and if instead we should do something like using `or` here: 

```python 
deps_ts = [self.tasks.get(key, None) or TaskState(key) for key in deps]
```
> what is inelegant about this in particular?

I wanted to confirm that this was due to (I think) work stealing and not a request for the "wrong" dependencies.  Tried to get a test that would reliably fail on this but didn't quite get there.  


> ```python
> deps_ts = [self.tasks.get(key, None) or TaskState(key) for key in deps]
> ```

yes, I think this is the right way to do this, I'll push up the PR now

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1998,7 +1998,7 @@ async def gather_dep(self, worker, dep, deps, total_nbytes, cause=None):
 
                 # dep states may have changed before gather_dep runs
                 # if a dep is no longer in-flight then don't fetch it
-                deps_ts = [self.tasks[key] for key in deps]
+                deps_ts = [self.tasks.get(key, None) or TaskState(key) for key in deps]
                 deps_ts = tuple(ts for ts in deps_ts if ts.state == "flight")
                 deps = [d.key for d in deps_ts]
 


Test patch: