Repository: dask/distributed

Issue:
Await on handle_stream raises missing delete_data await warning
For increased visibility, I'm reposting https://github.com/dask/distributed/pull/3847/files#r443766556 as an issue here:

We have a few tests in dask-cuda that check the behavior of Device<->Host<->Disk spilling and I noticed after the 2.19 release one of them has broken, I managed to track it down to one specific line of code in https://github.com/dask/distributed/blob/44b2358e33a0738c4c70ca96db4242636245e07d/distributed/core.py#L573, introduced by https://github.com/dask/distributed/pull/3847. The test in question happens in https://github.com/rapidsai/dask-cuda/blob/branch-0.15/dask_cuda/tests/test_spill.py#L409-L411, where we assert that the zict dictionaries are empty after deleting `cdf2`, which is the object being spilled. It seems that this is because we're not awaiting for `Worker.delete_data` somewhere, as per the warning below that doesn't happen if I comment `await gen.sleep(0)` out:

```python
dask_cuda/tests/test_spill.py::test_cudf_device_spill[params0]
  /datasets/pentschev/miniconda3/envs/r-102-0.14/lib/python3.7/inspect.py:732: RuntimeWarning: coroutine 'Worker.delete_data' was never awaited
    for modname, module in list(sys.modules.items()):
```

I _think_ that the only place where `Worker.delete_data` would be called and should be awaited is in https://github.com/dask/distributed/blob/4f878b420b349ee725de5ef64fd5e664dedb8aba/distributed/scheduler.py#L2791-L2800, but I don't have anything better than my guess at this time because it's really hard for me to understand all the async black magic. I'm gonna continue trying to figure this out, but any suggestions on how to pinpoint that are appreciated!


Discussion:
It's interesting as [`delete_data`]( https://github.com/dask/distributed/blob/5172678da28130ef7f4d86394d30566abc99328e/distributed/worker.py#L1344 ) is an `async` method, but we don't really call it from the worker (only the scheduler). So it doesn't get `await`ed anywhere AFAIK. Compare this to [`update_data`]( https://github.com/dask/distributed/blob/5172678da28130ef7f4d86394d30566abc99328e/distributed/worker.py#L1322 ), which is a similar method, but is not actually made `async`. I wonder if `delete_data` should have the `async` part removed.
cc @mrocklin @jrbourbeau (in case you have thoughts here 🙂)
Thanks @jakirkham for looking at that, I actually verified that applying your suggestion things work again:

```diff
diff --git a/distributed/worker.py b/distributed/worker.py
index 59cd285d..8bd45394 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1341,7 +1341,7 @@ class Worker(ServerNode):
         info = {"nbytes": {k: sizeof(v) for k, v in data.items()}, "status": "OK"}
         return info

-    async def delete_data(self, comm=None, keys=None, report=True):
+    def delete_data(self, comm=None, keys=None, report=True):
         if keys:
             for key in list(keys):
                 self.log.append((key, "delete"))
@@ -1355,7 +1355,7 @@ class Worker(ServerNode):
             if report:
                 logger.debug("Reporting loss of keys to scheduler")
                 # TODO: this route seems to not exist?
-                await self.scheduler.remove_keys(
+                self.scheduler.remove_keys(
                     address=self.contact_address, keys=list(keys)
                 )
         return "OK"
```

Possibly the second part can be removed/has to be fixed, as the comment above it suggests. There's no `remove_keys` anywhere in this repository.

Happy to file a PR if this change is reasonable.
Thanks for checking Peter! Let's see what others say 🙂
cc @quasiben (for vis)

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1341,7 +1341,7 @@ def update_data(self, comm=None, data=None, report=True, serializers=None):
         info = {"nbytes": {k: sizeof(v) for k, v in data.items()}, "status": "OK"}
         return info
 
-    async def delete_data(self, comm=None, keys=None, report=True):
+    def delete_data(self, comm=None, keys=None, report=True):
         if keys:
             for key in list(keys):
                 self.log.append((key, "delete"))
@@ -1352,12 +1352,6 @@ async def delete_data(self, comm=None, keys=None, report=True):
                     self.release_dep(key)
 
             logger.debug("Deleted %d keys", len(keys))
-            if report:
-                logger.debug("Reporting loss of keys to scheduler")
-                # TODO: this route seems to not exist?
-                await self.scheduler.remove_keys(
-                    address=self.contact_address, keys=list(keys)
-                )
         return "OK"
 
     async def set_resources(self, **resources):


Test patch: