fix segfault by gracefully printing warning if dask has already quit

rapidsai · Oct 30, 2023 · 1b53a08 · 1b53a08
1 parent e0c3a93
commit 1b53a08
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py b/python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py
@@ -320,7 +320,13 @@ def __init__(
     def __del__(self):
         if self.__is_graph_owner:
             if isinstance(self.__graph._plc_graph, dict):
-                distributed.get_client().unpublish_dataset("cugraph_graph")
+                try:
+                    distributed.get_client().unpublish_dataset("cugraph_graph")
+                except:
+                    warnings.warn(
+                        "Could not unpublish graph dataset, most likely because"
+                        " dask has already shut down."
+                    )
             del self.__graph
 
     def __make_offsets(self, input_dict):