Merge branch 'nxcg-notebooks' of https://github.com/nv-rliu/cugraph i…

…nto nxcg-notebooks

ci/test_python.sh

@@ -159,7 +159,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
       cugraph \
       cugraph-dgl \
       'dgl>=1.1.0.cu*,<=2.0.0.cu*' \
-      'pytorch>=2.0' \
+      'pytorch>=2.3,<2.4' \
       'cuda-version=11.8'
 
     rapids-print-env
@@ -198,10 +198,10 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     # TODO re-enable logic once CUDA 12 is testable
     #if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     CONDA_CUDA_VERSION="11.8"
-    PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu118.html"
+    PYG_URL="https://data.pyg.org/whl/torch-2.3.0+cu118.html"
     #else
     #  CONDA_CUDA_VERSION="12.1"
-    #  PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu121.html"
+    #  PYG_URL="https://data.pyg.org/whl/torch-2.3.0+cu121.html"
     #fi
 
     # Will automatically install built dependencies of cuGraph-PyG

ci/test_wheel_cugraph-dgl.sh

@@ -32,18 +32,8 @@ fi
 PYTORCH_URL="https://download.pytorch.org/whl/cu${PYTORCH_CUDA_VER}"
 DGL_URL="https://data.dgl.ai/wheels/cu${PYTORCH_CUDA_VER}/repo.html"
 
-# Starting from 2.2, PyTorch wheels depend on nvidia-nccl-cuxx>=2.19 wheel and
-# dynamically link to NCCL. RAPIDS CUDA 11 CI images have an older NCCL version that
-# might shadow the newer NCCL required by PyTorch during import (when importing
-# `cupy` before `torch`).
-if [[ "${NCCL_VERSION}" < "2.19" ]]; then
-  PYTORCH_VER="2.1.0"
-else
-  PYTORCH_VER="2.3.0"
-fi
-
 rapids-logger "Installing PyTorch and DGL"
-rapids-retry python -m pip install "torch==${PYTORCH_VER}" --index-url ${PYTORCH_URL}
+rapids-retry python -m pip install torch==2.3.0 --index-url ${PYTORCH_URL}
 rapids-retry python -m pip install dgl==2.0.0 --find-links ${DGL_URL}
 
 python -m pytest python/cugraph-dgl/tests

ci/test_wheel_cugraph-pyg.sh

@@ -29,13 +29,13 @@ export CI_RUN=1
 
 if [[ "${CUDA_VERSION}" == "11.8.0" ]]; then
   PYTORCH_URL="https://download.pytorch.org/whl/cu118"
-  PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu118.html"
+  PYG_URL="https://data.pyg.org/whl/torch-2.3.0+cu118.html"
 else
   PYTORCH_URL="https://download.pytorch.org/whl/cu121"
-  PYG_URL="https://data.pyg.org/whl/torch-2.1.0+cu121.html"
+  PYG_URL="https://data.pyg.org/whl/torch-2.3.0+cu121.html"
 fi
 rapids-logger "Installing PyTorch and PyG dependencies"
-rapids-retry python -m pip install torch==2.1.0 --index-url ${PYTORCH_URL}
+rapids-retry python -m pip install torch==2.3.0 --index-url ${PYTORCH_URL}
 rapids-retry python -m pip install "torch-geometric>=2.5,<2.6"
 rapids-retry python -m pip install \
   ogb \

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

@@ -33,13 +33,13 @@ dependencies:
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0
 - nbsphinx
-- nccl>=2.18.1.1
+- nccl>=2.19
 - networkx>=2.5.1
 - networkx>=3.0
 - ninja
 - notebook>=0.5.0
 - numba>=0.57
-- numpy>=1.23,<2.0a0
+- numpy>=1.23,<3.0a0
 - numpydoc
 - nvcc_linux-64=11.8
 - ogb
@@ -57,7 +57,7 @@ dependencies:
 - pytest-mpl
 - pytest-xdist
 - python-louvain
