Refactor updating graphs: nx.utils._update_cpu_gpu_graphs

Author: eriknw

nx_cugraph/classes/graph.py

@@ -86,6 +86,9 @@ def clear(self) -> None:
         self._graph._reify_networkx()
         super().clear()
 
+    def _clear_no_reify_networkx(self):
+        super().clear()
+
 
 class Graph(nx.Graph):
     # Tell networkx to dispatch calls with this object to nx-cugraph

nx_cugraph/drawing/layout.py

@@ -21,12 +21,12 @@
 import pylibcugraph as plc
 from networkx.utils import create_random_state
 
-import nx_cugraph as nxcg
 from nx_cugraph.convert import _to_graph
 from nx_cugraph.utils import (
     _dtype_param,
     _get_float_dtype,
     _seed_to_int,
+    _update_cpu_gpu_graphs,
     networkx_algorithm,
 )
 
@@ -75,6 +75,7 @@ def forceatlas2_layout(
     if len(G) == 0:
         return {}
 
+    # Mutate original graph if store_pos_as is given.
     G_orig = G
 
     if dim != 2:
@@ -149,29 +150,24 @@ def forceatlas2_layout(
     pos = G._nodearrays_to_dict(
         node_ids=vertices, values=pos_arr, values_as_arrays=True
     )
-
+    store_pos_as = "pos"
     if store_pos_as is not None:
-        if isinstance(G_orig, nxcg.Graph):
-            # Could be on GPU, CPU, or both. Update both GPU and CPU (if present)
-            if G_orig._is_on_gpu:
-                cuda_graph = G_orig._cudagraph
+
+        def update_cpu(graph):
+            nx.set_node_attributes(graph, pos, store_pos_as)
+
+        update_pos_array = True
+
+        def update_gpu(cuda_graph):
+            # Ensure vertices are in order with their positions.
+            # Use nonlocal variable to do this only once to ensure idempotency.
+            nonlocal update_pos_array
+            if update_pos_array:
                 pos_arr[vertices] = pos_arr
-                cuda_graph.node_values[store_pos_as] = pos_arr
-            else:
-                cuda_graph = None
-            if G_orig._is_on_cpu:
-                # This clears the cache (including on GPU)
-                nx.set_node_attributes(G_orig, pos, store_pos_as)
-                if cuda_graph is not None:
-                    # Add back to GPU
-                    G_orig._set_cudagraph(cuda_graph, clear_cpu=False)
-        elif isinstance(G_orig, nxcg.CudaGraph):
-            # ensure vertices are in order with their positions
-            pos_arr[vertices] = pos_arr
-            G_orig.node_values[store_pos_as] = pos_arr
-        else:
-            # Default: networkx graph
-            nx.set_node_attributes(G_orig, pos, store_pos_as)
+                update_pos_array = False
+            cuda_graph.node_values[store_pos_as] = pos_arr
+
+        _update_cpu_gpu_graphs(G_orig, update_cpu=update_cpu, update_gpu=update_gpu)
 
     return pos
 

nx_cugraph/utils/misc.py

@@ -21,8 +21,12 @@
 import cupy as cp
 import numpy as np
 
+import nx_cugraph as nxcg
+
 if TYPE_CHECKING:
-    import nx_cugraph as nxcg
+    from collections.abc import Callable
+
+    import networkx as nx
 
     from ..typing import Dtype, EdgeKey
 
@@ -46,6 +50,7 @@ def pairwise(it):
     "_get_float_dtype",
     "_dtype_param",
     "_cp_iscopied_asarray",
+    "_update_cpu_gpu_graphs",
 ]
 
 # This may switch to np.uint32 at some point
@@ -240,6 +245,50 @@ def _cp_iscopied_asarray(a, *args, orig_object=None, **kwargs):
     return True, arr
 
 
+def _update_cpu_gpu_graphs(
+    G,
+    *,
+    update_cpu: Callable[[nx.Graph | nxcg.Graph], None],
+    update_gpu: Callable[[nxcg.CudaGraph], None],
+) -> None:
+    """Update graph in-place whether it's on CPU or GPU (or both).
+
+    This works with nx.Graph, nxcg.Graph, and nxcg.CudaGraph objects.
+    nxcg.Graph instances will update both CPU and GPU data structures
+    if applicable.
+
+    Parameters
+    ----------
+    update_cpu : func
+        Function to modify a networkx-compatible graph in-place.
+
+    update_gpu: func
+        Function to modify a CudaGraph graph in-place. This also
+        updates the CudaGraph of an nxcg.Graph if it is on the GPU.
+    """
+    if isinstance(G, nxcg.Graph):
+        # Could be on GPU, CPU, or both. Update both GPU and CPU (if present)
+        if G._is_on_gpu:
+            cuda_graph = G._cudagraph
+            update_gpu(cuda_graph)
+            # Clear anything else in the cache; cache is invalidated.
+            # We will re-add the cuda_graph below.
+            G.__networkx_cache__._clear_no_reify_networkx()
+        else:
+            cuda_graph = None
+        if G._is_on_cpu:
+            # This clears the cache (including on GPU)
+            update_cpu(G)
+        if cuda_graph is not None:
+            # Add back to GPU
+            G._set_cudagraph(cuda_graph, clear_cpu=False)
+    elif isinstance(G, nxcg.CudaGraph):
+        update_gpu(G)
+    else:
+        # Default: networkx graph
+        update_cpu(G)
+
+
 class _And_NotImplementedError(NotImplementedError):
     """Additionally make an exception a ``NotImplementedError``.