[REVIEW] Updates to support nightly MG test automation

python/cugraph/tests/dask/mg_context.py

@@ -37,11 +37,12 @@ def get_visible_devices():
 
 
 def skip_if_not_enough_devices(required_devices):
-    visible_devices = get_visible_devices()
-    number_of_visible_devices = len(visible_devices)
-    if required_devices > number_of_visible_devices:
-        pytest.skip("Not enough devices available to "
-                    "test MG({})".format(required_devices))
+    if required_devices is not None:
+        visible_devices = get_visible_devices()
+        number_of_visible_devices = len(visible_devices)
+        if required_devices > number_of_visible_devices:
+            pytest.skip("Not enough devices available to "
+                        "test MG({})".format(required_devices))
 
 
 class MGContext:

python/cugraph/tests/dask/test_mg_batch_betweenness_centrality.py

@@ -37,7 +37,11 @@
 # Parameters
 # =============================================================================
 DATASETS = ["../datasets/karate.csv"]
-MG_DEVICE_COUNT_OPTIONS = [1, 2, 3, 4]
+MG_DEVICE_COUNT_OPTIONS = [pytest.param(1, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(2, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(3, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(4, marks=pytest.mark.preset_gpu_count),
+                           None]
 RESULT_DTYPE_OPTIONS = [np.float64]
 
 

python/cugraph/tests/dask/test_mg_batch_edge_betweenness_centrality.py

@@ -37,7 +37,11 @@
 # Parameters
 # =============================================================================
 DATASETS = ["../datasets/karate.csv"]
-MG_DEVICE_COUNT_OPTIONS = [1, 2, 4]
+MG_DEVICE_COUNT_OPTIONS = [pytest.param(1, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(2, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(3, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(4, marks=pytest.mark.preset_gpu_count),
+                           None]
 RESULT_DTYPE_OPTIONS = [np.float64]
 
 

python/cugraph/tests/dask/test_mg_bfs.py

@@ -27,6 +27,7 @@
 def client_connection():
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/cugraph/tests/dask/test_mg_comms.py

@@ -27,6 +27,7 @@
 def client_connection():
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/cugraph/tests/dask/test_mg_degree.py

@@ -28,6 +28,7 @@
 def client_connection():
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/cugraph/tests/dask/test_mg_katz_centrality.py

@@ -31,6 +31,7 @@
 def client_connection():
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/cugraph/tests/dask/test_mg_louvain.py

@@ -47,6 +47,7 @@ def client_connection():
     # setup
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/cugraph/tests/dask/test_mg_pagerank.py

@@ -56,6 +56,7 @@ def personalize(vertices, personalization_perc):
 def client_connection():
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/cugraph/tests/dask/test_mg_renumber.py

@@ -36,6 +36,7 @@
 def client_connection():
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/cugraph/tests/dask/test_mg_replication.py

@@ -24,8 +24,11 @@
 
 DATASETS_OPTIONS = utils.DATASETS_SMALL
 DIRECTED_GRAPH_OPTIONS = [False, True]
-# MG_DEVICE_COUNT_OPTIONS = [1, 2, 3, 4]
-MG_DEVICE_COUNT_OPTIONS = [1]
+MG_DEVICE_COUNT_OPTIONS = [pytest.param(1, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(2, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(3, marks=pytest.mark.preset_gpu_count),
+                           pytest.param(4, marks=pytest.mark.preset_gpu_count),
+                           None]
 
 
 @pytest.mark.skipif(

python/cugraph/tests/dask/test_mg_sssp.py

@@ -27,6 +27,7 @@
 def client_connection():
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/cugraph/tests/dask/test_mg_utility.py

@@ -39,6 +39,7 @@ def setup_function():
 def client_connection():
     cluster = LocalCUDACluster()
     client = Client(cluster)
+    client.wait_for_workers(None)  # number of devices None = all vsble devices
     Comms.initialize(p2p=True)
 
     yield client

python/pytest.ini

@@ -11,6 +11,7 @@ markers =
           managedmem_off: RMM managed memory disabled
           poolallocator_on: RMM pool allocator enabled
           poolallocator_off: RMM pool allocator disabled
+          preset_gpu_count: Use a hard-coded number of GPUs for specific MG tests
           ETL: benchmarks for ETL steps
           small: small datasets
           tiny: tiny datasets