Merge remote-tracking branch 'upstream/branch-24.02' into branch-24.0…

…2-nx_cugraph_benchmarks

benchmarks/cugraph/standalone/bulk_sampling/cugraph_bulk_sampling.py

@@ -28,7 +28,7 @@
 )
 
 from cugraph.structure.symmetrize import symmetrize
-from cugraph.experimental.gnn import BulkSampler
+from cugraph.gnn import BulkSampler
 
 import cugraph
 

ci/build_python.sh

@@ -19,7 +19,7 @@ echo "${version}" > VERSION
 rapids-logger "Begin py build"
 
 package_dir="python"
-for package_name in pylibcugraph cugraph nx-cugraph cugraph-pyg cugraph-dgl; do 
+for package_name in pylibcugraph cugraph nx-cugraph cugraph-pyg cugraph-dgl; do
   underscore_package_name=$(echo "${package_name}" | tr "-" "_")
   sed -i "/^__git_commit__/ s/= .*/= \"${git_commit}\"/g" "${package_dir}/${package_name}/${underscore_package_name}/_version.py"
 done

ci/test_cpp.sh

@@ -19,6 +19,7 @@ conda activate test
 set -u
 
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+
 RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
 mkdir -p "${RAPIDS_TESTS_DIR}"
 

ci/test_python.sh

@@ -111,6 +111,11 @@ popd
 rapids-logger "pytest networkx using nx-cugraph backend"
 pushd python/nx-cugraph
 ./run_nx_tests.sh
+# Individually run tests that are skipped above b/c they may run out of memory
+PYTEST_NO_SKIP=True ./run_nx_tests.sh --cov-append -k "TestDAG and test_antichains"
+PYTEST_NO_SKIP=True ./run_nx_tests.sh --cov-append -k "TestMultiDiGraph_DAGLCA and test_all_pairs_lca_pairs_without_lca"
+PYTEST_NO_SKIP=True ./run_nx_tests.sh --cov-append -k "TestDAGLCA and test_all_pairs_lca_pairs_without_lca"
+PYTEST_NO_SKIP=True ./run_nx_tests.sh --cov-append -k "TestEfficiency and test_using_ego_graph"
 # run_nx_tests.sh outputs coverage data, so check that total coverage is >0.0%
 # in case nx-cugraph failed to load but fallback mode allowed the run to pass.
 _coverage=$(coverage report|grep "^TOTAL")
@@ -208,7 +213,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
       "cugraph-pyg" \
       "pytorch>=2.0,<2.1" \
       "pytorch-cuda=11.8"
-    
+
     # Install pyg dependencies (which requires pip)
     pip install \
         pyg_lib \

cpp/CMakeLists.txt

@@ -142,8 +142,8 @@ rapids_cpm_init()
 ###
 # Linking to the `raft::raft` target implicitly links cugraph targets to the
 # following public header-only  raft dependencies:
+# * CCCL
 # * RMM
-# * Thrust
 # * GTest/GMock
 #
 # The CMakeLists.txt for each of these projects are properly configured
@@ -153,16 +153,14 @@ rapids_cpm_init()
 # lags behind.
 ###
 
-# Need to make sure rmm is found before cuco so that rmm patches the libcudacxx
-# directory to be found by cuco.
+# Need CCCL, then rmm, then cuCollections, then RAFT.
+# This ensures that libraries can be overridden for testing.
+include(cmake/thirdparty/get_cccl.cmake)
 include(${rapids-cmake-dir}/cpm/rmm.cmake)
 rapids_cpm_rmm(BUILD_EXPORT_SET cugraph-exports
-               INSTALL_EXPORT_SET  cugraph-exports)
-# Putting this before raft to override RAFT from pulling them in.
-include(cmake/thirdparty/get_libcudacxx.cmake)
+               INSTALL_EXPORT_SET cugraph-exports)
 include(${rapids-cmake-dir}/cpm/cuco.cmake)
 rapids_cpm_cuco(BUILD_EXPORT_SET cugraph-exports INSTALL_EXPORT_SET cugraph-exports)
