Merge branch 'branch-23.12' into branch-24.02-merge-branch-23.12

.github/workflows/build.yaml

@@ -22,7 +22,7 @@ on:
         default: nightly
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event_name }}
   cancel-in-progress: true
 
 jobs:

ci/build_wheel.sh

@@ -40,8 +40,11 @@ for dep in rmm cudf raft-dask pylibcugraph pylibraft ucx-py; do
     sed -r -i "s/${dep}==(.*)\"/${dep}${PACKAGE_CUDA_SUFFIX}==\1${alpha_spec}\"/g" ${pyproject_file}
 done
 
-# dask-cuda doesn't get a suffix, but it does get an alpha spec.
-sed -r -i "s/dask-cuda==(.*)\"/dask-cuda==\1${alpha_spec}\"/g" ${pyproject_file}
+# dask-cuda & rapids-dask-dependency doesn't get a suffix, but it does get an alpha spec.
+for dep in dask-cuda rapids-dask-dependency; do
+    sed -r -i "s/${dep}==(.*)\"/${dep}==\1${alpha_spec}\"/g" ${pyproject_file}
+done
+
 
 if [[ $PACKAGE_CUDA_SUFFIX == "-cu12" ]]; then
     sed -i "s/cupy-cuda11x/cupy-cuda12x/g" ${pyproject_file}

ci/release/update-version.sh

