Added cluster.shutdown to dask distributed tests

Author: seanlaw

tests/test_aamp_ostinato.py

@@ -2,6 +2,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy import core
@@ -17,7 +18,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 @pytest.mark.parametrize(

tests/test_aamp_stimp.py

@@ -2,6 +2,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy.aamp_stimp import aamp_stimp, aamp_stimped
@@ -23,7 +24,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 @pytest.mark.parametrize("T", T)

tests/test_aampdist.py

@@ -2,6 +2,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy.aampdist import _aampdist_vect, aampdist, aampdisted
@@ -16,7 +17,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 test_data = [

tests/test_aamped.py

@@ -3,6 +3,7 @@
 import numpy.testing as npt
 import pandas as pd
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy import config
@@ -18,7 +19,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 test_data = [

tests/test_maamped.py

@@ -3,6 +3,7 @@
 import numpy.testing as npt
 import pandas as pd
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy import config
@@ -18,7 +19,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 test_data = [

tests/test_mpdist.py

@@ -4,6 +4,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy.mpdist import _mpdist_vect, mpdist, mpdisted
@@ -18,7 +19,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 test_data = [

tests/test_mstumped.py

@@ -5,6 +5,7 @@
 import numpy.testing as npt
 import pandas as pd
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy import config
@@ -20,7 +21,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 test_data = [

tests/test_non_normalized_decorator.py

@@ -1,6 +1,7 @@
 import naive
 import numpy as np
 import numpy.testing as npt
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 from numba import cuda
 
@@ -37,6 +38,7 @@
     from stumpy.core import _gpu_ostinato_driver_not_found as gpu_ostinato  # noqa: F401
     from stumpy.core import _gpu_stimp_driver_not_found as gpu_stimp  # noqa: F401
     from stumpy.core import _gpu_stump_driver_not_found as gpu_stump  # noqa: F401
+
 from stumpy.maamp import maamp, maamp_mdl, maamp_multi_distance_profile, maamp_subspace
 from stumpy.maamped import maamped
 from stumpy.mmotifs import mmotifs
@@ -70,7 +72,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 test_data = [

tests/test_ostinato.py

@@ -4,6 +4,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy import core
@@ -19,7 +20,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 @pytest.mark.parametrize(

tests/test_stimp.py

@@ -4,6 +4,7 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
+import tornado.ioloop
 from dask.distributed import Client, LocalCluster
 
 from stumpy.stimp import stimp, stimped
@@ -23,7 +24,10 @@ def dask_cluster():
         worker_dashboard_address=None,
     )
     yield cluster.scheduler_address
-    cluster.close(timeout=60)
+    try:
+        cluster.close(timeout=60)
+    except tornado.ioloop.TimeoutError:  # pragma: no cover
+        cluster.shutdown()
 
 
 @pytest.mark.parametrize("T", T)