Minor refactor to cleanup teardown steps for the local dask cluster b…

…y adding a teardown utility.

python/cugraph/dask/common/mg_utils.py

@@ -62,9 +62,23 @@ def get_visible_devices():
 
 
 def setup_local_dask_cluster(p2p=True):
+    """
+    Performs steps to setup a Dask cluster using LocalCUDACluster and returns
+    the LocalCUDACluster and corresponding client instance.
+    """
     cluster = LocalCUDACluster()
     client = Client(cluster)
     client.wait_for_workers(len(get_visible_devices()))
     Comms.initialize(p2p)
 
-    return (Comms, client, cluster)
+    return (cluster, client)
+
+
+def teardown_local_dask_cluster(cluster, client):
+    """
+    Performs steps to destroy a Dask cluster and a corresponding client
+    instance.
+    """
+    Comms.destroy()
+    client.close()
+    cluster.close()

python/cugraph/tests/dask/test_mg_bfs.py

@@ -18,20 +18,15 @@
 import dask_cudf
 import cudf
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(

python/cugraph/tests/dask/test_mg_comms.py

@@ -18,20 +18,15 @@
 import dask_cudf
 import cudf
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(

python/cugraph/tests/dask/test_mg_degree.py

@@ -17,20 +17,15 @@
 import cugraph
 import dask_cudf
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(

python/cugraph/tests/dask/test_mg_katz_centrality.py

@@ -19,20 +19,15 @@
 import dask_cudf
 import cudf
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(

python/cugraph/tests/dask/test_mg_louvain.py

@@ -18,7 +18,8 @@
 import dask_cudf
 from cugraph.tests import utils
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 try:
     from rapids_pytest_benchmark import setFixtureParamNames
@@ -42,15 +43,9 @@ def setFixtureParamNames(*args, **kwargs):
 # Fixtures
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(

python/cugraph/tests/dask/test_mg_pagerank.py

@@ -18,7 +18,8 @@
 import dask_cudf
 import cudf
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 
 # The function selects personalization_perc% of accessible vertices in graph M
@@ -52,15 +53,9 @@ def personalize(vertices, personalization_perc):
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(

python/cugraph/tests/dask/test_mg_renumber.py

@@ -27,20 +27,15 @@
 from cugraph.tests import utils
 from cugraph.structure.number_map import NumberMap
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 # Test all combinations of default/managed and pooled/non-pooled allocation

python/cugraph/tests/dask/test_mg_sssp.py

@@ -18,20 +18,15 @@
 import dask_cudf
 import cudf
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(

python/cugraph/tests/dask/test_mg_utility.py

@@ -21,7 +21,8 @@
 from cugraph.dask.common.part_utils import concat_within_workers
 from cugraph.dask.common.read_utils import get_n_workers
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 import os
 import time
 import numpy as np
@@ -37,15 +38,9 @@ def setup_function():
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(

python/cugraph/tests/test_symmetrize.py

@@ -20,7 +20,8 @@
 import cugraph
 from cugraph.tests import utils
 from cugraph.dask.common.mg_utils import (is_single_gpu,
-                                          setup_local_dask_cluster)
+                                          setup_local_dask_cluster,
+                                          teardown_local_dask_cluster)
 
 
 def test_version():
@@ -188,15 +189,9 @@ def test_symmetrize_weighted(graph_file):
 
 @pytest.fixture(scope="module")
 def client_connection():
-    # setup
-    (comms, client, cluster) = setup_local_dask_cluster(p2p=True)
-
+    (cluster, client) = setup_local_dask_cluster(p2p=True)
     yield client
-
-    # teardown
-    comms.destroy()
-    client.close()
-    cluster.close()
+    teardown_local_dask_cluster(cluster, client)
 
 
 @pytest.mark.skipif(