-- pytorch>=2.0,<2.2.0a0
+- pytorch>=2.3,<2.4.0a0
 - raft-dask==24.10.*,>=0.0.0a0
 - rapids-build-backend>=0.3.1,<0.4.0.dev0
 - rapids-dask-dependency==24.10.*,>=0.0.0a0

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

@@ -39,13 +39,13 @@ dependencies:
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0
 - nbsphinx
-- nccl>=2.18.1.1
+- nccl>=2.19
 - networkx>=2.5.1
 - networkx>=3.0
 - ninja
 - notebook>=0.5.0
 - numba>=0.57
-- numpy>=1.23,<2.0a0
+- numpy>=1.23,<3.0a0
 - numpydoc
 - ogb
 - openmpi
@@ -62,7 +62,7 @@ dependencies:
 - pytest-mpl
 - pytest-xdist
 - python-louvain
-- pytorch>=2.0,<2.2.0a0
+- pytorch>=2.3,<2.4.0a0
 - raft-dask==24.10.*,>=0.0.0a0
 - rapids-build-backend>=0.3.1,<0.4.0.dev0
 - rapids-dask-dependency==24.10.*,>=0.0.0a0

conda/recipes/cugraph-dgl/meta.yaml

@@ -27,11 +27,11 @@ requirements:
     - cugraph ={{ version }}
     - dgl >=1.1.0.cu*
     - numba >=0.57
-    - numpy >=1.23,<2.0a0
+    - numpy >=1.23,<3.0a0
     - pylibcugraphops ={{ minor_version }}
     - tensordict >=0.1.2
     - python
-    - pytorch >=2.0
+    - pytorch >=2.3,<2.4.0a0
     - cupy >=12.0.0
 
 tests:

conda/recipes/cugraph-pyg/meta.yaml

@@ -29,9 +29,9 @@ requirements:
   run:
     - rapids-dask-dependency ={{ minor_version }}
     - numba >=0.57
-    - numpy >=1.23,<2.0a0
+    - numpy >=1.23,<3.0a0
     - python
-    - pytorch >=2.0
+    - pytorch >=2.3,<2.4.0a0
     - cupy >=12.0.0
     - cugraph ={{ version }}
     - pylibcugraphops ={{ minor_version }}

conda/recipes/cugraph-service/meta.yaml

@@ -63,7 +63,7 @@ outputs:
         - dask-cuda ={{ minor_version }}
         - dask-cudf ={{ minor_version }}
         - numba >=0.57
-        - numpy >=1.23,<2.0a0
+        - numpy >=1.23,<3.0a0
         - python
         - rapids-dask-dependency ={{ minor_version }}
         - thriftpy2 >=0.4.15,!=0.5.0,!=0.5.1

conda/recipes/libcugraph/conda_build_config.yaml

@@ -17,7 +17,7 @@ doxygen_version:
   - ">=1.8.11"
 
 nccl_version:
-  - ">=2.18.1.1"
+  - ">=2.19"
 
 c_stdlib:
   - sysroot

cpp/src/c_api/graph_generators.cpp

