Merge branch 'branch-24.02' into cugraph-equivariant

rapidsai · Jan 23, 2024 · 9547bd5 · 9547bd5
2 parents 0a9c4b8 + f20219e
commit 9547bd5
Show file tree

Hide file tree

Showing 32 changed files with 1,026 additions and 165 deletions.
diff --git a/cpp/include/cugraph/mtmg/resource_manager.hpp b/cpp/include/cugraph/mtmg/resource_manager.hpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -106,9 +106,9 @@ class resource_manager_t {
     auto per_device_it = per_device_rmm_resources_.insert(
       std::pair{global_rank, std::make_shared<rmm::mr::cuda_memory_resource>()});
 #else
-    auto const [free, total] = rmm::detail::available_device_memory();
+    auto const [free, total] = rmm::available_device_memory();
     auto const min_alloc =
-      rmm::detail::align_down(std::min(free, total / 6), rmm::detail::CUDA_ALLOCATION_ALIGNMENT);
+      rmm::align_down(std::min(free, total / 6), rmm::CUDA_ALLOCATION_ALIGNMENT);
 
     auto per_device_it = per_device_rmm_resources_.insert(
       std::pair{global_rank,

diff --git a/cpp/tests/utilities/base_fixture.hpp b/cpp/tests/utilities/base_fixture.hpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -73,9 +73,8 @@ inline auto make_pool()
   // run more than 2 tests in parallel at the same time. Changes to this value could
   // effect the maximum amount of parallel tests, and therefore `tests/CMakeLists.txt`
   // `_CUGRAPH_TEST_PERCENT` default value will need to be audited.
-  auto const [free, total] = rmm::detail::available_device_memory();
-  auto const min_alloc =
-    rmm::detail::align_down(std::min(free, total / 6), rmm::detail::CUDA_ALLOCATION_ALIGNMENT);
+  auto const [free, total] = rmm::available_device_memory();
+  auto const min_alloc = rmm::align_down(std::min(free, total / 6), rmm::CUDA_ALLOCATION_ALIGNMENT);
   return rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(make_cuda(), min_alloc);
 }
 

diff --git a/python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py b/python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2023, NVIDIA CORPORATION.
+# Copyright (c) 2021-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -35,11 +35,9 @@
 from cugraph.structure.number_map import NumberMap
 from cugraph.structure.symmetrize import symmetrize
 from cugraph.dask.common.part_utils import (
-    get_persisted_df_worker_map,
     persist_dask_df_equal_parts_per_worker,
 )
 from cugraph.dask.common.mg_utils import run_gc_on_dask_cluster
-from cugraph.dask import get_n_workers
 import cugraph.dask.comms.comms as Comms
 
 
@@ -825,12 +823,13 @@ def get_two_hop_neighbors(self, start_vertices=None):
         _client = default_client()
 
         def _call_plc_two_hop_neighbors(sID, mg_graph_x, start_vertices):
-            return pylibcugraph_get_two_hop_neighbors(
+            results_ = pylibcugraph_get_two_hop_neighbors(
                 resource_handle=ResourceHandle(Comms.get_handle(sID).getHandle()),
                 graph=mg_graph_x,
                 start_vertices=start_vertices,
                 do_expensive_check=False,
             )
+            return results_
 
         if isinstance(start_vertices, int):
             start_vertices = [start_vertices]
@@ -845,31 +844,31 @@ def _call_plc_two_hop_neighbors(sID, mg_graph_x, start_vertices):
             else:
                 start_vertices_type = self.input_df.dtypes[0]
 
-            if not isinstance(start_vertices, (dask_cudf.Series)):
-                start_vertices = dask_cudf.from_cudf(
+            start_vertices = start_vertices.astype(start_vertices_type)
+
+            def create_iterable_args(
+                session_id, input_graph, start_vertices=None, npartitions=None
+            ):
+                session_id_it = [session_id] * npartitions
+                graph_it = input_graph.values()
+                start_vertices = cp.array_split(start_vertices.values, npartitions)
+                return [
+                    session_id_it,
+                    graph_it,
                     start_vertices,
-                    npartitions=min(self._npartitions, len(start_vertices)),
-                )
-                start_vertices = start_vertices.astype(start_vertices_type)
+                ]
 
-            n_workers = get_n_workers()
-            start_vertices = start_vertices.repartition(npartitions=n_workers)
-            start_vertices = persist_dask_df_equal_parts_per_worker(
-                start_vertices, _client
+            result = _client.map(
+                _call_plc_two_hop_neighbors,
+                *create_iterable_args(
+                    Comms.get_session_id(),
+                    self._plc_graph,
+                    start_vertices,
+                    self._npartitions,
+                ),
+                pure=False,
             )
-            start_vertices = get_persisted_df_worker_map(start_vertices, _client)
 
-            result = [
-                _client.submit(
-                    _call_plc_two_hop_neighbors,
-                    Comms.get_session_id(),
-                    self._plc_graph[w],
-                    start_vertices[w][0],
-                    workers=[w],
-                    allow_other_workers=False,
-                )
-                for w in start_vertices.keys()
-            ]
         else:
             result = [
                 _client.submit(
@@ -896,7 +895,8 @@ def convert_to_cudf(cp_arrays):
             return df
 
         cudf_result = [
-            _client.submit(convert_to_cudf, cp_arrays) for cp_arrays in result
+            _client.submit(convert_to_cudf, cp_arrays, pure=False)
+            for cp_arrays in result
         ]
 
         wait(cudf_result)

diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -30,6 +30,7 @@
     "functions": {
         # BEGIN: functions
         "ancestors",
+        "average_clustering",
         "barbell_graph",
         "betweenness_centrality",
         "bfs_edges",
@@ -41,10 +42,12 @@
         "caveman_graph",
         "chvatal_graph",
         "circular_ladder_graph",
+        "clustering",
         "complete_bipartite_graph",
         "complete_graph",
         "complete_multipartite_graph",
         "connected_components",
+        "core_number",
         "cubical_graph",
         "cycle_graph",
         "davis_southern_women_graph",
@@ -68,8 +71,15 @@
         "house_x_graph",
         "icosahedral_graph",
         "in_degree_centrality",
+        "is_arborescence",
+        "is_bipartite",
+        "is_branching",
         "is_connected",
+        "is_forest",
         "is_isolate",
+        "is_strongly_connected",
+        "is_tree",
+        "is_weakly_connected",
         "isolates",
         "k_truss",
         "karate_club_graph",
@@ -85,44 +95,56 @@
         "number_connected_components",
         "number_of_isolates",
         "number_of_selfloops",
+        "number_strongly_connected_components",
+        "number_weakly_connected_components",
         "octahedral_graph",
         "out_degree_centrality",
+        "overall_reciprocity",
         "pagerank",
         "pappus_graph",
         "path_graph",
         "petersen_graph",
+        "reciprocity",
         "sedgewick_maze_graph",
         "single_source_shortest_path_length",
         "single_target_shortest_path_length",
         "star_graph",
+        "strongly_connected_components",
         "tadpole_graph",
         "tetrahedral_graph",
+        "transitivity",
+        "triangles",
         "trivial_graph",
         "truncated_cube_graph",
         "truncated_tetrahedron_graph",
         "turan_graph",
         "tutte_graph",
+        "weakly_connected_components",
         "wheel_graph",
         # END: functions
     },
     "extra_docstrings": {
         # BEGIN: extra_docstrings
-        "betweenness_centrality": "`weight` parameter is not yet supported.",
+        "average_clustering": "Directed graphs and `weight` parameter are not yet supported.",
+        "betweenness_centrality": "`weight` parameter is not yet supported, and RNG with seed may be different.",
         "bfs_edges": "`sort_neighbors` parameter is not yet supported.",
         "bfs_predecessors": "`sort_neighbors` parameter is not yet supported.",
         "bfs_successors": "`sort_neighbors` parameter is not yet supported.",
         "bfs_tree": "`sort_neighbors` parameter is not yet supported.",
-        "edge_betweenness_centrality": "`weight` parameter is not yet supported.",
+        "clustering": "Directed graphs and `weight` parameter are not yet supported.",
+        "core_number": "Directed graphs are not yet supported.",
+        "edge_betweenness_centrality": "`weight` parameter is not yet supported, and RNG with seed may be different.",
         "eigenvector_centrality": "`nstart` parameter is not used, but it is checked for validity.",
-        "from_pandas_edgelist": "cudf.DataFrame inputs also supported.",
+        "from_pandas_edgelist": "cudf.DataFrame inputs also supported; value columns with str is unsuppported.",
         "generic_bfs_edges": "`neighbors` and `sort_neighbors` parameters are not yet supported.",
         "k_truss": (
             "Currently raises `NotImplementedError` for graphs with more than one connected\n"
             "component when k >= 3. We expect to fix this soon."
         ),
         "katz_centrality": "`nstart` isn't used (but is checked), and `normalized=False` is not supported.",
-        "louvain_communities": "`seed` parameter is currently ignored.",
+        "louvain_communities": "`seed` parameter is currently ignored, and self-loops are not yet supported.",
         "pagerank": "`dangling` parameter is not supported, but it is checked for validity.",
+        "transitivity": "Directed graphs are not yet supported.",
         # END: extra_docstrings
     },
     "extra_parameters": {

diff --git a/python/nx-cugraph/lint.yaml b/python/nx-cugraph/lint.yaml
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 #
 # https://pre-commit.com/
 #
@@ -36,7 +36,7 @@ repos:
       - id: autoflake
         args: [--in-place]
   - repo: https://github.com/pycqa/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - id: isort
   - repo: https://github.com/asottile/pyupgrade
@@ -45,23 +45,23 @@ repos:
       - id: pyupgrade
         args: [--py39-plus]
   - repo: https://github.com/psf/black
-    rev: 23.11.0
+    rev: 23.12.1
     hooks:
       - id: black
       # - id: black-jupyter
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.7
+    rev: v0.1.13
     hooks:
       - id: ruff
         args: [--fix-only, --show-fixes]  # --unsafe-fixes]
   - repo: https://github.com/PyCQA/flake8
-    rev: 6.1.0
+    rev: 7.0.0
     hooks:
       - id: flake8
         args: ['--per-file-ignores=_nx_cugraph/__init__.py:E501', '--extend-ignore=SIM105']  # Why is this necessary?
         additional_dependencies: &flake8_dependencies
           # These versions need updated manually
-          - flake8==6.1.0
+          - flake8==7.0.0
           - flake8-bugbear==23.12.2
           - flake8-simplify==0.21.0
   - repo: https://github.com/asottile/yesqa
@@ -77,7 +77,7 @@ repos:
         additional_dependencies: [tomli]
         files: ^(nx_cugraph|docs)/
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.7
+    rev: v0.1.13
     hooks:
       - id: ruff
   - repo: https://github.com/pre-commit/pre-commit-hooks

diff --git a/python/nx-cugraph/nx_cugraph/_version.py b/python/nx-cugraph/nx_cugraph/_version.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,12 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-
 import importlib.resources
 
 # Read VERSION file from the module that is symlinked to VERSION file
-# in the root of the repo at build time or copied to the moudle at
+# in the root of the repo at build time or copied to the module at
 # installation. VERSION is a separate file that allows CI build-time scripts
 # to update version info (including commit hashes) without modifying
 # source files.

diff --git a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -13,18 +13,23 @@
 from . import (
     bipartite,
     centrality,
+    cluster,
     community,
     components,
     link_analysis,
     shortest_paths,
     traversal,
+    tree,
 )
-from .bipartite import complete_bipartite_graph
+from .bipartite import complete_bipartite_graph, is_bipartite
 from .centrality import *
+from .cluster import *
 from .components import *
 from .core import *
 from .dag import *
 from .isolate import *
 from .link_analysis import *
+from .reciprocity import *
 from .shortest_paths import *
 from .traversal import *
+from .tree.recognition import *
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/bipartite/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/bipartite/__init__.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -10,4 +10,5 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from .basic import *
 from .generators import *
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/bipartite/basic.py b/python/nx-cugraph/nx_cugraph/algorithms/bipartite/basic.py
@@ -0,0 +1,31 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import cupy as cp
+
+from nx_cugraph.algorithms.cluster import _triangles
+from nx_cugraph.convert import _to_graph
+from nx_cugraph.utils import networkx_algorithm
+
+__all__ = [
+    "is_bipartite",
+]
+
+
+@networkx_algorithm(plc="triangle_count", version_added="24.02")
+def is_bipartite(G):
+    G = _to_graph(G)
+    # Counting triangles may not be the fastest way to do this, but it is simple.
+    node_ids, triangles, is_single_node = _triangles(
+        G, None, symmetrize="union" if G.is_directed() else None
+    )
+    return int(cp.count_nonzero(triangles)) == 0