[BUG] Fix Incorrect Edge Index, Directory Selection in cuGraph-PyG Loader

ci/test_python.sh

@@ -200,15 +200,21 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
     # Install pytorch
     rapids-mamba-retry install \
       --force-reinstall \
-      --channel pyg \
       --channel pytorch \
       --channel nvidia \
-      'pyg=2.3' \
-      'pytorch=2.0.0' \
+      --channel conda-forge \
+      'pytorch=2.1.0' \
       '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
+    pip install \
+        torch_geometric==2.4.0
+        pyg_lib \
+        torch_scatter \
+        torch_sparse \
+        torch_cluster \
+        torch_spline_conv \
+      -f https://data.pyg.org/whl/torch-2.1.0+cu118.html
 
     rapids-mamba-retry install \
       --channel "${CPP_CHANNEL}" \

dependencies.yaml

@@ -503,9 +503,9 @@ dependencies:
       - output_types: [conda]
         packages:
           - cugraph==23.12.*
-          - pytorch==2.0
+          - pytorch==2.1
           - pytorch-cuda==11.8
-          - pyg=2.3.1=*torch_2.0.0*cu118*
+          - pyg=2.4.0=*torch_2.1.0*cu118*
 
   depends_on_rmm:
     common:

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

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

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
@@ -437,11 +455,10 @@ 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:
+                src = out[edge_type].edge_index[0]
+                dst = out[edge_type].edge_index[1]
+                out[edge_type].edge_index = torch.stack([dst, src])
 
         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/test_cugraph_loader.py

@@ -18,6 +18,7 @@
 
 import cudf
 import cupy
+import numpy as np
 
 from cugraph_pyg.loader import CuGraphNeighborLoader
 from cugraph_pyg.loader import BulkSampleLoader
@@ -27,6 +28,8 @@
 from cugraph.gnn import FeatureStore
 from cugraph.utilities.utils import import_optional, MissingModule
 
+from typing import Dict, Tuple
+
 torch = import_optional("torch")
 torch_geometric = import_optional("torch_geometric")
 trim_to_layer = import_optional("torch_geometric.utils.trim_to_layer")
@@ -40,7 +43,9 @@
 
 
 @pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
-def test_cugraph_loader_basic(karate_gnn):
+def test_cugraph_loader_basic(
+    karate_gnn: Tuple[FeatureStore, Dict[str, np.ndarray], Dict[str, int]]
+):
     F, G, N = karate_gnn
     cugraph_store = CuGraphStore(F, G, N, order="CSR")
     loader = CuGraphNeighborLoader(
@@ -66,7 +71,9 @@ def test_cugraph_loader_basic(karate_gnn):
 
 
 @pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
-def test_cugraph_loader_hetero(karate_gnn):
+def test_cugraph_loader_hetero(
+    karate_gnn: Tuple[FeatureStore, Dict[str, np.ndarray], Dict[str, int]]
+):
     F, G, N = karate_gnn
     cugraph_store = CuGraphStore(F, G, N, order="CSR")
     loader = CuGraphNeighborLoader(
@@ -342,7 +349,7 @@ def test_cugraph_loader_e2e_coo():
 @pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
 @pytest.mark.skipif(not HAS_TORCH_SPARSE, reason="torch-sparse not available")
 @pytest.mark.parametrize("framework", ["pyg", "cugraph-ops"])
-def test_cugraph_loader_e2e_csc(framework):
+def test_cugraph_loader_e2e_csc(framework: str):
     m = [2, 9, 99, 82, 9, 3, 18, 1, 12]
     x = torch.randint(3000, (256, 256)).to(torch.float32)
     F = FeatureStore()
@@ -442,3 +449,38 @@ def test_cugraph_loader_e2e_csc(framework):
         x = x.narrow(dim=0, start=0, length=s - num_sampled_nodes[1])
 
         assert list(x.shape) == [1, 1]
+
+
+@pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
+@pytest.mark.parametrize("directory", ["local", "temp"])
+def test_load_directory(
+    karate_gnn: Tuple[FeatureStore, Dict[str, np.ndarray], Dict[str, int]],
+    directory: str,
+):
+    if directory == "local":
+        local_dir = tempfile.TemporaryDirectory(dir=".")
+
+    cugraph_store = CuGraphStore(*karate_gnn)
+    cugraph_loader = CuGraphNeighborLoader(
+        (cugraph_store, cugraph_store),
+        torch.arange(8, dtype=torch.int64),
+        2,
+        num_neighbors=[8, 4, 2],
+        random_state=62,
+        replace=False,
+        directory=None if directory == "temp" else local_dir.name,
+        batches_per_partition=1,
+    )
+
+    it = iter(cugraph_loader)
+    next_batch = next(it)
+    assert next_batch is not None
+
+    if directory == "local":
+        assert len(os.listdir(local_dir.name)) == 4
+
+    count = 1
+    while next(it, None) is not None:
+        count += 1
+
+    assert count == 4