@@ -124,32 +124,41 @@ cugraph_error_code_t cugraph_generate_rmat_edgelists(
 extern "C" cugraph_type_erased_device_array_view_t* cugraph_coo_get_sources(cugraph_coo_t* coo)
 {
   auto internal_pointer = reinterpret_cast<cugraph::c_api::cugraph_coo_t*>(coo);
-  return reinterpret_cast<cugraph_type_erased_device_array_view_t*>(internal_pointer->src_->view());
+  return (internal_pointer->src_) ? reinterpret_cast<cugraph_type_erased_device_array_view_t*>(
+                                      internal_pointer->src_->view())
+                                  : nullptr;
 }
 
 extern "C" cugraph_type_erased_device_array_view_t* cugraph_coo_get_destinations(cugraph_coo_t* coo)
 {
   auto internal_pointer = reinterpret_cast<cugraph::c_api::cugraph_coo_t*>(coo);
-  return reinterpret_cast<cugraph_type_erased_device_array_view_t*>(internal_pointer->dst_->view());
+  return (internal_pointer->dst_) ? reinterpret_cast<cugraph_type_erased_device_array_view_t*>(
+                                      internal_pointer->dst_->view())
+                                  : nullptr;
 }
 
 extern "C" cugraph_type_erased_device_array_view_t* cugraph_coo_get_edge_weights(cugraph_coo_t* coo)
 {
   auto internal_pointer = reinterpret_cast<cugraph::c_api::cugraph_coo_t*>(coo);
-  return reinterpret_cast<cugraph_type_erased_device_array_view_t*>(internal_pointer->wgt_->view());
+  return (internal_pointer->wgt_) ? reinterpret_cast<cugraph_type_erased_device_array_view_t*>(
+                                      internal_pointer->wgt_->view())
+                                  : nullptr;
 }
 
 extern "C" cugraph_type_erased_device_array_view_t* cugraph_coo_get_edge_id(cugraph_coo_t* coo)
 {
   auto internal_pointer = reinterpret_cast<cugraph::c_api::cugraph_coo_t*>(coo);
-  return reinterpret_cast<cugraph_type_erased_device_array_view_t*>(internal_pointer->id_->view());
+  return (internal_pointer->id_) ? reinterpret_cast<cugraph_type_erased_device_array_view_t*>(
+                                     internal_pointer->id_->view())
+                                 : nullptr;
 }
 
 extern "C" cugraph_type_erased_device_array_view_t* cugraph_coo_get_edge_type(cugraph_coo_t* coo)
 {
   auto internal_pointer = reinterpret_cast<cugraph::c_api::cugraph_coo_t*>(coo);
-  return reinterpret_cast<cugraph_type_erased_device_array_view_t*>(
-    internal_pointer->type_->view());
+  return (internal_pointer->type_) ? reinterpret_cast<cugraph_type_erased_device_array_view_t*>(
+                                       internal_pointer->type_->view())
+                                   : nullptr;
 }
 
 extern "C" size_t cugraph_coo_list_size(const cugraph_coo_list_t* coo_list)

dependencies.yaml

@@ -524,13 +524,13 @@ dependencies:
           - &dask rapids-dask-dependency==24.10.*,>=0.0.0a0
           - &dask_cuda dask-cuda==24.10.*,>=0.0.0a0
           - &numba numba>=0.57
-          - &numpy numpy>=1.23,<2.0a0
+          - &numpy numpy>=1.23,<3.0a0
       - output_types: conda
         packages:
           - aiohttp
           - fsspec>=0.6.0
           - requests
-          - nccl>=2.18.1.1
+          - nccl>=2.19
           - ucx-proc=*=gpu
           - &ucx_py_unsuffixed ucx-py==0.40.*,>=0.0.0a0
       - output_types: pyproject
@@ -695,7 +695,9 @@ dependencies:
       - output_types: [conda]
         packages:
           - *cugraph_unsuffixed
-          - pytorch>=2.0
+          # ceiling could be removed when this is fixed:
+          # https://github.com/conda-forge/pytorch-cpu-feedstock/issues/254
+          - &pytorch_conda pytorch>=2.3,<2.4.0a0
           - pytorch-cuda==11.8
           - &tensordict tensordict>=0.1.2
           - dgl>=1.1.0.cu*
@@ -704,7 +706,7 @@ dependencies:
       - output_types: [conda]
         packages:
           - *cugraph_unsuffixed
-          - pytorch>=2.0
+          - *pytorch_conda
           - pytorch-cuda==11.8
           - *tensordict
           - pyg>=2.5,<2.6
@@ -713,7 +715,7 @@ dependencies:
     common:
       - output_types: [conda]
         packages:
-          - &pytorch_unsuffixed pytorch>=2.0,<2.2.0a0
+          - *pytorch_conda
           - torchdata
           - pydantic
           - ogb
@@ -733,7 +735,7 @@ dependencies:
         matrices:
           - matrix: {cuda: "12.*"}
             packages:
-              - &pytorch_pip torch>=2.0,<2.2.0a0
+              - &pytorch_pip torch>=2.3,<2.4.0a0
               - *tensordict
           - matrix: {cuda: "11.*"}
             packages:

python/cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml

@@ -19,7 +19,7 @@ dependencies:
 - pytest-cov
 - pytest-xdist
 - pytorch-cuda==11.8
-- pytorch>=2.0
+- pytorch>=2.3,<2.4.0a0
 - scipy
 - tensordict>=0.1.2
 name: cugraph_dgl_dev_cuda-118

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

@@ -197,10 +197,8 @@ def sample(
 
         if g.is_homogeneous:
             indices = torch.concat(list(indices))
-            ds.sample_from_nodes(indices.long(), batch_size=batch_size)
-            return HomogeneousSampleReader(
-                ds.get_reader(), self.output_format, self.edge_dir
-            )
+            reader = ds.sample_from_nodes(indices.long(), batch_size=batch_size)
+            return HomogeneousSampleReader(reader, self.output_format, self.edge_dir)
 
         raise ValueError(
             "Sampling heterogeneous graphs is currently"

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

@@ -20,7 +20,6 @@
     create_homogeneous_sampled_graphs_from_tensors_csc,
 )
 
-from cugraph.gnn import DistSampleReader
 
 from cugraph.utilities.utils import import_optional
 
@@ -33,14 +32,18 @@ class SampleReader:
     Iterator that processes results from the cuGraph distributed sampler.
     """
 
-    def __init__(self, base_reader: DistSampleReader, output_format: str = "dgl.Block"):
+    def __init__(
+        self,
+        base_reader: Iterator[Tuple[Dict[str, "torch.Tensor"], int, int]],
+        output_format: str = "dgl.Block",
+    ):
         """
         Constructs a new SampleReader.
 
         Parameters
         ----------
-        base_reader: DistSampleReader
-            The reader responsible for loading saved samples produced by
+        base_reader: Iterator[Tuple[Dict[str, "torch.Tensor"], int, int]]
+            The iterator responsible for loading saved samples produced by
             the cuGraph distributed sampler.
         """
         self.__output_format = output_format
@@ -83,7 +86,7 @@ class HomogeneousSampleReader(SampleReader):
 
     def __init__(
         self,
-        base_reader: DistSampleReader,
+        base_reader: Iterator[Tuple[Dict[str, "torch.Tensor"], int, int]],
         output_format: str = "dgl.Block",
         edge_dir="in",
     ):
@@ -92,7 +95,7 @@ def __init__(
 
         Parameters
         ----------
-        base_reader: DistSampleReader
+        base_reader: Iterator[Tuple[Dict[str, "torch.Tensor"], int, int]]
             The reader responsible for loading saved samples produced by
             the cuGraph distributed sampler.
         output_format: str

python/cugraph-dgl/cugraph_dgl/graph.py

@@ -620,9 +620,6 @@ def _get_n_emb(
             )
 
         try:
-            print(
-                u,
-            )
             return self.__ndata_storage[ntype, emb_name].fetch(
                 _cast_to_torch_tensor(u), "cuda"
             )

python/cugraph-dgl/pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 dependencies = [
     "cugraph==24.10.*,>=0.0.0a0",
     "numba>=0.57",
-    "numpy>=1.23,<2.0a0",
+    "numpy>=1.23,<3.0a0",
     "pylibcugraphops==24.10.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
@@ -40,7 +40,7 @@ test = [
     "pytest-xdist",
     "scipy",
     "tensordict>=0.1.2",
-    "torch>=2.0,<2.2.0a0",
+    "torch>=2.3,<2.4.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project.urls]

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

@@ -19,7 +19,7 @@ dependencies:
 - pytest-cov
 - pytest-xdist
 - pytorch-cuda==11.8
-- pytorch>=2.0
+- pytorch>=2.3,<2.4.0a0
 - scipy
 - tensordict>=0.1.2
 name: cugraph_pyg_dev_cuda-118

python/cugraph-pyg/cugraph_pyg/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2023, NVIDIA CORPORATION.
+# Copyright (c) 2019-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
@@ -12,3 +12,8 @@
 # limitations under the License.
 
 from cugraph_pyg._version import __git_commit__, __version__
+
+import cugraph_pyg.data
+import cugraph_pyg.loader
+import cugraph_pyg.sampler
+import cugraph_pyg.nn