[FEA] New Graph Interface and Loaders for Distributed Sampling in DGL

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

@@ -46,6 +46,7 @@ dependencies:
 - packaging>=21
 - pandas
 - pre-commit
+- pydantic
 - pydata-sphinx-theme
 - pylibcugraphops==24.8.*,>=0.0.0a0
 - pylibraft==24.8.*,>=0.0.0a0
@@ -72,6 +73,7 @@ dependencies:
 - sphinx<6
 - sphinxcontrib-websupport
 - thriftpy2<=0.5.0
+- torchdata
 - ucx-proc=*=gpu
 - ucx-py==0.39.*,>=0.0.0a0
 - wget

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

@@ -51,6 +51,7 @@ dependencies:
 - packaging>=21
 - pandas
 - pre-commit
+- pydantic
 - pydata-sphinx-theme
 - pylibcugraphops==24.8.*,>=0.0.0a0
 - pylibraft==24.8.*,>=0.0.0a0
@@ -77,6 +78,7 @@ dependencies:
 - sphinx<6
 - sphinxcontrib-websupport
 - thriftpy2<=0.5.0
+- torchdata
 - ucx-proc=*=gpu
 - ucx-py==0.39.*,>=0.0.0a0
 - wget

conda/recipes/cugraph-dgl/meta.yaml

@@ -28,8 +28,10 @@ requirements:
     - numba >=0.57
     - numpy >=1.23,<2.0a0
     - pylibcugraphops ={{ minor_version }}
+    - tensordict >=0.1.2
     - python
-    - pytorch
+    - pytorch >=2.0
+    - cupy >=12.0.0
 
 tests:
   imports:

dependencies.yaml

@@ -647,6 +647,7 @@ dependencies:
           - cugraph==24.8.*,>=0.0.0a0
           - pytorch>=2.0
           - pytorch-cuda==11.8
+          - &tensordict tensordict>=0.1.2
           - dgl>=1.1.0.cu*
   cugraph_pyg_dev:
     common:
@@ -655,14 +656,16 @@ dependencies:
           - cugraph==24.8.*,>=0.0.0a0
           - pytorch>=2.0
           - pytorch-cuda==11.8
-          - &tensordict tensordict>=0.1.2
+          - *tensordict
           - pyg>=2.5,<2.6
 
   depends_on_pytorch:
     common:
       - output_types: [conda]
         packages:
           - &pytorch_conda pytorch>=2.0,<2.2.0a0
+          - torchdata
+          - pydantic
 
     specific:
       - output_types: [requirements]

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

@@ -21,4 +21,5 @@ dependencies:
 - pytorch-cuda==11.8
 - pytorch>=2.0
 - scipy
+- tensordict>=0.1.2
 name: cugraph_dgl_dev_cuda-118

python/cugraph-dgl/cugraph_dgl/__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
@@ -15,8 +15,12 @@
 
 # to prevent rapids context being created when importing cugraph_dgl
 os.environ["RAPIDS_NO_INITIALIZE"] = "1"
+from cugraph_dgl.graph import Graph
 from cugraph_dgl.cugraph_storage import CuGraphStorage
-from cugraph_dgl.convert import cugraph_storage_from_heterograph
+from cugraph_dgl.convert import (
+    cugraph_storage_from_heterograph,
+    cugraph_dgl_graph_from_heterograph,
+)
 import cugraph_dgl.dataloading
 import cugraph_dgl.nn
 

python/cugraph-dgl/cugraph_dgl/convert.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022-2023, NVIDIA CORPORATION.
+# Copyright (c) 2022-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,6 +12,8 @@
 # limitations under the License.
 from __future__ import annotations
 from cugraph.utilities.utils import import_optional
+
+import cugraph_dgl
 from cugraph_dgl import CuGraphStorage
 from cugraph_dgl.utils.cugraph_conversion_utils import (
     get_edges_dict_from_dgl_HeteroGraph,
@@ -39,3 +41,53 @@ def cugraph_storage_from_heterograph(
     add_ndata_from_dgl_HeteroGraph(gs, g)
     add_edata_from_dgl_HeteroGraph(gs, g)
     return gs
+
+
+def cugraph_dgl_graph_from_heterograph(
+    input_graph: dgl.DGLGraph,
+    single_gpu: bool = True,
+    ndata_storage: str = "torch",
+    edata_storage: str = "torch",
+    **kwargs,
+) -> cugraph_dgl.Graph:
+    """
+    Converts a DGL Graph to a cuGraph-DGL Graph.
+    """
+
+    output_graph = cugraph_dgl.Graph(
+        is_multi_gpu=(not single_gpu),
+        ndata_storage=ndata_storage,
+        edata_storage=edata_storage,
+        **kwargs,
+    )
+
+    # Calling is_homogeneous does not work here
+    if len(input_graph.ntypes) <= 1:
+        output_graph.add_nodes(
+            input_graph.num_nodes(), data=input_graph.ndata, ntype=input_graph.ntypes[0]
+        )
+    else:
+        for ntype in input_graph.ntypes:
+            data = {
+                k: v_dict[ntype]
+                for k, v_dict in input_graph.ndata.items()
+                if ntype in v_dict
+            }
+            output_graph.add_nodes(input_graph.num_nodes(ntype), data=data, ntype=ntype)
+
+    if len(input_graph.canonical_etypes) <= 1:
+        can_etype = input_graph.canonical_etypes[0]
+        src_t, dst_t = input_graph.edges(form="uv", etype=can_etype)
+        output_graph.add_edges(src_t, dst_t, input_graph.edata, etype=can_etype)
+    else:
+        for can_etype in input_graph.canonical_etypes:
+            data = {
+                k: v_dict[can_etype]
+                for k, v_dict in input_graph.edata.items()
+                if can_etype in v_dict
+            }
+
+            src_t, dst_t = input_graph.edges(form="uv", etype=can_etype)
+            output_graph.add_edges(src_t, dst_t, data=data, etype=can_etype)
+
+    return output_graph

python/cugraph-dgl/cugraph_dgl/dataloading/__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
@@ -11,9 +11,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 from cugraph_dgl.dataloading.dataset import (
     HomogenousBulkSamplerDataset,
     HeterogenousBulkSamplerDataset,
 )
+
+from cugraph_dgl.dataloading.sampler import Sampler
 from cugraph_dgl.dataloading.neighbor_sampler import NeighborSampler
-from cugraph_dgl.dataloading.dataloader import DataLoader
+
+from cugraph_dgl.dataloading.dask_dataloader import DaskDataLoader
+from cugraph_dgl.dataloading.dataloader import DataLoader as FutureDataLoader
+
+
+def DataLoader(*args, **kwargs):
+    warnings.warn(
+        "DataLoader has been renamed to DaskDataLoader.  "
+        "In Release 24.10, cugraph_dgl.dataloading.FutureDataLoader "
+        "will take over the DataLoader name.",
+        FutureWarning,
+    )
+    return DaskDataLoader(*args, **kwargs)