Repository: dask/distributed

Issue:
RuntimeWarning: coroutine 'Client._close' was never awaited
I'm using a wrapper function to add semaphore functionality to a function I want to run in parallel and it's causing RuntimeWarnings to appear when I close the Client. Aside from the warnings, it seems to be working just how I intended.

Using dask version 2.15.0 and distributed version 2.15.2

Minimal reproducible example:
```
import time
import logging
import sys
import tracemalloc
from dask.distributed import Client, Semaphore


logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
tracemalloc.start()


def test(index):
    time.sleep(1)
    print(f"{index} done")


def wrapper(func, sem):
    def inner(*args, **kwargs):
        with sem:
            return func(*args, **kwargs)
    return inner


def main():
    client = Client(
        n_workers=4,
        threads_per_worker=1,
        dashboard_address=None,
    )
    sem = Semaphore(max_leases=2, client=client)
    futures = client.map(wrapper(test, sem), range(10))
    client.gather(futures)
    sem.close()
    client.close()


if __name__ == '__main__':
    main()
```

output:

```
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
DEBUG:asyncio:Using selector: KqueueSelector
2 done
3 done
1 done
0 done
6 done
7 done
5 done
4 done
9 done
8 done
/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'wait_for' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py", lineno 1357
    future = asyncio.wait_for(future, timeout)
/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'wait_for' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py", lineno 1357
    future = asyncio.wait_for(future, timeout)
/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'Client._close' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py", lineno 1355
    future = self._close()
/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'wait_for' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py", lineno 1357
    future = asyncio.wait_for(future, timeout)
/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'Client._close' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py", lineno 1355
    future = self._close()
/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'Client._close' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py", lineno 1355
    future = self._close()
/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'wait_for' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py", lineno 1357
    future = asyncio.wait_for(future, timeout)
/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'Client._close' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/anaconda3/lib/python3.7/site-packages/distributed/client.py", lineno 1355
    future = self._close()

Process finished with exit code 0

```


Discussion:
I guess this should be fixed by https://github.com/dask/distributed/pull/3713
You're right, that appears to be the fix.
I've updated to 2.18 and I'm still seeing this. I didn't pay close enough attention before, but the other PR just fixes the issue in semaphore, what I'm seeing here is in client.

<details>
<summary> Traceback </summary >

```
/Users/shannon/opt/miniconda3/envs/pywakes/lib/python3.8/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'wait_for' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/miniconda3/envs/pywakes/lib/python3.8/site-packages/distributed/client.py", lineno 1357
    future = asyncio.wait_for(future, timeout)
/Users/shannon/opt/miniconda3/envs/pywakes/lib/python3.8/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'wait_for' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/miniconda3/envs/pywakes/lib/python3.8/site-packages/distributed/client.py", lineno 1357
    future = asyncio.wait_for(future, timeout)
/Users/shannon/opt/miniconda3/envs/pywakes/lib/python3.8/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'Client._close' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/miniconda3/envs/pywakes/lib/python3.8/site-packages/distributed/client.py", lineno 1355
    future = self._close()
/Users/shannon/opt/miniconda3/envs/pywakes/lib/python3.8/site-packages/distributed/client.py:4721: RuntimeWarning: coroutine 'Client._close' was never awaited
  c.close(timeout=2)
Object allocated at (most recent call last):
  File "/Users/shannon/opt/miniconda3/envs/pywakes/lib/python3.8/site-packages/distributed/client.py", lineno 1355
    future = self._close()
```
</details>

@fjetter this issue needs to be re-opened, but I can't because you closed it.
Indeed, this is still not fixed. Issue comes from the client which is attached to the worker. Will open another PR for this

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -715,6 +715,7 @@ def __init__(
         )
 
         self._start_arg = address
