Repository: dask/distributed

Issue:
KeyError in `Worker.handle_compute_task` (causes deadlock)
```
Traceback (most recent call last):
  File "/opt/conda/envs/coiled/lib/python3.9/site-packages/distributed/worker.py", line 1237, in handle_scheduler
    await self.handle_stream(
  File "/opt/conda/envs/coiled/lib/python3.9/site-packages/distributed/core.py", line 564, in handle_stream
    handler(**merge(extra, msg))
  File "/opt/conda/envs/coiled/lib/python3.9/site-packages/distributed/worker.py", line 1937, in handle_compute_task
    self.tasks[key].nbytes = value
KeyError: "('slices-40d6794b777d639e9440f8f518224cfd', 2, 1)"
distributed.worker - INFO - Connection to scheduler broken.  Reconnecting...
```

I may be able to reproduce this if necessary. I was running a stackstac example notebook [on binder](https://mybinder.org/v2/gh/gjoseph92/stackstac/use-binder?urlpath=lab/tree/examples/cluster.ipynb) against a Coiled cluster over `wss`, where the particular versions of things were causing a lot of errors (unrelated to dask). So I was frequently rerunning the same tasks, cancelling them, restarting the client, rerunning, etc. Perhaps this cancelling, restarting, rerunning is related?

@fjetter  says
> The only reason this keyerror can appear is if the compute instruction transitions its own dependency into a _forgotten_ state which is very, very wrong.

Relevant code, ending at the line where the error occurs: https://github.com/dask/distributed/blob/11c41b525267cc144caa8d077b84a0939fecae97/distributed/worker.py#L1852-L1937

Scheduler code producing the message which causes this error:
https://github.com/dask/distributed/blob/11c41b525267cc144caa8d077b84a0939fecae97/distributed/scheduler.py#L7953-L7994


Discussion:
> So I was frequently rerunning the same tasks, cancelling them, restarting the client, rerunning, etc. Perhaps this cancelling, restarting, rerunning is related?

Probably. What you're describing is basically what motivated https://github.com/dask/distributed/pull/5046, see specifically https://github.com/dask/distributed/pull/5046#issuecomment-895305582 and to some extend https://github.com/dask/distributed/issues/4413#issuecomment-817820890 although the latter is a bit outdated. I'll need to update both (xref https://github.com/dask/distributed/issues/5413)

The code should be much more robust to this situation but I guess I missed an edge case. Therefore, a transition log of the deadlocked worker would be very valuable.
> @fjetter says
> 
> > The only reason this keyerror can appear is if the compute instruction transitions its own dependency into a forgotten state which is very, very wrong.

Explanation below

1. Schedule ensures that `msg["who_has"].keys() == msg["nbytes"].keys()`

https://github.com/dask/distributed/blob/11c41b525267cc144caa8d077b84a0939fecae97/distributed/scheduler.py#L7977-L7980

2. Worker ensures there exists a `TaskState` in `Worker.tasks` for every key in `who_has`

https://github.com/dask/distributed/blob/76495965cf8d3fb5f54bb4b8d20279ae402e0957/distributed/worker.py#L1907-L1916

3. We perform a bunch of transitions. This is the only place where can actually drop another task again.

https://github.com/dask/distributed/blob/76495965cf8d3fb5f54bb4b8d20279ae402e0957/distributed/worker.py#L1937

4. This is where the KeyError is raised.

https://github.com/dask/distributed/blob/76495965cf8d3fb5f54bb4b8d20279ae402e0957/distributed/worker.py#L1943-L1944


Since 1.) ensures `nbytes.keys() == who_has.keys()` and 2.) ensure there exists a task for all keys in who_has, therefore it ensures there is a task for every key in nbytes, 3. is the only way we could drop/forget a key since 4. is still in the same synchronous block.

The transition logic for that broken key could be either of

A) released -> waiting -> ready -> executing -> cancelled(executing)
B) released -> fetch -> flight -> cancelled(flight)

Above 2.) should then be a no-op 3.) should then trigger a 

A) `cancelled->fetch` transitions to `flight`
B) `cancelled->fetch` transitions to `resumed(flight)`. This one triggers a `release_key` and is as far as I can see the only source which could generate a `forgotten` recommendation.

I _think_ this is some not well covered edge case in https://github.com/dask/distributed/blame/76495965cf8d3fb5f54bb4b8d20279ae402e0957/distributed/worker.py#L2965-L2966

As I said, a transition log would be helpful since it would remove some of the hypothetical of this argumentation. It should basically just show us the series of decisions leading up to the event.
At the request of @gjoseph92  I am posting my transition logs as a result of a deadlock:

