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:
fjetter: Thanks for the pointers! I wasn't aware that the client may be shared (my head often still thinks "worker ~ process"). Will have another look
matteobachetti: @fjetter, @mrocklin: I'm having this issue as well. Are there any news?
zhqu1148980644: Any updates？
fjetter: For the test I needed to fix the `set_as_default` flag of the client which was not reliable. I can break this change off into another PR if desired
fjetter: @jrbourbeau I'd appreciate a review on this

Patch:
diff --git a/distributed/client.py b/distributed/client.py
index f214a303..919d9181 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -715,6 +715,7 @@ class Client:
         )
 
         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 @@ class Client:
             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 @@ class Client:
         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:
@@ -1419,7 +1420,7 @@ class Client:
         if timeout == no_default:
             timeout = self._timeout * 2
         # XXX handling of self.status here is not thread-safe
-        if self.status == "closed":
+        if self.status in ["closed", "newly-created"]:
             if self.asynchronous:
                 future = asyncio.Future()
                 future.set_result(None)
diff --git a/distributed/comm/ucx.py b/distributed/comm/ucx.py
index b0bb62b5..4d10a898 100644
--- a/distributed/comm/ucx.py
+++ b/distributed/comm/ucx.py
@@ -11,15 +11,7 @@ import weakref
 
 import dask
 
-from ..utils import (
-    CancelledError,
-    ensure_ip,
-    get_ip,
-    get_ipv6,
-    log_errors,
-    nbytes,
-    parse_bytes,
-)
+from ..utils import ensure_ip, get_ip, get_ipv6, log_errors, nbytes, parse_bytes
 from .addressing import parse_host_port, unparse_host_port
 from .core import Comm, CommClosedError, Connector, Listener
 from .registry import Backend, backends
@@ -128,12 +120,15 @@ def init_once():
         ucx_create_endpoint = ucp.create_endpoint
         ucx_create_listener = ucp.create_listener
     else:
-        if dask.config.get("ucx.reuse-endpoints"):
-            ucx_create_endpoint = EndpointReuse.create_endpoint
-            ucx_create_listener = EndpointReuse.create_listener
-        else:
+        reuse_endpoints = dask.config.get("ucx.reuse-endpoints")
+        if (
+            reuse_endpoints is None and ucp.get_ucx_version() >= (1, 11, 0)
+        ) or reuse_endpoints is False:
             ucx_create_endpoint = ucp.create_endpoint
             ucx_create_listener = ucp.create_listener
+        else:
+            ucx_create_endpoint = EndpointReuse.create_endpoint
+            ucx_create_listener = EndpointReuse.create_listener
 
 
 class UCX(Comm):
@@ -266,7 +261,7 @@ class UCX(Comm):
                 (shutdown, nframes) = struct.unpack("?Q", msg)
 
                 if shutdown:  # The writer is closing the connection
-                    raise CancelledError("Connection closed by writer")
+                    raise CommClosedError("Connection closed by writer")
 
                 # Recv which frames are CUDA (bool) and
                 # how large each frame is (uint64)
@@ -275,9 +270,12 @@ class UCX(Comm):
                 await self.ep.recv(header)
                 header = struct.unpack(header_fmt, header)
                 cuda_frames, sizes = header[:nframes], header[nframes:]
-            except (ucp.exceptions.UCXBaseException, CancelledError):
+            except (
+                ucp.exceptions.UCXCloseError,
+                ucp.exceptions.UCXCanceled,
+            ) + (getattr(ucp.exceptions, "UCXConnectionReset", ()),):
                 self.abort()
-                raise CommClosedError("While reading, the connection was closed")
+                raise CommClosedError("Connection closed by writer")
             else:
                 # Recv frames
                 frames = [
@@ -311,7 +309,11 @@ class UCX(Comm):
         if self._ep is not None:
             try:
                 await self.ep.send(struct.pack("?Q", True, 0))
-            except ucp.exceptions.UCXError:
+            except (
+                ucp.exceptions.UCXError,
+                ucp.exceptions.UCXCloseError,
+                ucp.exceptions.UCXCanceled,
+            ) + (getattr(ucp.exceptions, "UCXConnectionReset", ()),):
                 # If the other end is in the process of closing,
                 # UCX will sometimes raise a `Input/output` error,
                 # which we can ignore.
@@ -344,7 +346,13 @@ class UCXConnector(Connector):
         logger.debug("UCXConnector.connect: %s", address)
         ip, port = parse_host_port(address)
         init_once()
-        ep = await ucx_create_endpoint(ip, port)
+        try:
+            ep = await ucx_create_endpoint(ip, port)
+        except (
+            ucp.exceptions.UCXCloseError,
+            ucp.exceptions.UCXCanceled,
+        ) + (getattr(ucp.exceptions, "UCXConnectionReset", ()),):
+            raise CommClosedError("Connection closed before handshake completed")
         return self.comm_class(
             ep,
             local_addr=None,
diff --git a/distributed/core.py b/distributed/core.py
index cfd84bce..d769611a 100644
--- a/distributed/core.py
+++ b/distributed/core.py
@@ -136,7 +136,6 @@ class Server:
         connection_args=None,
         timeout=None,
         io_loop=None,
-        **kwargs,
     ):
         self.handlers = {
             "identity": self.identity,
@@ -238,8 +237,6 @@ class Server:
 
         self.__stopped = False
 
-        super().__init__(**kwargs)
-
     @property
     def status(self):
         return self._status
@@ -645,7 +642,7 @@ async def send_recv(comm, reply=True, serializers=None, deserializers=None, **kw
             response = await comm.read(deserializers=deserializers)
         else:
             response = None
-    except EnvironmentError:
+    except (EnvironmentError, CommClosedError):
         # On communication errors, we should simply close the communication
         force_close = True
         raise
diff --git a/distributed/distributed-schema.yaml b/distributed/distributed-schema.yaml
index 0e8faf2f..eac459d7 100644
--- a/distributed/distributed-schema.yaml
+++ b/distributed/distributed-schema.yaml
@@ -855,6 +855,10 @@ properties:
           and compiled with hwloc support. Unexpected errors can occur when using
           ``"auto"`` if any interfaces are disconnected or improperly configured.
       reuse-endpoints:
-        type: boolean
-        description: | 
-          Whether to reuse endpoints or not.
+        type: [boolean, 'null']
+        description: |
+          Enable UCX-Py reuse endpoints mechanism if ``True`` or if it's not specified and
+          UCX < 1.11 is installed, otherwise disable reuse endpoints. This was primarily
+          introduced to resolve an issue with CUDA IPC that has been fixed in UCX 1.10, but
+          can cause establishing endpoints to be very slow, this is particularly noticeable in
+          clusters of more than a few dozen workers.
diff --git a/distributed/distributed.yaml b/distributed/distributed.yaml
index c71845e3..ff34e48b 100644
--- a/distributed/distributed.yaml
+++ b/distributed/distributed.yaml
@@ -196,4 +196,4 @@ ucx:
   infiniband: False # enable Infiniband
   rdmacm: False # enable RDMACM
   net-devices: null  # define what interface to use for UCX comm
-  reuse-endpoints: True  # enable endpoint reuse
+  reuse-endpoints: null  # enable endpoint reuse
diff --git a/distributed/nanny.py b/distributed/nanny.py
index 127906b1..3d9fb185 100644
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -9,6 +9,7 @@ import warnings
 import weakref
 from contextlib import suppress
 from multiprocessing.queues import Empty
+from time import sleep as sync_sleep
 
 import psutil
 from tornado import gen
@@ -362,7 +363,6 @@ class Nanny(ServerNode):
                 config=self.config,
             )
 
-        self.auto_restart = True
         if self.death_timeout:
             try:
                 result = await asyncio.wait_for(
@@ -378,7 +378,11 @@ class Nanny(ServerNode):
                 raise
 
         else:
-            result = await self.process.start()
+            try:
+                result = await self.process.start()
+            except Exception:
+                await self.close()
+                raise
         return result
 
     async def restart(self, comm=None, timeout=2, executor_wait=True):
@@ -414,9 +418,10 @@ class Nanny(ServerNode):
         """Track worker's memory.  Restart if it goes above terminate fraction"""
         if self.status != Status.running:
             return
+        if self.process is None or self.process.process is None:
+            return None
         process = self.process.process
-        if process is None:
-            return
+
         try:
             proc = self._psutil_process
             memory = proc.memory_info().rss
@@ -519,6 +524,9 @@ class Nanny(ServerNode):
 
 
 class WorkerProcess:
+    # The interval how often to check the msg queue for init
+    _init_msg_interval = 0.05
+
     def __init__(
         self,
         worker_kwargs,
@@ -584,9 +592,14 @@ class WorkerProcess:
         except OSError:
             logger.exception("Nanny failed to start process", exc_info=True)
             self.process.terminate()
-            return
-
-        msg = await self._wait_until_connected(uid)
+            self.status = Status.failed
+            return self.status
+        try:
+            msg = await self._wait_until_connected(uid)
+        except Exception:
+            self.status = Status.failed
+            self.process.terminate()
+            raise
         if not msg:
             return self.status
         self.worker_address = msg["address"]
@@ -683,14 +696,15 @@ class WorkerProcess:
                 logger.error("Failed to kill worker process: %s", e)
 
     async def _wait_until_connected(self, uid):
-        delay = 0.05
         while True:
             if self.status != Status.starting:
                 return
+            # This is a multiprocessing queue and we'd block the event loop if
+            # we simply called get
             try:
                 msg = self.init_result_q.get_nowait()
             except Empty:
-                await asyncio.sleep(delay)
+                await asyncio.sleep(self._init_msg_interval)
                 continue
 
             if msg["uid"] != uid:  # ensure that we didn't cross queues
@@ -700,7 +714,6 @@ class WorkerProcess:
                 logger.error(
                     "Failed while trying to start worker process: %s", msg["exception"]
                 )
-                await self.process.join()
                 raise msg["exception"]
             else:
                 return msg
@@ -718,88 +731,108 @@ class WorkerProcess:
         config,
         Worker,
     ):  # pragma: no cover
-        os.environ.update(env)
-        dask.config.set(config)
         try:
-            from dask.multiprocessing import initialize_worker_process
-        except ImportError:  # old Dask version
-            pass
-        else:
-            initialize_worker_process()
+            os.environ.update(env)
+            dask.config.set(config)
+            try:
+                from dask.multiprocessing import initialize_worker_process
+            except ImportError:  # old Dask version
+                pass
+            else:
+                initialize_worker_process()
 
-        if silence_logs:
-            logger.setLevel(silence_logs)
+            if silence_logs:
+                logger.setLevel(silence_logs)
 
-        IOLoop.clear_instance()
-        loop = IOLoop()
-        loop.make_current()
-        worker = Worker(**worker_kwargs)
+            IOLoop.clear_instance()
+            loop = IOLoop()
+            loop.make_current()
+            worker = Worker(**worker_kwargs)
 
-        async def do_stop(timeout=5, executor_wait=True):
-            try:
-                await worker.close(
-                    report=True,
-                    nanny=False,
-                    safe=True,  # TODO: Graceful or not?
-                    executor_wait=executor_wait,
-                    timeout=timeout,
-                )
-            finally:
-                loop.stop()
-
-        def watch_stop_q():
-            """
-            Wait for an incoming stop message and then stop the
-            worker cleanly.
-            """
-            while True:
-                try:
-                    msg = child_stop_q.get(timeout=1000)
-                except Empty:
-                    pass
-                else:
-                    child_stop_q.close()
-                    assert msg.pop("op") == "stop"
-                    loop.add_callback(do_stop, **msg)
-                    break
-
-        t = threading.Thread(target=watch_stop_q, name="Nanny stop queue watch")
-        t.daemon = True
-        t.start()
-
-        async def run():
-            """
-            Try to start worker and inform parent of outcome.
-            """
-            try:
-                await worker
-            except Exception as e:
-                logger.exception("Failed to start worker")
-                init_result_q.put({"uid": uid, "exception": e})
-                init_result_q.close()
-            else:
+            async def do_stop(timeout=5, executor_wait=True):
                 try:
-                    assert worker.address
-                except ValueError:
-                    pass
-                else:
-                    init_result_q.put(
-                        {
-                            "address": worker.address,
-                            "dir": worker.local_directory,
-                            "uid": uid,
-                        }
+                    await worker.close(
+                        report=True,
+                        nanny=False,
+                        safe=True,  # TODO: Graceful or not?
+                        executor_wait=executor_wait,
+                        timeout=timeout,
                     )
+                finally:
+                    loop.stop()
+
+            def watch_stop_q():
+                """
+                Wait for an incoming stop message and then stop the
+                worker cleanly.
+                """
+                while True:
+                    try:
+                        msg = child_stop_q.get(timeout=1000)
+                    except Empty:
+                        pass
+                    else:
+                        child_stop_q.close()
+                        assert msg.pop("op") == "stop"
+                        loop.add_callback(do_stop, **msg)
+                        break
+
+            t = threading.Thread(target=watch_stop_q, name="Nanny stop queue watch")
+            t.daemon = True
+            t.start()
+
+            async def run():
+                """
+                Try to start worker and inform parent of outcome.
+                """
+                try:
+                    await worker
+                except Exception as e:
+                    logger.exception("Failed to start worker")
+                    init_result_q.put({"uid": uid, "exception": e})
                     init_result_q.close()
-                    await worker.finished()
-                    logger.info("Worker closed")
-
-        try:
-            loop.run_sync(run)
-        except (TimeoutError, gen.TimeoutError):
-            # Loop was stopped before wait_until_closed() returned, ignore
-            pass
-        except KeyboardInterrupt:
-            # At this point the loop is not running thus we have to run
-            # do_stop() explicitly.
-            loop.run_sync(do_stop)
+                    # If we hit an exception here we need to wait for a least
+                    # one interval for the outside to pick up this message.
+                    # Otherwise we arrive in a race condition where the process
+                    # cleanup wipes the queue before the exception can be
+                    # properly handled. See also
+                    # WorkerProcess._wait_until_connected (the 2 is for good
+                    # measure)
+                    sync_sleep(cls._init_msg_interval * 2)
+                else:
+                    try:
+                        assert worker.address
+                    except ValueError:
+                        pass
+                    else:
+                        init_result_q.put(
+                            {
+                                "address": worker.address,
+                                "dir": worker.local_directory,
+                                "uid": uid,
+                            }
+                        )
+                        init_result_q.close()
+                        await worker.finished()
+                        logger.info("Worker closed")
+
+        except Exception as e:
+            logger.exception("Failed to initialize Worker")
+            init_result_q.put({"uid": uid, "exception": e})
+            init_result_q.close()
+            # If we hit an exception here we need to wait for a least one
+            # interval for the outside to pick up this message. Otherwise we
+            # arrive in a race condition where the process cleanup wipes the
+            # queue before the exception can be properly handled. See also
+            # WorkerProcess._wait_until_connected (the 2 is for good measure)
+            sync_sleep(cls._init_msg_interval * 2)
+        else:
+            try:
+                loop.run_sync(run)
+            except (TimeoutError, gen.TimeoutError):
+                # Loop was stopped before wait_until_closed() returned, ignore
+                pass
+            except KeyboardInterrupt:
+                # At this point the loop is not running thus we have to run
+                # do_stop() explicitly.
+                loop.run_sync(do_stop)
diff --git a/distributed/worker.py b/distributed/worker.py
index 494df279..e7746b43 100644
--- 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 @@ class Worker(ServerNode):
 
             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 @@ class Worker(ServerNode):
                 name="worker",
                 timeout=timeout,
             )
+            Worker._initialized_clients.add(self._client)
             if not asynchronous:
                 assert self._client.status == "running"
 


Test patch:
diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py
index a8b08513..d12c8bfc 100644
--- a/distributed/deploy/tests/test_local.py
+++ b/distributed/deploy/tests/test_local.py
@@ -1063,3 +1063,21 @@ async def test_cluster_names():
 
         async with LocalCluster(processes=False, asynchronous=True) as unnamed_cluster2:
             assert unnamed_cluster2 != unnamed_cluster
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("nanny", [True, False])
+async def test_local_cluster_redundant_kwarg(nanny):
+    with pytest.raises(TypeError, match="unexpected keyword argument"):
+        # Extra arguments are forwarded to the worker class. Depending on
+        # whether we use the nanny or not, the error treatment is quite
+        # different and we should assert that an exception is raised
+        async with await LocalCluster(
+            typo_kwarg="foo", processes=nanny, n_workers=1
+        ) as cluster:
+
+            # This will never work but is a reliable way to block without hard
+            # coding any sleep values
+            async with Client(cluster) as c:
+                f = c.submit(sleep, 0)
+                await f
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index ec595dc8..6fbbae95 100644
--- 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_core.py b/distributed/tests/test_core.py
index 6dd76f90..de0218e7 100644
--- a/distributed/tests/test_core.py
+++ b/distributed/tests/test_core.py
@@ -880,3 +880,9 @@ async def test_close_properly():
 
         # weakref set/dict should be cleaned up
         assert not len(server._ongoing_coroutines)
+
+
+@pytest.mark.asyncio
+async def test_server_redundant_kwarg():
+    with pytest.raises(TypeError, match="unexpected keyword argument"):
+        await Server({}, typo_kwarg="foo")
diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py
index 779cbdc4..1203fb3b 100644
--- a/distributed/tests/test_nanny.py
+++ b/distributed/tests/test_nanny.py
@@ -14,7 +14,6 @@ from tornado.ioloop import IOLoop
 import dask
 
 from distributed import Client, Nanny, Scheduler, Worker, rpc, wait, worker
-from distributed.compatibility import MACOS
 from distributed.core import CommClosedError, Status
 from distributed.diagnostics import SchedulerPlugin
 from distributed.metrics import time
@@ -565,10 +564,19 @@ class BrokenWorker(worker.Worker):
         raise StartException("broken")
 
 
-@pytest.mark.flaky(reruns=10, reruns_delay=5, condition=MACOS)
 @pytest.mark.asyncio
 async def test_worker_start_exception(cleanup):
     # make sure this raises the right Exception:
     with pytest.raises(StartException):
         async with Nanny("tcp://localhost:1", worker_class=BrokenWorker) as n:
             await n.start()
+
+
+@pytest.mark.asyncio
+async def test_failure_during_worker_initialization(cleanup):
+    with captured_logger(logger="distributed.nanny", level=logging.WARNING) as logs:
+        async with Scheduler() as s:
+            with pytest.raises(Exception):
+                async with Nanny(s.address, foo="bar") as n:
+                    await n
+        assert "Restarting worker" not in logs.getvalue()
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
index b1b5eb8b..7a9feffc 100644
--- 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()