+        self._set_as_default = set_as_default
         if set_as_default:
             self._set_config = dask.config.set(
                 scheduler="dask.distributed", shuffle="tasks"
@@ -958,8 +959,8 @@ def start(self, **kwargs):
             return
 
         self._loop_runner.start()
-
-        _set_global_client(self)
+        if self._set_as_default:
+            _set_global_client(self)
         self.status = "connecting"
 
         if self.asynchronous:
@@ -1150,8 +1151,8 @@ async def _ensure_connected(self, timeout=None):
         bcomm = BatchedSend(interval="10ms", loop=self.loop)
         bcomm.start(comm)
         self.scheduler_comm = bcomm
-
-        _set_global_client(self)
+        if self._set_as_default:
+            _set_global_client(self)
         self.status = "running"
 
         for msg in self._pending_msg_buffer:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -363,6 +363,7 @@ class Worker(ServerNode):
     """
 
     _instances = weakref.WeakSet()
+    _initialized_clients = weakref.WeakSet()
 
     def __init__(
         self,
@@ -1236,6 +1237,29 @@ async def close(
 
             for pc in self.periodic_callbacks.values():
                 pc.stop()
+
+            if self._client:
+                # If this worker is the last one alive, clean up the worker
+                # initialized clients
+                if not any(
+                    w
+                    for w in Worker._instances
+                    if w != self and w.status == Status.running
+                ):
+                    for c in Worker._initialized_clients:
+                        # Regardless of what the client was initialized with
+                        # we'll require the result as a future. This is
+                        # necessary since the heursitics of asynchronous are not
+                        # reliable and we might deadlock here
+                        c._asynchronous = True
+                        if c.asynchronous:
+                            await c.close()
+                        else:
+                            # There is still the chance that even with us
+                            # telling the client to be async, itself will decide
+                            # otherwise
+                            c.close()
+
             with suppress(EnvironmentError, TimeoutError):
                 if report and self.contact_address is not None:
                     await asyncio.wait_for(
@@ -3256,6 +3280,7 @@ def _get_client(self, timeout=None):
                 name="worker",
                 timeout=timeout,
             )
+            Worker._initialized_clients.add(self._client)
             if not asynchronous:
                 assert self._client.status == "running"
 


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -3344,6 +3344,32 @@ def test_default_get():
         assert dask.base.get_scheduler() == pre_get
 
 
+@gen_cluster()
+async def test_set_as_default(s, a, b):
+    with pytest.raises(ValueError):
+        default_client()
+
+    async with Client(s.address, set_as_default=False, asynchronous=True) as c1:
+        with pytest.raises(ValueError):
+            default_client()
+        async with Client(s.address, set_as_default=True, asynchronous=True) as c2:
+            assert default_client() is c2
+            async with Client(s.address, set_as_default=True, asynchronous=True) as c3:
+                assert default_client() is c3
+                async with Client(
+                    s.address, set_as_default=False, asynchronous=True
+                ) as c4:
+                    assert default_client() is c3
+
+                    await c4.scheduler_comm.close()
+                    while c4.status != "running":
+                        await asyncio.sleep(0.01)
+                    assert default_client() is c3
+
+    with pytest.raises(ValueError):
+        default_client()
+
+
 @gen_cluster(client=True)
 async def test_get_processing(c, s, a, b):
     processing = await c.processing()
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -1897,3 +1897,102 @@ async def test_gather_dep_one_worker_always_busy(c, s, a, b):
     # naturally
     while any(["Worker.gather_dep" in str(t) for t in asyncio.all_tasks()]):
         await asyncio.sleep(0.05)
+
+
+@gen_cluster(client=True, nthreads=[("127.0.0.1", 0)])
+async def test_worker_client_uses_default_no_close(c, s, a):
+    """
+    If a default client is available in the process, the worker will pick this
+    one and will not close it if it is closed
+    """
+    assert default_client() is c
+    existing_client = c.id
+
+    def get_worker_client_id():
+        def_client = get_client()
+        return def_client.id
+
+    worker_client = await c.submit(get_worker_client_id)
+    assert worker_client == existing_client
+
+    assert not Worker._initialized_clients
+
+    await a.close()
+
+    assert len(Client._instances) == 1
+    assert c.status == "running"
+    c_def = default_client()
+    assert c is c_def
+
+
+@gen_cluster(nthreads=[("127.0.0.1", 0)])
+async def test_worker_client_closes_if_created_on_worker_one_worker(s, a):
+    async with Client(s.address, set_as_default=False, asynchronous=True) as c:
+
+        with pytest.raises(ValueError):
+            default_client()
+
+        def get_worker_client_id():
+            def_client = get_client()
+            return def_client.id
+
+        new_client_id = await c.submit(get_worker_client_id)
+        default_client_id = await c.submit(get_worker_client_id)
+        assert new_client_id != c.id
+        assert new_client_id == default_client_id
+
+        new_client = default_client()
+        assert new_client_id == new_client.id
+        assert new_client.status == "running"
+
+        # If a worker closes, all clients created on it should close as well
+        await a.close()
+        assert new_client.status == "closed"
+
+        assert len(Client._instances) == 2
+
+        assert c.status == "running"
+
+        with pytest.raises(ValueError):
+            default_client()
+
+
+@gen_cluster(timeout=3600)
+async def test_worker_client_closes_if_created_on_worker_last_worker_alive(s, a, b):
+    async with Client(s.address, set_as_default=False, asynchronous=True) as c:
+
+        with pytest.raises(ValueError):
+            default_client()
+
+        def get_worker_client_id():
+            def_client = get_client()
+            return def_client.id
+
+        new_client_id = await c.submit(get_worker_client_id, workers=[a.address])
+        default_client_id = await c.submit(get_worker_client_id, workers=[a.address])
+
+        default_client_id_b = await c.submit(get_worker_client_id, workers=[b.address])
+        assert not b._comms
+        assert new_client_id != c.id
+        assert new_client_id == default_client_id
+        assert new_client_id == default_client_id_b
+
+        new_client = default_client()
+        assert new_client_id == new_client.id
+        assert new_client.status == "running"
+
+        # We'll close A. This should *not* close the client since the client is also used by B
+        await a.close()
+        assert new_client.status == "running"
+
+        client_id_b_after = await c.submit(get_worker_client_id, workers=[b.address])
+        assert client_id_b_after == default_client_id_b
+
+        assert len(Client._instances) == 2
+        await b.close()
+        assert new_client.status == "closed"
+
+        assert c.status == "running"
+
+        with pytest.raises(ValueError):
+            default_client()