Merge branch 'branch-24.10' into sample-buffer

rapidsai · Sep 26, 2024 · cc44b3c · cc44b3c
2 parents 337fa16 + 36c190a
commit cc44b3c
Show file tree

Hide file tree

Showing 61 changed files with 1,434 additions and 444 deletions.
diff --git a/benchmarks/nx-cugraph/pytest-based/run-main-benchmarks.sh b/benchmarks/nx-cugraph/pytest-based/run-main-benchmarks.sh
@@ -30,7 +30,7 @@ algos="
 datasets="
    netscience
    email_Eu_core
-   cit_patents
+   cit-patents
    hollywood
    soc-livejournal
 "

diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
@@ -45,8 +45,8 @@ function sed_runner() {
 echo "${NEXT_FULL_TAG}" > VERSION
 
 # Need to distutils-normalize the original version
-NEXT_SHORT_TAG_PEP440=$(python -c "from setuptools.extern import packaging; print(packaging.version.Version('${NEXT_SHORT_TAG}'))")
-NEXT_UCXX_SHORT_TAG_PEP440=$(python -c "from setuptools.extern import packaging; print(packaging.version.Version('${NEXT_UCXX_SHORT_TAG}'))")
+NEXT_SHORT_TAG_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_SHORT_TAG}'))")
+NEXT_UCXX_SHORT_TAG_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_UCXX_SHORT_TAG}'))")
 
 DEPENDENCIES=(
   cudf

diff --git a/ci/run_nx_cugraph_pytests.sh b/ci/run_nx_cugraph_pytests.sh
@@ -6,4 +6,5 @@ set -euo pipefail
 # Support invoking run_nx_cugraph_pytests.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/nx-cugraph/nx_cugraph
 
-pytest --capture=no --cache-clear --benchmark-disable "$@" tests
+NX_CUGRAPH_USE_COMPAT_GRAPHS=False pytest --capture=no --cache-clear --benchmark-disable "$@" tests
+NX_CUGRAPH_USE_COMPAT_GRAPHS=True pytest --capture=no --cache-clear --benchmark-disable "$@" tests
diff --git a/ci/test_python.sh b/ci/test_python.sh
@@ -108,7 +108,7 @@ echo "nx-cugraph coverage from networkx tests: $_coverage"
 echo $_coverage | awk '{ if ($NF == "0.0%") exit 1 }'
 # Ensure all algorithms were called by comparing covered lines to function lines.
 # Run our tests again (they're fast enough) to add their coverage, then create coverage.json
-pytest \
+NX_CUGRAPH_USE_COMPAT_GRAPHS=False pytest \
   --pyargs nx_cugraph \
   --config-file=../pyproject.toml \
   --cov-config=../pyproject.toml \

diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
@@ -37,6 +37,7 @@ else
     DASK_DISTRIBUTED__SCHEDULER__WORKER_TTL="1000s" \
     DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT="1000s" \
     DASK_CUDA_WAIT_WORKERS_MIN_TIMEOUT="1000s" \
+    NX_CUGRAPH_USE_COMPAT_GRAPHS=False \
     python -m pytest \
        -v \
        --import-mode=append \

diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -33,7 +33,7 @@ dependencies:
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0
 - nbsphinx
-- nccl>=2.9.9
+- nccl>=2.18.1.1
 - networkx>=2.5.1
 - networkx>=3.0
 - ninja
@@ -44,7 +44,6 @@ dependencies:
 - nvcc_linux-64=11.8
 - ogb
 - openmpi
-- packaging>=21
 - pandas
 - pre-commit
 - pydantic

diff --git a/conda/environments/all_cuda-125_arch-x86_64.yaml b/conda/environments/all_cuda-125_arch-x86_64.yaml
@@ -39,7 +39,7 @@ dependencies:
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0
 - nbsphinx
-- nccl>=2.9.9
+- nccl>=2.18.1.1
 - networkx>=2.5.1
 - networkx>=3.0
 - ninja
@@ -49,7 +49,6 @@ dependencies:
 - numpydoc
 - ogb
 - openmpi
-- packaging>=21
 - pandas
 - pre-commit
 - pydantic

diff --git a/conda/recipes/libcugraph/conda_build_config.yaml b/conda/recipes/libcugraph/conda_build_config.yaml
@@ -17,7 +17,7 @@ doxygen_version:
   - ">=1.8.11"
 
 nccl_version:
-  - ">=2.9.9"
+  - ">=2.18.1.1"
 
 c_stdlib:
   - sysroot

diff --git a/cpp/include/cugraph/mtmg/instance_manager.hpp b/cpp/include/cugraph/mtmg/instance_manager.hpp
@@ -20,6 +20,8 @@
 
 #include <raft/comms/std_comms.hpp>
 
+#include <nccl.h>
+
 #include <vector>
 
 namespace cugraph {

diff --git a/cpp/include/cugraph/mtmg/resource_manager.hpp b/cpp/include/cugraph/mtmg/resource_manager.hpp
@@ -27,6 +27,8 @@
 #include <rmm/mr/device/owning_wrapper.hpp>
 #include <rmm/mr/device/pool_memory_resource.hpp>
 
+#include <nccl.h>
+
 #include <execution>
 
 namespace cugraph {

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
@@ -220,6 +220,7 @@ function(ConfigureTestMG CMAKE_TEST_NAME)
         GPUS ${GPU_COUNT}
         PERCENT 100
         INSTALL_COMPONENT_SET testing_mg
+        INSTALL_TARGET ${CMAKE_TEST_NAME}
     )
     set_tests_properties(${CMAKE_TEST_NAME} PROPERTIES LABELS "CUGRAPH_MG")
 
@@ -302,6 +303,7 @@ function(ConfigureCTestMG CMAKE_TEST_NAME)
         GPUS ${GPU_COUNT}
         PERCENT 100
         INSTALL_COMPONENT_SET testing_mg
+        INSTALL_TARGET ${CMAKE_TEST_NAME}
     )
     set_tests_properties(${CMAKE_TEST_NAME} PROPERTIES LABELS "CUGRAPH_C_MG")
 

diff --git a/cpp/tests/mtmg/multi_node_threaded_test.cu b/cpp/tests/mtmg/multi_node_threaded_test.cu
@@ -39,6 +39,7 @@
 #include <thrust/unique.h>
 
 #include <gtest/gtest.h>
+#include <nccl.h>
 
 #include <filesystem>
 #include <fstream>

diff --git a/dependencies.yaml b/dependencies.yaml
@@ -530,7 +530,7 @@ dependencies:
           - aiohttp
           - fsspec>=0.6.0
           - requests
-          - nccl>=2.9.9
+          - nccl>=2.18.1.1
           - ucx-proc=*=gpu
           - &ucx_py_unsuffixed ucx-py==0.40.*,>=0.0.0a0
       - output_types: pyproject
@@ -688,7 +688,6 @@ dependencies:
     common:
       - output_types: [conda, pyproject]
         packages:
-          - packaging>=21
             # not needed by nx-cugraph tests, but is required for running networkx tests
           - pytest-mpl
   cugraph_dgl_dev:

diff --git a/python/cugraph-dgl/cugraph_dgl/dataloading/neighbor_sampler.py b/python/cugraph-dgl/cugraph_dgl/dataloading/neighbor_sampler.py
@@ -18,7 +18,7 @@
 
 from typing import Sequence, Optional, Union, List, Tuple, Iterator
 
-from cugraph.gnn import UniformNeighborSampler, DistSampleWriter
+from cugraph.gnn import UniformNeighborSampler, BiasedNeighborSampler, DistSampleWriter
 from cugraph.utilities.utils import import_optional
 
 import cugraph_dgl
@@ -93,7 +93,6 @@ def __init__(
             If provided, the probability of each neighbor being
             sampled is proportional to the edge feature
             with the given name.  Mutually exclusive with mask.
-            Currently unsupported.
         mask: str
             Optional.
             If proivided, only neighbors where the edge mask
@@ -133,10 +132,6 @@ def __init__(
             raise NotImplementedError(
                 "Edge masking is currently unsupported by cuGraph-DGL"
             )
-        if prob:
-            raise NotImplementedError(
-                "Edge masking is currently unsupported by cuGraph-DGL"
-            )
         if prefetch_edge_feats:
             warnings.warn("'prefetch_edge_feats' is ignored by cuGraph-DGL")
         if prefetch_node_feats:
@@ -146,6 +141,8 @@ def __init__(
         if fused:
             warnings.warn("'fused' is ignored by cuGraph-DGL")
 
+        self.__prob_attr = prob
+
         self.fanouts = fanouts_per_layer
         reverse_fanouts = fanouts_per_layer.copy()
         reverse_fanouts.reverse()
@@ -180,8 +177,14 @@ def sample(
             format=kwargs.pop("format", "parquet"),
         )
 
-        ds = UniformNeighborSampler(
-            g._graph(self.edge_dir),
+        sampling_clx = (
+            UniformNeighborSampler
+            if self.__prob_attr is None
+            else BiasedNeighborSampler
+        )
+
+        ds = sampling_clx(
+            g._graph(self.edge_dir, prob_attr=self.__prob_attr),
             writer,
             compression="CSR",
             fanout=self._reversed_fanout_vals,