Fix MG PLC algos intermittent hang

conda/recipes/cugraph/meta.yaml

@@ -42,6 +42,8 @@ requirements:
     - ucx-proc=*=gpu
     - cudatoolkit {{ cuda_version }}.*
     - libraft-headers {{ minor_version }}
+    # FIXME: this pin can be removed once we move to the GitHub Actions build process
+    - setuptools<=65.2.0
   run:
     - python x.x
     - pylibcugraph={{ version }}

python/cugraph/cugraph/dask/cores/core_number.py

@@ -114,6 +114,7 @@ def core_number(input_graph,
             degree_type,
             do_expensive_check,
             workers=[w],
+            allow_other_workers=False,
         )
         for w in Comms.get_workers()
     ]
@@ -126,7 +127,17 @@ def core_number(input_graph,
 
     wait(cudf_result)
 
-    ddf = dask_cudf.from_delayed(cudf_result)
+    ddf = dask_cudf.from_delayed(cudf_result).persist()
+    wait(ddf)
+
+    # FIXME: Dask doesn't always release it fast enough.
+    # For instance if the algo is run several times with
+    # the same PLC graph, the current iteration might try to cache
+    # the past iteration's futures and this can cause a hang if some
+    # of those futures get released midway
+    del result
+    del cudf_result
+
     if input_graph.renumbered:
         ddf = input_graph.unrenumber(ddf, "vertex")
 

python/cugraph/cugraph/dask/link_analysis/pagerank.py

@@ -286,6 +286,7 @@ def pagerank(input_graph,
                 max_iter,
                 do_expensive_check,
                 workers=[w],
+                allow_other_workers=False,
             )
             for w, data_personalization in data_prsztn.worker_to_parts.items()
         ]
@@ -304,6 +305,7 @@ def pagerank(input_graph,
                 max_iter,
                 do_expensive_check,
                 workers=[w],
+                allow_other_workers=False,
             )
             for w in Comms.get_workers()
         ]
@@ -316,7 +318,17 @@ def pagerank(input_graph,
 
     wait(cudf_result)
 
-    ddf = dask_cudf.from_delayed(cudf_result)
+    ddf = dask_cudf.from_delayed(cudf_result).persist()
+    wait(ddf)
+
+    # FIXME: Dask doesn't always release it fast enough.
+    # For instance if the algo is run several times with
+    # the same PLC graph, the current iteration might try to cache
+    # the past iteration's futures and this can cause a hang if some
+    # of those futures get released midway
+    del result
+    del cudf_result
+
     if input_graph.renumbered:
         ddf = input_graph.unrenumber(ddf, "vertex")
 

python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py

@@ -152,6 +152,7 @@ def uniform_neighbor_sample(input_graph,
             fanout_vals,
             with_replacement,
             workers=[w],
+            allow_other_workers=False,
         )
         for w in Comms.get_workers()
     ]
@@ -164,7 +165,16 @@ def uniform_neighbor_sample(input_graph,
 
     wait(cudf_result)
 
-    ddf = dask_cudf.from_delayed(cudf_result)
+    ddf = dask_cudf.from_delayed(cudf_result).persist()
+    wait(ddf)
+
+    # FIXME: Dask doesn't always release it fast enough.
+    # For instance if the algo is run several times with
+    # the same PLC graph, the current iteration might try to cache
+    # the past iteration's futures and this can cause a hang if some
+    # of those futures get released midway
+    del result
+    del cudf_result
 
     if input_graph.renumbered:
         ddf = input_graph.unrenumber(ddf, "sources", preserve_order=True)

python/cugraph/cugraph/dask/traversal/bfs.py

@@ -176,7 +176,8 @@ def bfs(input_graph,
             st[0],
             depth_limit,
             return_distances,
-            workers=[w]
+            workers=[w],
+            allow_other_workers=False,
         )
         for w, st in data_start.worker_to_parts.items()
     ]
@@ -188,7 +189,16 @@ def bfs(input_graph,
                    for cp_arrays in cupy_result]
     wait(cudf_result)
 
-    ddf = dask_cudf.from_delayed(cudf_result)
+    ddf = dask_cudf.from_delayed(cudf_result).persist()
+    wait(ddf)
+
+    # FIXME: Dask doesn't always release it fast enough.
+    # For instance if the algo is run several times with
+    # the same PLC graph, the current iteration might try to cache
+    # the past iteration's futures and this can cause a hang if some
+    # of those futures get released midway
+    del cupy_result
+    del cudf_result
 
     if input_graph.renumbered:
         ddf = input_graph.unrenumber(ddf, 'vertex')