diff --git a/python/nx-cugraph/nx_cugraph/classes/digraph.py b/python/nx-cugraph/nx_cugraph/classes/digraph.py index 0ef90b4e88..a73ca8d0b2 100644 --- a/python/nx-cugraph/nx_cugraph/classes/digraph.py +++ b/python/nx-cugraph/nx_cugraph/classes/digraph.py @@ -34,6 +34,7 @@ __all__ = ["CudaDiGraph", "DiGraph"] networkx_api = nxcg.utils.decorators.networkx_class(nx.DiGraph) +gpu_cpu_api = nxcg.utils.decorators._gpu_cpu_api(nx.DiGraph, __name__) class DiGraph(nx.DiGraph, Graph): @@ -110,21 +111,37 @@ def to_networkx_class(cls) -> type[nx.DiGraph]: ########################## # Dispatch to nx.DiGraph or CudaDiGraph - __contains__ = Graph.__dict__["__contains__"] - __len__ = Graph.__dict__["__len__"] - __iter__ = Graph.__dict__["__iter__"] - get_edge_data = Graph.__dict__["get_edge_data"] - has_edge = Graph.__dict__["has_edge"] - neighbors = Graph.__dict__["neighbors"] - has_node = Graph.__dict__["has_node"] - nbunch_iter = Graph.__dict__["nbunch_iter"] - number_of_nodes = Graph.__dict__["number_of_nodes"] - order = Graph.__dict__["order"] - successors = Graph.__dict__["neighbors"] # Alias - - clear = Graph.clear - clear_edges = Graph.clear_edges + __contains__ = gpu_cpu_api("__contains__") + __len__ = gpu_cpu_api("__len__") + __iter__ = gpu_cpu_api("__iter__") + + @networkx_api + def clear(self) -> None: + cudagraph = self._cudagraph if self._is_on_gpu else None + if self._is_on_cpu: + super().clear() + if cudagraph is not None: + cudagraph.clear() + self._set_cudagraph(cudagraph, clear_cpu=False) + + @networkx_api + def clear_edges(self) -> None: + cudagraph = self._cudagraph if self._is_on_gpu else None + if self._is_on_cpu: + super().clear_edges() + if cudagraph is not None: + cudagraph.clear_edges() + self._set_cudagraph(cudagraph, clear_cpu=False) + + get_edge_data = gpu_cpu_api("get_edge_data", edge_data=True) + has_edge = gpu_cpu_api("has_edge") + neighbors = gpu_cpu_api("neighbors") + has_node = gpu_cpu_api("has_node") + nbunch_iter = gpu_cpu_api("nbunch_iter") number_of_edges = Graph.number_of_edges + number_of_nodes = gpu_cpu_api("number_of_nodes") + order = gpu_cpu_api("order") + successors = gpu_cpu_api("successors") class CudaDiGraph(CudaGraph): diff --git a/python/nx-cugraph/nx_cugraph/classes/graph.py b/python/nx-cugraph/nx_cugraph/classes/graph.py index 9c92890289..759a6a10cd 100644 --- a/python/nx-cugraph/nx_cugraph/classes/graph.py +++ b/python/nx-cugraph/nx_cugraph/classes/graph.py @@ -48,6 +48,7 @@ __all__ = ["CudaGraph", "Graph"] networkx_api = nxcg.utils.decorators.networkx_class(nx.Graph) +gpu_cpu_api = nxcg.utils.decorators._gpu_cpu_api(nx.Graph, __name__) # The "everything" cache key is an internal implementation detail of NetworkX # that may change between releases. @@ -86,56 +87,6 @@ def clear(self) -> None: super().clear() -class _graph_property: - """Dispatch property to NetworkX or CudaGraph based on cache. - - For example, this will use any cached CudaGraph for ``len(G)``, which - prevents creating NetworkX data structures. - """ - - def __init__(self, attr, *, edge_data=False, node_data=False): - self._attr = attr - self._edge_data = edge_data - self._node_data = node_data - - def __get__(self, instance, owner=None): - nx_class = owner.to_networkx_class() - if instance is None: - # Let's handle e.g. `nxcg.Graph.__len__` to look and behave correctly. - # - # If you want the instance of `_graph_property`, get it from the class dict: - # >>> nxcg.Graph.__dict__["__len__"] - # - # Alternatives: - # - `return op.methodcaller(self._attr)` - # - This dispatches, but does not have e.g. __name__ - # - `return getattr(nx_class, self._attr)` - # - This does not dispatch--it always uses networkx--but does have attrs - prop = owner.__dict__[self._attr] - - def inner(self, *args, **kwargs): - return prop.__get__(self, owner)(*args, **kwargs) - - # Standard function-wrapping - nx_func = getattr(nx_class, self._attr) - inner.__name__ = nx_func.__name__ - inner.__doc__ = nx_func.__doc__ - inner.__qualname__ = nx_func.__qualname__ - inner.__defaults__ = nx_func.__defaults__ - inner.__kwdefaults__ = nx_func.__kwdefaults__ - inner.__dict__.update(nx_func.__dict__) - inner.__module__ = owner.__module__ - inner.__wrapped__ = nx_func - return inner - - cuda_graph = instance._get_cudagraph( - edge_data=self._edge_data, node_data=self._node_data - ) - if cuda_graph is not None: - return getattr(cuda_graph, self._attr) - return getattr(nx_class, self._attr).__get__(instance, owner) - - class Graph(nx.Graph): # Tell networkx to dispatch calls with this object to nx-cugraph __networkx_backend__: ClassVar[str] = "cugraph" # nx >=3.2 @@ -592,9 +543,9 @@ def from_dcsc( ########################## # Dispatch to nx.Graph or CudaGraph - __contains__ = _graph_property("__contains__") - __len__ = _graph_property("__len__") - __iter__ = _graph_property("__iter__") + __contains__ = gpu_cpu_api("__contains__") + __len__ = gpu_cpu_api("__len__") + __iter__ = gpu_cpu_api("__iter__") @networkx_api def clear(self) -> None: @@ -614,11 +565,11 @@ def clear_edges(self) -> None: cudagraph.clear_edges() self._set_cudagraph(cudagraph, clear_cpu=False) - get_edge_data = _graph_property("get_edge_data", edge_data=True) - has_edge = _graph_property("has_edge") - neighbors = _graph_property("neighbors") - has_node = _graph_property("has_node") - nbunch_iter = _graph_property("nbunch_iter") + get_edge_data = gpu_cpu_api("get_edge_data", edge_data=True) + has_edge = gpu_cpu_api("has_edge") + neighbors = gpu_cpu_api("neighbors") + has_node = gpu_cpu_api("has_node") + nbunch_iter = gpu_cpu_api("nbunch_iter") @networkx_api def number_of_edges( @@ -628,10 +579,11 @@ def number_of_edges( # NotImplemented by CudaGraph nx_class = self.to_networkx_class() return nx_class.number_of_edges(self, u, v) - return _graph_property("number_of_edges").__get__(self, self.__class__)() + return self._number_of_edges(u, v) - number_of_nodes = _graph_property("number_of_nodes") - order = _graph_property("order") + _number_of_edges = gpu_cpu_api("number_of_edges") + number_of_nodes = gpu_cpu_api("number_of_nodes") + order = gpu_cpu_api("order") # Future work: implement more graph methods, and handle e.g. `copy` diff --git a/python/nx-cugraph/nx_cugraph/classes/multidigraph.py b/python/nx-cugraph/nx_cugraph/classes/multidigraph.py index 643e2df5ed..0671d21a85 100644 --- a/python/nx-cugraph/nx_cugraph/classes/multidigraph.py +++ b/python/nx-cugraph/nx_cugraph/classes/multidigraph.py @@ -23,6 +23,7 @@ __all__ = ["CudaMultiDiGraph", "MultiDiGraph"] networkx_api = nxcg.utils.decorators.networkx_class(nx.MultiDiGraph) +gpu_cpu_api = nxcg.utils.decorators._gpu_cpu_api(nx.MultiDiGraph, __name__) class MultiDiGraph(nx.MultiDiGraph, MultiGraph, DiGraph): @@ -55,21 +56,20 @@ def to_networkx_class(cls) -> type[nx.MultiDiGraph]: ########################## # Dispatch to nx.MultiDiGraph or CudaMultiDiGraph - __contains__ = Graph.__dict__["__contains__"] - __len__ = Graph.__dict__["__len__"] - __iter__ = Graph.__dict__["__iter__"] - get_edge_data = Graph.__dict__["get_edge_data"] - has_edge = Graph.__dict__["has_edge"] - neighbors = Graph.__dict__["neighbors"] - has_node = Graph.__dict__["has_node"] - nbunch_iter = Graph.__dict__["nbunch_iter"] - number_of_nodes = Graph.__dict__["number_of_nodes"] - order = Graph.__dict__["order"] - successors = Graph.__dict__["neighbors"] # Alias - - clear = Graph.clear - clear_edges = Graph.clear_edges - number_of_edges = Graph.number_of_edges + __contains__ = gpu_cpu_api("__contains__") + __len__ = gpu_cpu_api("__len__") + __iter__ = gpu_cpu_api("__iter__") + clear = DiGraph.clear + clear_edges = DiGraph.clear_edges + get_edge_data = gpu_cpu_api("get_edge_data", edge_data=True) + has_edge = gpu_cpu_api("has_edge") + neighbors = gpu_cpu_api("neighbors") + has_node = gpu_cpu_api("has_node") + nbunch_iter = gpu_cpu_api("nbunch_iter") + number_of_edges = MultiGraph.number_of_edges + number_of_nodes = gpu_cpu_api("number_of_nodes") + order = gpu_cpu_api("order") + successors = gpu_cpu_api("successors") class CudaMultiDiGraph(CudaMultiGraph, CudaDiGraph): diff --git a/python/nx-cugraph/nx_cugraph/classes/multigraph.py b/python/nx-cugraph/nx_cugraph/classes/multigraph.py index e91103c1f6..3f0204f699 100644 --- a/python/nx-cugraph/nx_cugraph/classes/multigraph.py +++ b/python/nx-cugraph/nx_cugraph/classes/multigraph.py @@ -37,6 +37,7 @@ __all__ = ["MultiGraph", "CudaMultiGraph"] networkx_api = nxcg.utils.decorators.networkx_class(nx.MultiGraph) +gpu_cpu_api = nxcg.utils.decorators._gpu_cpu_api(nx.MultiGraph, __name__) class MultiGraph(nx.MultiGraph, Graph): @@ -282,20 +283,30 @@ def from_dcsc( ########################## # Dispatch to nx.MultiGraph or CudaMultiGraph - __contains__ = Graph.__dict__["__contains__"] - __len__ = Graph.__dict__["__len__"] - __iter__ = Graph.__dict__["__iter__"] - get_edge_data = Graph.__dict__["get_edge_data"] - has_edge = Graph.__dict__["has_edge"] - neighbors = Graph.__dict__["neighbors"] - has_node = Graph.__dict__["has_node"] - nbunch_iter = Graph.__dict__["nbunch_iter"] - number_of_nodes = Graph.__dict__["number_of_nodes"] - order = Graph.__dict__["order"] - + __contains__ = gpu_cpu_api("__contains__") + __len__ = gpu_cpu_api("__len__") + __iter__ = gpu_cpu_api("__iter__") clear = Graph.clear clear_edges = Graph.clear_edges - number_of_edges = Graph.number_of_edges + get_edge_data = gpu_cpu_api("get_edge_data", edge_data=True) + has_edge = gpu_cpu_api("has_edge") + neighbors = gpu_cpu_api("neighbors") + has_node = gpu_cpu_api("has_node") + nbunch_iter = gpu_cpu_api("nbunch_iter") + + @networkx_api + def number_of_edges( + self, u: NodeKey | None = None, v: NodeKey | None = None + ) -> int: + if u is not None or v is not None: + # NotImplemented by CudaGraph + nx_class = self.to_networkx_class() + return nx_class.number_of_edges(self, u, v) + return self._number_of_edges(u, v) + + _number_of_edges = gpu_cpu_api("number_of_edges") + number_of_nodes = gpu_cpu_api("number_of_nodes") + order = gpu_cpu_api("order") class CudaMultiGraph(CudaGraph): diff --git a/python/nx-cugraph/nx_cugraph/tests/test_graph_methods.py b/python/nx-cugraph/nx_cugraph/tests/test_graph_methods.py index b989f91b77..abc71ed775 100644 --- a/python/nx-cugraph/nx_cugraph/tests/test_graph_methods.py +++ b/python/nx-cugraph/nx_cugraph/tests/test_graph_methods.py @@ -90,21 +90,26 @@ def test_multidigraph_to_undirected(): ("nbunch_iter", ([0, 1],)), ], ) -def test_method_does_not_create_host_data(create_using, method): +@pytest.mark.parametrize("where", ["gpu", "cpu"]) +def test_method_does_not_convert_to_cpu_or_gpu(create_using, method, where): attr, args = method if attr == "successors" and not create_using.is_directed(): return G = nxcg.complete_graph(3, create_using=create_using) - assert G._is_on_gpu - assert not G._is_on_cpu + is_on_gpu = where == "gpu" + is_on_cpu = where == "cpu" + if is_on_cpu: + G.add_edge(10, 20) + assert G._is_on_gpu == is_on_gpu + assert G._is_on_cpu == is_on_cpu getattr(G, attr)(*args) - assert G._is_on_gpu - assert not G._is_on_cpu + assert G._is_on_gpu == is_on_gpu + assert G._is_on_cpu == is_on_cpu # Also usable from the class and dispatches correctly func = getattr(create_using, attr) func(G, *args) - assert G._is_on_gpu - assert not G._is_on_cpu + assert G._is_on_gpu == is_on_gpu + assert G._is_on_cpu == is_on_cpu # Basic "looks like networkx" checks nx_class = create_using.to_networkx_class() nx_func = getattr(nx_class, attr) diff --git a/python/nx-cugraph/nx_cugraph/utils/decorators.py b/python/nx-cugraph/nx_cugraph/utils/decorators.py index 15cb010e9b..26fce73503 100644 --- a/python/nx-cugraph/nx_cugraph/utils/decorators.py +++ b/python/nx-cugraph/nx_cugraph/utils/decorators.py @@ -167,3 +167,31 @@ def _default_should_run(*args, **kwargs): def _restore_networkx_dispatched(name): return getattr(BackendInterface, name) + + +def _gpu_cpu_api(nx_class, module_name): + def _gpu_cpu_graph_method(attr, *, edge_data=False, node_data=False): + """Dispatch property to NetworkX or CudaGraph based on cache. + + For example, this will use any cached CudaGraph for ``len(G)``, which + prevents creating NetworkX data structures. + """ + nx_func = getattr(nx_class, attr) + + def inner(self, *args, **kwargs): + cuda_graph = self._get_cudagraph(edge_data=edge_data, node_data=node_data) + if cuda_graph is None: + return nx_func(self, *args, **kwargs) + return getattr(cuda_graph, attr)(*args, **kwargs) + + inner.__name__ = nx_func.__name__ + inner.__doc__ = nx_func.__doc__ + inner.__qualname__ = nx_func.__qualname__ + inner.__defaults__ = nx_func.__defaults__ + inner.__kwdefaults__ = nx_func.__kwdefaults__ + inner.__module__ = module_name + inner.__dict__.update(nx_func.__dict__) + inner.__wrapped__ = nx_func + return inner + + return _gpu_cpu_graph_method