@@ -88,6 +88,7 @@ DEPENDENCIES=(
   raft-dask
   rmm
   ucx-py
+  rapids-dask-dependency
 )
 for DEP in "${DEPENDENCIES[@]}"; do
   for FILE in dependencies.yaml conda/environments/*.yaml python/cugraph-{pyg,dgl}/conda/*.yaml; do

ci/test_python.sh

@@ -197,27 +197,26 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     conda activate test_cugraph_pyg
     set -u
 
-    # Install pytorch
+    # Will automatically install built dependencies of cuGraph-PyG
     rapids-mamba-retry install \
-      --force-reinstall \
-      --channel pyg \
+      --channel "${CPP_CHANNEL}" \
+      --channel "${PYTHON_CHANNEL}" \
       --channel pytorch \
       --channel nvidia \
-      'pyg=2.3' \
-      'pytorch=2.0.0' \
-      'pytorch-cuda=11.8'
+      --channel pyg \
+      --channel rapidsai-nightly \
+      "cugraph-pyg" \
+      "pytorch>=2.0,<2.1" \
+      "pytorch-cuda=11.8"
 
     # Install pyg dependencies (which requires pip)
-    pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.0.0+cu118.html
-
-    rapids-mamba-retry install \
-      --channel "${CPP_CHANNEL}" \
-      --channel "${PYTHON_CHANNEL}" \
-      libcugraph \
-      pylibcugraph \
-      pylibcugraphops \
-      cugraph \
-      cugraph-pyg
+    pip install \
+        pyg_lib \
+        torch_scatter \
+        torch_sparse \
+        torch_cluster \
+        torch_spline_conv \
+      -f https://data.pyg.org/whl/torch-2.0.0+cu118.html
 
     rapids-print-env
 

ci/test_wheel_cugraph.sh

@@ -8,7 +8,4 @@ RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
 RAPIDS_PY_WHEEL_NAME="pylibcugraph_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./local-pylibcugraph-dep
 python -m pip install --no-deps ./local-pylibcugraph-dep/pylibcugraph*.whl
 
-# Always install latest dask for testing
-python -m pip install git+https://github.com/dask/dask.git@main git+https://github.com/dask/distributed.git@main
-
 ./ci/test_wheel.sh cugraph python/cugraph

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -20,11 +20,8 @@ dependencies:
 - cupy>=12.0.0
 - cxx-compiler
 - cython>=3.0.0
-- dask-core>=2023.9.2
 - dask-cuda==24.2.*
 - dask-cudf==24.2.*
-- dask>=2023.7.1
-- distributed>=2023.7.1
 - doxygen
 - fsspec>=0.6.0
 - gcc_linux-64=11.*
@@ -62,6 +59,7 @@ dependencies:
 - pytest-xdist
 - python-louvain
 - raft-dask==24.2.*
+- rapids-dask-dependency==24.2.*
 - recommonmark
 - requests
 - rmm==24.2.*

conda/environments/all_cuda-120_arch-x86_64.yaml

@@ -20,11 +20,8 @@ dependencies:
 - cupy>=12.0.0
 - cxx-compiler
 - cython>=3.0.0
-- dask-core>=2023.9.2
 - dask-cuda==24.2.*
 - dask-cudf==24.2.*
-- dask>=2023.7.1
-- distributed>=2023.7.1
 - doxygen
 - fsspec>=0.6.0
 - gcc_linux-64=11.*
@@ -61,6 +58,7 @@ dependencies:
 - pytest-xdist
 - python-louvain
 - raft-dask==24.2.*
+- rapids-dask-dependency==24.2.*
 - recommonmark
 - requests
 - rmm==24.2.*

conda/recipes/cugraph-dgl/meta.yaml

@@ -26,7 +26,7 @@ requirements:
     - dgl >=1.1.0.cu*
     - numba >=0.57
     - numpy >=1.21
-    - pylibcugraphops ={{ version }}
+    - pylibcugraphops ={{ minor_version }}
     - python
     - pytorch
 

conda/recipes/cugraph-pyg/meta.yaml

@@ -26,15 +26,15 @@ requirements:
     - python
     - scikit-build >=0.13.1
   run:
-    - distributed >=2023.9.2
+    - rapids-dask-dependency ={{ minor_version }}
     - numba >=0.57
     - numpy >=1.21
     - python
     - pytorch >=2.0
     - cupy >=12.0.0
     - cugraph ={{ version }}
-    - pylibcugraphops ={{ version }}
-    - pyg >=2.3,<2.4
+    - pylibcugraphops ={{ minor_version }}
+    - pyg >=2.3,<2.5
 
 tests:
   imports:

conda/recipes/cugraph-service/meta.yaml

@@ -59,10 +59,10 @@ outputs:
         - cupy >=12.0.0
         - dask-cuda ={{ minor_version }}
         - dask-cudf ={{ minor_version }}
-        - distributed >=2023.9.2
         - numba >=0.57
         - numpy >=1.21
         - python
+        - rapids-dask-dependency ={{ minor_version }}
         - thriftpy2 >=0.4.15
         - ucx-py {{ ucx_py_version }}
 

conda/recipes/cugraph/meta.yaml

@@ -76,15 +76,13 @@ requirements:
     - cupy >=12.0.0
     - dask-cuda ={{ minor_version }}
     - dask-cudf ={{ minor_version }}
-    - dask >=2023.9.2
-    - dask-core >=2023.9.2
-    - distributed >=2023.9.2
     - fsspec>=0.6.0
     - libcugraph ={{ version }}
     - pylibcugraph ={{ version }}
     - pylibraft ={{ minor_version }}
     - python
     - raft-dask ={{ minor_version }}
+    - rapids-dask-dependency ={{ minor_version }}
     - requests
     - ucx-proc=*=gpu
     - ucx-py {{ ucx_py_version }}

cpp/CMakeLists.txt

@@ -153,6 +153,11 @@ 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.
+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)
 include(${rapids-cmake-dir}/cpm/cuco.cmake)

dependencies.yaml

@@ -377,15 +377,13 @@ dependencies:
     common:
       - output_types: [conda, pyproject]
         packages:
-          - &dask dask>=2023.7.1
-          - &distributed distributed>=2023.7.1
+          - &dask rapids-dask-dependency==24.2.*
           - &dask_cuda dask-cuda==24.2.*
           - &numba numba>=0.57
           - &ucx_py ucx-py==0.36.*
       - output_types: conda
         packages:
           - aiohttp
-          - &dask-core_conda dask-core>=2023.9.2
           - fsspec>=0.6.0
           - libcudf==24.2.*
           - requests
@@ -431,14 +429,10 @@ dependencies:
         packages:
           - *dask
           - *dask_cuda
-          - *distributed
           - *numba
           - *numpy
           - *thrift
           - *ucx_py
-      - output_types: conda
-        packages:
-          - *dask-core_conda
       - output_types: pyproject
         packages:
           - *cugraph
@@ -503,9 +497,9 @@ dependencies:
       - output_types: [conda]
         packages:
           - cugraph==24.2.*
-          - pytorch==2.0
+          - pytorch>=2.0
           - pytorch-cuda==11.8
-          - pyg=2.3.1=*torch_2.0.0*cu118*
+          - pyg>=2.4.0
 
   depends_on_rmm:
     common:

python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml

@@ -13,13 +13,13 @@ dependencies:
 - cugraph==24.2.*
 - pandas
 - pre-commit
-- pyg=2.3.1=*torch_2.0.0*cu118*
+- pyg>=2.4.0
 - pylibcugraphops==24.2.*
 - pytest
 - pytest-benchmark
 - pytest-cov
 - pytest-xdist
 - pytorch-cuda==11.8
-- pytorch==2.0
+- pytorch>=2.0
 - scipy
 name: cugraph_pyg_dev_cuda-118

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

@@ -210,7 +210,10 @@ class EXPERIMENTAL__CuGraphStore:
     def __init__(
         self,
         F: cugraph.gnn.FeatureStore,
-        G: Union[Dict[str, Tuple[TensorType]], Dict[str, int]],
+        G: Union[
+            Dict[Tuple[str, str, str], Tuple[TensorType]],
+            Dict[Tuple[str, str, str], int],
+        ],
         num_nodes_dict: Dict[str, int],
         *,
         multi_gpu: bool = False,
@@ -744,7 +747,7 @@ def _subgraph(self, edge_types: List[tuple] = None) -> cugraph.MultiGraph:
 
     def _get_vertex_groups_from_sample(
         self, nodes_of_interest: TensorType, is_sorted: bool = False
-    ) -> dict:
+    ) -> Dict[str, torch.Tensor]:
         """
         Given a tensor of nodes of interest, this
         method a single dictionary, noi_index.
@@ -808,7 +811,10 @@ def _get_sample_from_vertex_groups(
 
     def _get_renumbered_edge_groups_from_sample(
         self, sampling_results: cudf.DataFrame, noi_index: dict
-    ) -> Tuple[dict, dict]:
+    ) -> Tuple[
+        Dict[Tuple[str, str, str], torch.Tensor],
+        Tuple[Dict[Tuple[str, str, str], torch.Tensor]],
+    ]:
         """
         Given a cudf (NOT dask_cudf) DataFrame of sampling results and a dictionary
         of non-renumbered vertex ids grouped by vertex type, this method

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

@@ -15,6 +15,7 @@
 
 import os
 import re
+import warnings
 
 import cupy
 import cudf
@@ -159,23 +160,34 @@ def __init__(
         if batch_size is None or batch_size < 1:
             raise ValueError("Batch size must be >= 1")
 
-        self.__directory = tempfile.TemporaryDirectory(dir=directory)
+        self.__directory = (
+            tempfile.TemporaryDirectory() if directory is None else directory
+        )
 
         if isinstance(num_neighbors, dict):
             raise ValueError("num_neighbors dict is currently unsupported!")
 
-        renumber = (
-            True
-            if (
-                (len(self.__graph_store.node_types) == 1)
-                and (len(self.__graph_store.edge_types) == 1)
+        if "renumber" in kwargs:
+            warnings.warn(
+                "Setting renumbering manually could result in invalid output,"
+                " please ensure you intended to do this."
+            )
+            renumber = kwargs.pop("renumber")
+        else:
+            renumber = (
+                True
+                if (
+                    (len(self.__graph_store.node_types) == 1)
+                    and (len(self.__graph_store.edge_types) == 1)
+                )
+                else False
             )
-            else False
-        )
 
         bulk_sampler = BulkSampler(
             batch_size,
-            self.__directory.name,
+            self.__directory
+            if isinstance(self.__directory, str)
+            else self.__directory.name,
             self.__graph_store._subgraph(edge_types),
             fanout_vals=num_neighbors,
             with_replacement=replace,
@@ -219,7 +231,13 @@ def __init__(
             )
 
         bulk_sampler.flush()
-        self.__input_files = iter(os.listdir(self.__directory.name))
+        self.__input_files = iter(
+            os.listdir(
+                self.__directory
+                if isinstance(self.__directory, str)
+                else self.__directory.name
+            )
+        )
 
     def __next__(self):
         from time import perf_counter
@@ -423,9 +441,6 @@ def __next__(self):
                 sampler_output.edge,
             )
         else:
-            if self.__graph_store.order == "CSR":
-                raise ValueError("CSR format incompatible with CSC output")
-
             out = filter_cugraph_store_csc(
                 self.__feature_store,
                 self.__graph_store,
@@ -437,11 +452,8 @@ def __next__(self):
 
         # Account for CSR format in cuGraph vs. CSC format in PyG
         if self.__coo and self.__graph_store.order == "CSC":
-            for node_type in out.edge_index_dict:
-                out[node_type].edge_index[0], out[node_type].edge_index[1] = (
-                    out[node_type].edge_index[1],
-                    out[node_type].edge_index[0],
-                )
+            for edge_type in out.edge_index_dict:
+                out[edge_type].edge_index = out[edge_type].edge_index.flip(dims=[0])
 
         out.set_value_dict("num_sampled_nodes", sampler_output.num_sampled_nodes)
         out.set_value_dict("num_sampled_edges", sampler_output.num_sampled_edges)

python/cugraph-pyg/cugraph_pyg/tests/mg/test_mg_cugraph_store.py

@@ -120,7 +120,7 @@ def test_get_edge_index(graph, edge_index_type, dask_client):
             G[et][0] = dask_cudf.from_cudf(cudf.Series(G[et][0]), npartitions=1)
             G[et][1] = dask_cudf.from_cudf(cudf.Series(G[et][1]), npartitions=1)
 
-    cugraph_store = CuGraphStore(F, G, N, multi_gpu=True)
+    cugraph_store = CuGraphStore(F, G, N, order="CSC", multi_gpu=True)
 
     for pyg_can_edge_type in G:
         src, dst = cugraph_store.get_edge_index(