-
 include(cmake/thirdparty/get_raft.cmake)
 
 if(USE_CUGRAPH_OPS)

cpp/cmake/thirdparty/get_libcudacxx.cmake → cpp/cmake/thirdparty/get_cccl.cmake

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2020-2022, NVIDIA CORPORATION.
+# Copyright (c) 2020-2023, 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
@@ -12,12 +12,10 @@
 # the License.
 # =============================================================================
 
-# This function finds libcudacxx and sets any additional necessary environment variables.
-function(find_and_configure_libcudacxx)
-  include(${rapids-cmake-dir}/cpm/libcudacxx.cmake)
-
-  rapids_cpm_libcudacxx(BUILD_EXPORT_SET cugraph-exports)
-
+# This function finds CCCL and sets any additional necessary environment variables.
+function(find_and_configure_cccl)
+  include(${rapids-cmake-dir}/cpm/cccl.cmake)
+  rapids_cpm_cccl(BUILD_EXPORT_SET cugraph-exports INSTALL_EXPORT_SET cugraph-exports)
 endfunction()
 
-find_and_configure_libcudacxx()
+find_and_configure_cccl()

cpp/src/community/flatten_dendrogram.hpp

@@ -75,7 +75,7 @@ void leiden_partition_at_level(raft::handle_t const& handle,
     thrust::make_counting_iterator<size_t>(0),
     thrust::make_counting_iterator<size_t>((level - 1) / 2),
     [&handle, &dendrogram, &local_vertex_ids_v, &d_partition, local_num_verts](size_t l) {
-      cugraph::relabel<vertex_t, false>(
+      cugraph::relabel<vertex_t, multi_gpu>(
         handle,
         std::tuple<vertex_t const*, vertex_t const*>(dendrogram.get_level_ptr_nocheck(2 * l + 1),
                                                      dendrogram.get_level_ptr_nocheck(2 * l + 2)),

python/cugraph-dgl/cugraph_dgl/dataloading/dataloader.py

@@ -17,7 +17,7 @@
 import cupy as cp
 import cudf
 from cugraph.utilities.utils import import_optional
-from cugraph.experimental import BulkSampler
+from cugraph.gnn import BulkSampler
 from dask.distributed import default_client, Event
 from cugraph_dgl.dataloading import (
     HomogenousBulkSamplerDataset,

python/cugraph-pyg/cugraph_pyg/data/__init__.py

@@ -11,8 +11,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cugraph.utilities.api_tools import experimental_warning_wrapper
-
-from cugraph_pyg.data.cugraph_store import EXPERIMENTAL__CuGraphStore
-
-CuGraphStore = experimental_warning_wrapper(EXPERIMENTAL__CuGraphStore)
+from cugraph_pyg.data.cugraph_store import CuGraphStore

python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py

@@ -199,7 +199,7 @@ def cast(cls, *args, **kwargs):
         return cls(*args, **kwargs)
 
 
-class EXPERIMENTAL__CuGraphStore:
+class CuGraphStore:
     """
     Duck-typed version of PyG's GraphStore and FeatureStore.
     """

python/cugraph-pyg/cugraph_pyg/loader/__init__.py

@@ -11,14 +11,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cugraph.utilities.api_tools import experimental_warning_wrapper
+from cugraph_pyg.loader.cugraph_node_loader import CuGraphNeighborLoader
 
-from cugraph_pyg.loader.cugraph_node_loader import EXPERIMENTAL__CuGraphNeighborLoader
-
-CuGraphNeighborLoader = experimental_warning_wrapper(
-    EXPERIMENTAL__CuGraphNeighborLoader
-)
-
-from cugraph_pyg.loader.cugraph_node_loader import EXPERIMENTAL__BulkSampleLoader
-
-BulkSampleLoader = experimental_warning_wrapper(EXPERIMENTAL__BulkSampleLoader)
+from cugraph_pyg.loader.cugraph_node_loader import BulkSampleLoader

python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py

@@ -20,7 +20,7 @@
 import cupy
 import cudf
 
-from cugraph.experimental.gnn import BulkSampler
+from cugraph.gnn import BulkSampler
 from cugraph.utilities.utils import import_optional, MissingModule
 
 from cugraph_pyg.data import CuGraphStore
@@ -42,7 +42,7 @@
 )
 
 
-class EXPERIMENTAL__BulkSampleLoader:
+class BulkSampleLoader:
 
     __ex_parquet_file = re.compile(r"batch=([0-9]+)\-([0-9]+)\.parquet")
 
@@ -478,7 +478,7 @@ def __iter__(self):
         return self
 
 
-class EXPERIMENTAL__CuGraphNeighborLoader:
+class CuGraphNeighborLoader:
     def __init__(
         self,
         data: Union[CuGraphStore, Tuple[CuGraphStore, CuGraphStore]],
@@ -527,7 +527,7 @@ def batch_size(self) -> int:
         return self.__batch_size
 
     def __iter__(self):
-        self.current_loader = EXPERIMENTAL__BulkSampleLoader(
+        self.current_loader = BulkSampleLoader(
             self.__feature_store,
             self.__graph_store,
             self.__input_nodes,

python/cugraph/cugraph/experimental/__init__.py

@@ -48,9 +48,9 @@
     experimental_warning_wrapper(EXPERIMENTAL__find_bicliques)
 )
 
-from cugraph.gnn.data_loading import EXPERIMENTAL__BulkSampler
+from cugraph.gnn.data_loading import BulkSampler
 
-BulkSampler = experimental_warning_wrapper(EXPERIMENTAL__BulkSampler)
+BulkSampler = promoted_experimental_warning_wrapper(BulkSampler)
 
 
 from cugraph.link_prediction.jaccard import jaccard, jaccard_coefficient

python/cugraph/cugraph/experimental/gnn/__init__.py

@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cugraph.gnn.data_loading import EXPERIMENTAL__BulkSampler
-from cugraph.utilities.api_tools import experimental_warning_wrapper
+from cugraph.gnn.data_loading import BulkSampler
+from cugraph.utilities.api_tools import promoted_experimental_warning_wrapper
 
-BulkSampler = experimental_warning_wrapper(EXPERIMENTAL__BulkSampler)
+BulkSampler = promoted_experimental_warning_wrapper(BulkSampler)

python/cugraph/cugraph/gnn/__init__.py

@@ -12,3 +12,4 @@
 # limitations under the License.
 
 from .feature_storage.feat_storage import FeatureStore
+from .data_loading.bulk_sampler import BulkSampler

python/cugraph/cugraph/gnn/data_loading/__init__.py

@@ -11,4 +11,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cugraph.gnn.data_loading.bulk_sampler import EXPERIMENTAL__BulkSampler
+from cugraph.gnn.data_loading.bulk_sampler import BulkSampler

python/cugraph/cugraph/gnn/data_loading/bulk_sampler.py

@@ -31,7 +31,7 @@
 import time
 
 
-class EXPERIMENTAL__BulkSampler:
+class BulkSampler:
     """
     Performs sampling based on input seeds grouped into batches by
     a batch id.  Writes the output minibatches to parquet, with
@@ -158,7 +158,7 @@ def add_batches(
         Examples
         --------
         >>> import cudf
-        >>> from cugraph.experimental.gnn import BulkSampler
+        >>> from cugraph.gnn import BulkSampler
         >>> from cugraph.datasets import karate
         >>> import tempfile
         >>> df = cudf.DataFrame({

python/cugraph/cugraph/tests/sampling/test_bulk_sampler.py

@@ -17,7 +17,7 @@
 import cupy
 import cugraph
 from cugraph.datasets import karate, email_Eu_core
-from cugraph.experimental.gnn import BulkSampler
+from cugraph.gnn import BulkSampler
 from cugraph.utilities.utils import create_directory_with_overwrite
 
 import os

python/cugraph/cugraph/tests/sampling/test_bulk_sampler_mg.py

@@ -22,7 +22,7 @@
 import cugraph
 import dask_cudf
 from cugraph.datasets import karate, email_Eu_core
-from cugraph.experimental import BulkSampler
+from cugraph.gnn import BulkSampler
 from cugraph.utilities.utils import create_directory_with_overwrite
 
 

python/nx-cugraph/_nx_cugraph/__init__.py

@@ -29,8 +29,14 @@
     # "description": "TODO",
     "functions": {
         # BEGIN: functions
+        "ancestors",
         "barbell_graph",
         "betweenness_centrality",
+        "bfs_edges",
+        "bfs_layers",
+        "bfs_predecessors",
+        "bfs_successors",
+        "bfs_tree",
         "bull_graph",
         "caveman_graph",
         "chvatal_graph",
@@ -44,6 +50,8 @@
         "davis_southern_women_graph",
         "degree_centrality",
         "desargues_graph",
+        "descendants",
+        "descendants_at_distance",
         "diamond_graph",
         "dodecahedral_graph",
         "edge_betweenness_centrality",
@@ -53,6 +61,7 @@
         "from_pandas_edgelist",
         "from_scipy_sparse_array",
         "frucht_graph",
+        "generic_bfs_edges",
         "heawood_graph",
         "hits",
         "house_graph",
@@ -99,9 +108,14 @@
     "extra_docstrings": {
         # BEGIN: extra_docstrings
         "betweenness_centrality": "`weight` parameter is not yet supported.",
+        "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.",
         "eigenvector_centrality": "`nstart` parameter is not used, but it is checked for validity.",
         "from_pandas_edgelist": "cudf.DataFrame inputs also supported.",
+        "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."

python/nx-cugraph/nx_cugraph/algorithms/__init__.py

@@ -17,11 +17,14 @@
     components,
     link_analysis,
     shortest_paths,
+    traversal,
 )
 from .bipartite import complete_bipartite_graph
 from .centrality import *
 from .components import *
 from .core import *
+from .dag import *
 from .isolate import *
 from .link_analysis import *
 from .shortest_paths import *
+from .traversal import *

python/nx-cugraph/nx_cugraph/algorithms/core.py

@@ -15,7 +15,12 @@
 import pylibcugraph as plc
 
 import nx_cugraph as nxcg
-from nx_cugraph.utils import _get_int_dtype, networkx_algorithm, not_implemented_for
+from nx_cugraph.utils import (
+    _get_int_dtype,
+    index_dtype,
+    networkx_algorithm,
+    not_implemented_for,
+)
 
 __all__ = ["k_truss"]
 
@@ -81,10 +86,8 @@ def k_truss(G, k):
         edge_values = {key: val[edge_indices] for key, val in G.edge_values.items()}
         edge_masks = {key: val[edge_indices] for key, val in G.edge_masks.items()}
     # Renumber step 2: edge indices
-    mapper = cp.zeros(len(G), src_indices.dtype)
-    mapper[node_indices] = cp.arange(node_indices.size, dtype=mapper.dtype)
-    src_indices = mapper[src_indices]
-    dst_indices = mapper[dst_indices]
+    src_indices = cp.searchsorted(node_indices, src_indices).astype(index_dtype)
+    dst_indices = cp.searchsorted(node_indices, dst_indices).astype(index_dtype)
     # Renumber step 3: node values
     node_values = {key: val[node_indices] for key, val in G.node_values.items()}
     node_masks = {key: val[node_indices] for key, val in G.node_masks.items()}