[log.txt](https://github.com/dask/distributed/files/7470870/log.txt)



^ For background @bennnym ran into this same `KeyError`:
```
distributed.worker - INFO - -------------------------------------------------
distributed.worker - INFO - Run out-of-band function 'lambda'
distributed.worker - ERROR - "('split-shuffle-1-80cde6aa3e7938b28ad0dfe2387b063e', 5, (2, 3))"
Traceback (most recent call last):
  File "/opt/conda/envs/coiled/lib/python3.8/site-packages/distributed/worker.py", line 1237, in handle_scheduler
    await self.handle_stream(
  File "/opt/conda/envs/coiled/lib/python3.8/site-packages/distributed/core.py", line 564, in handle_stream
    handler(**merge(extra, msg))
  File "/opt/conda/envs/coiled/lib/python3.8/site-packages/distributed/worker.py", line 1937, in handle_compute_task
    self.tasks[key].nbytes = value
KeyError: "('split-shuffle-1-80cde6aa3e7938b28ad0dfe2387b063e', 5, (2, 3))"
distributed.worker - INFO - Connection to scheduler broken.  Reconnecting...
```
This error then caused a deadlock (https://github.com/dask/distributed/issues/5480). We confirmed it was also due to https://github.com/dask/distributed/issues/5481 by checking for queued messages in the worker's `batched_stream`:
```
>>> client.run(
...     lambda dask_worker: (
...         len(dask_worker.batched_stream.buffer),
...         dask_worker.batched_stream.please_stop
...     )
... )
{'tls://10.12.16.108:35519': (0, False),
 'tls://10.12.16.88:42647': (0, False),
 'tls://10.12.18.155:41501': (0, False),
 'tls://10.12.19.84:40143': (0, False),
 'tls://10.12.20.221:40545': (0, False),
 'tls://10.12.21.210:45285': (0, False),
 'tls://10.12.21.217:35583': (0, False),
 'tls://10.12.21.39:43355': (0, False),
 'tls://10.12.23.137:44833': (0, False),
 'tls://10.12.23.225:36857': (0, False),
 'tls://10.12.23.48:36147': (620, True),  <-- this one has enqueued messages and a stopped `_background_send` coroutine
 'tls://10.12.24.247:42369': (0, False),
 'tls://10.12.26.157:42411': (0, False),
 'tls://10.12.27.196:44821': (0, False),
 'tls://10.12.27.70:40899': (0, False),
 'tls://10.12.28.196:32955': (0, False),
 'tls://10.12.29.12:40313': (0, False),
 'tls://10.12.30.116:45129': (0, False),
 'tls://10.12.30.157:34201': (0, False),
 'tls://10.12.31.58:36025': (0, False)}
```

So we then got transition logs for the key in question, which Ben has attached above.

@bennnym I forgot, could you also post the output of `client.get_versions`?
@gjoseph92  my code was running in a notebook and my kernell died, sorry. That is all I have

Do you know offhand what versions of dask and distributed you were running? It may show in your Coiled software environment.
I was just installing `dask[distributed]` so I assume the latest version


Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -652,7 +652,7 @@ def __init__(
             ("resumed", "waiting"): self.transition_rescheduled_next,
             ("resumed", "fetch"): self.transition_rescheduled_next,
             ("constrained", "executing"): self.transition_constrained_executing,
-            ("constrained", "released"): self.transition_constrained_released,
+            ("constrained", "released"): self.transition_generic_released,
             ("error", "released"): self.transition_generic_released,
             ("executing", "error"): self.transition_executing_error,
             ("executing", "long-running"): self.transition_executing_long_running,
@@ -1810,7 +1810,7 @@ def handle_free_keys(self, comm=None, keys=None, stimulus_id=None):
         for key in keys:
             ts = self.tasks.get(key)
             if ts:
-                recommendations[ts] = "released" if ts.dependents else "forgotten"
+                recommendations[ts] = "released"
 
         self.transitions(recommendations, stimulus_id=stimulus_id)
 
@@ -1843,7 +1843,7 @@ def handle_remove_replicas(self, keys, stimulus_id):
                 continue
             if not ts.is_protected():
                 self.log.append((ts.key, "remove-replica-confirmed", stimulus_id))
-                recommendations[ts] = "released" if ts.dependents else "forgotten"
+                recommendations[ts] = "released"
             else:
                 rejected.append(key)
 
@@ -2022,9 +2022,11 @@ def transition_missing_fetch(self, ts, *, stimulus_id):
 
     def transition_missing_released(self, ts, *, stimulus_id):
         self._missing_dep_flight.discard(ts)
-        recommendations = self.release_key(ts.key, reason="missing->released")
+        recommendations, smsgs = self.transition_generic_released(
+            ts, stimulus_id=stimulus_id
+        )
         assert ts.key in self.tasks
-        return recommendations, []
+        return recommendations, smsgs
 
     def transition_fetch_missing(self, ts, *, stimulus_id):
         # handle_missing will append to self.data_needed if new workers are found
@@ -2041,7 +2043,18 @@ def transition_released_fetch(self, ts, *, stimulus_id):
         return {}, []
 
     def transition_generic_released(self, ts, *, stimulus_id):
-        recs = self.release_key(ts.key, reason=stimulus_id)
+        self.release_key(ts.key, reason=stimulus_id)
+        recs = {}
+        for dependency in ts.dependencies:
+            if (
+                not dependency.waiters
+                and dependency.state not in READY | PROCESSING | {"memory"}
+            ):
+                recs[dependency] = "released"
+
+        if not ts.dependents:
+            recs[ts] = "forgotten"
+
         return recs, []
 
     def transition_released_waiting(self, ts, *, stimulus_id):
@@ -2081,8 +2094,8 @@ def transition_fetch_flight(self, ts, worker, *, stimulus_id):
         return {}, []
 
     def transition_memory_released(self, ts, *, stimulus_id):
-        recs = self.release_key(ts.key, reason=stimulus_id)
-        smsgs = [{"op": "release-worker-data", "key": ts.key}]
+        recs, smsgs = self.transition_generic_released(ts, stimulus_id=stimulus_id)
+        smsgs.append({"op": "release-worker-data", "key": ts.key})
         return recs, smsgs
 
     def transition_waiting_constrained(self, ts, *, stimulus_id):
@@ -2155,9 +2168,9 @@ def transition_executing_error(
 
     def transition_rescheduled_next(self, ts, *, stimulus_id):
         next_state = ts._next
-        recs = self.release_key(ts.key, reason=stimulus_id)
+        recs, smsgs = self.transition_generic_released(ts, stimulus_id=stimulus_id)
         recs[ts] = next_state
-        return recs, []
+        return recs, smsgs
 
     def transition_cancelled_fetch(self, ts, *, stimulus_id):
         if ts.done:
@@ -2200,9 +2213,11 @@ def transition_cancelled_released(self, ts, *, stimulus_id):
 
         for resource, quantity in ts.resource_restrictions.items():
             self.available_resources[resource] += quantity
-        recommendations = self.release_key(ts.key, reason=stimulus_id)
-        recommendations[ts] = next_state or "released"
-        return recommendations, []
+        recommendations, smsgs = self.transition_generic_released(
+            ts, stimulus_id=stimulus_id
+        )
+        recommendations[ts] = next_state or recommendations.get(ts, "released")
+        return recommendations, smsgs
 
     def transition_executing_released(self, ts, *, stimulus_id):
         ts._previous = ts.state
@@ -2243,10 +2258,6 @@ def transition_executing_memory(self, ts, value=no_value, *, stimulus_id):
         self.executed_count += 1
         return self.transition_generic_memory(ts, value=value, stimulus_id=stimulus_id)
 
-    def transition_constrained_released(self, ts, *, stimulus_id):
-        recs = self.release_key(ts.key, reason=stimulus_id)
-        return recs, []
-
     def transition_constrained_executing(self, ts, *, stimulus_id):
         if self.validate:
             assert not ts.waiting_for_data
@@ -2983,7 +2994,7 @@ def handle_steal_request(self, key, stimulus_id):
             # If task is marked as "constrained" we haven't yet assigned it an
             # `available_resources` to run on, that happens in
             # `transition_constrained_executing`
-            self.transition(ts, "forgotten", stimulus_id=stimulus_id)
+            self.transition(ts, "released", stimulus_id=stimulus_id)
 
     def release_key(
         self,
@@ -2991,8 +3002,7 @@ def release_key(
         cause: TaskState | None = None,
         reason: str | None = None,
         report: bool = True,
-    ):
-        recommendations = {}
+    ) -> None:
         try:
             if self.validate:
                 assert not isinstance(key, TaskState)
@@ -3032,9 +3042,6 @@ def release_key(
                 ts.waiting_for_data.discard(d)
                 d.waiters.discard(ts)
 
-                if not d.waiters and d.state in {"flight", "fetch", "missing"}:
-                    recommendations[d] = "forgotten"
-
             ts.waiting_for_data.clear()
             ts.nbytes = None
             ts._previous = None
@@ -3058,8 +3065,6 @@ def release_key(
                 pdb.set_trace()
             raise
 
-        return recommendations
-
     ################
     # Execute Task #
     ################


Test patch: