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:

Resolving patch:
diff --git a/distributed/semaphore.py b/distributed/semaphore.py
--- a/distributed/semaphore.py
+++ b/distributed/semaphore.py
@@ -39,9 +39,9 @@ class SemaphoreExtension:
 
     * semaphore_acquire
     * semaphore_release
-    * semaphore_create
     * semaphore_close
     * semaphore_refresh_leases
+    * semaphore_register
     """
 
     def __init__(self, scheduler):
@@ -56,7 +56,7 @@ def __init__(self, scheduler):
 
         self.scheduler.handlers.update(
             {
-                "semaphore_create": self.create,
+                "semaphore_register": self.create,
                 "semaphore_acquire": self.acquire,
                 "semaphore_release": self.release,
                 "semaphore_close": self.close,
@@ -285,6 +285,12 @@ class Semaphore:
     client: Client (optional)
         Client to use for communication with the scheduler.  If not given, the
         default global client will be used.
+    register: bool
+        If True, register the semaphore with the scheduler. This needs to be
+        done before any leases can be acquired. If not done during
+        initialization, this can also be done by calling the register method of
+        this class.
+        When registering, this needs to be awaited.
 
     Examples
     --------
@@ -320,18 +326,13 @@ class Semaphore:
 
     """
 
-    def __init__(self, max_leases=1, name=None, client=None):
+    def __init__(self, max_leases=1, name=None, client=None, register=True):
         self.client = client or get_client()
         self.name = name or "semaphore-" + uuid.uuid4().hex
         self.max_leases = max_leases
         self.id = uuid.uuid4().hex
         self._leases = deque()
 
-        self._started = self.client.sync(
-            self.client.scheduler.semaphore_create,
-            name=self.name,
-            max_leases=max_leases,
-        )
         # this should give ample time to refresh without introducing another
         # config parameter since this *must* be smaller than the timeout anyhow
         refresh_leases_interval = (
@@ -352,9 +353,28 @@ def __init__(self, max_leases=1, name=None, client=None):
         pc.start()
         self.refresh_leases = True
 
+        self._registered = None
+        if register:
+            self._registered = self.register()
+
+    def register(self):
+        """
+        Register the semaphore on scheduler side
+
+        This will register the semaphore on scheduler side and ensure that all necessary data structures exist.
+        """
+        if self._registered is None:
+            self._registered = self.client.sync(
+                self.client.scheduler.semaphore_register,
+                name=self.name,
+                max_leases=self.max_leases,
+            )
+        return self._registered
+
     def __await__(self):
         async def create_semaphore():
-            await self._started
+            if self._registered:
+                await self._registered
             return self
 
         return create_semaphore().__await__()
@@ -444,7 +464,7 @@ def __getstate__(self):
     def __setstate__(self, state):
         name, max_leases = state
         client = get_client()
-        self.__init__(name=name, client=client, max_leases=max_leases)
+        self.__init__(name=name, client=client, max_leases=max_leases, register=False)
 
     def close(self):
         return self.client.sync(self.client.scheduler.semaphore_close, name=self.name)


Test patch: