Close clusters at exit

distributed/deploy/spec.py

@@ -1,4 +1,5 @@
 import asyncio
+import atexit
 import weakref
 
 from tornado import gen
@@ -97,6 +98,8 @@ class does handle all of the logic around asynchronously cleanly setting up
     specifications into the same dictionary.
     """
 
+    _instances = weakref.WeakSet()
+
     def __init__(
         self,
         workers=None,
@@ -133,6 +136,7 @@ def __init__(
             loop=self.loop, **self.scheduler_spec["options"]
         )
         self.status = "created"
+        self._instances.add(self)
         self._correct_state_waiting = None
 
         if not self.asynchronous:
@@ -248,9 +252,9 @@ async def _close(self):
 
         self.status = "closed"
 
-    def close(self):
+    def close(self, timeout=None):
         with ignoring(RuntimeError):  # loop closed during process shutdown
-            return self.sync(self._close)
+            return self.sync(self._close, callback_timeout=timeout)
 
     def __del__(self):
         if self.status != "closed":
@@ -295,3 +299,10 @@ def __repr__(self):
             self.scheduler_address,
             len(self.workers),
         )
+
+
+@atexit.register
+def close_clusters():
+    for cluster in list(SpecCluster._instances):
+        with ignoring(gen.TimeoutError):
+            cluster.close(timeout=10)

distributed/deploy/tests/test_spec_cluster.py

@@ -1,4 +1,5 @@
 from dask.distributed import SpecCluster, Worker, Client, Scheduler
+from distributed.deploy.spec import close_clusters
 from distributed.utils_test import loop  # noqa: F401
 import pytest
 
@@ -113,3 +114,13 @@ async def test_broken_worker():
             pass
 
     assert "Broken" in str(info.value)
+
+
+@pytest.mark.slow
+def test_spec_close_clusters(loop):
+    workers = {0: {"cls": Worker}}
+    scheduler = {"cls": Scheduler, "options": {"port": 0}}
+    cluster = SpecCluster(workers=workers, scheduler=scheduler, loop=loop)
+    assert cluster in SpecCluster._instances
+    close_clusters()
+    assert cluster.status == "closed"