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:

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

[gforsyth]

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]
 

[gforsyth]

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.

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