Deprecate Distributed Sampling API

python/cugraph/cugraph/gnn/data_loading/__init__.py

@@ -1,9 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
+import warnings
+
 from cugraph.gnn.data_loading.dist_sampler import (
-    NeighborSampler,
-    DistSampler,
+    DEPRECATED__NeighborSampler,
+    DEPRECATED__DistSampler,
 )
 from cugraph.gnn.data_loading.dist_io import (
     DistSampleWriter,
@@ -12,6 +14,22 @@
 )
 
 
+def DistSampler(*args, **kwargs):
+    warnings.warn(
+        "DistSampler is deprecated and will be removed in a future release.  Please migrate to the distributed sampling API in cuGraph-PyG.",
+        FutureWarning,
+    )
+    return DEPRECATED__DistSampler(*args, **kwargs)
+
+
+def NeighborSampler(*args, **kwargs):
+    warnings.warn(
+        "NeighborSampler is deprecated and will be removed in a future release.  Please migrate to the distributed sampling API in cuGraph-PyG.",
+        FutureWarning,
+    )
+    return DEPRECATED__NeighborSampler(*args, **kwargs)
+
+
 def UniformNeighborSampler(*args, **kwargs):
     return NeighborSampler(
         *args,

python/cugraph/cugraph/gnn/data_loading/dist_io/__init__.py

@@ -1,6 +1,31 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
+import warnings
 
-from .reader import BufferedSampleReader, DistSampleReader
-from .writer import DistSampleWriter
+from .reader import DEPRECATED__BufferedSampleReader, DEPRECATED__DistSampleReader
+from .writer import DEPRECATED__DistSampleWriter
+
+
+def BufferedSampleReader(*args, **kwargs):
+    warnings.warn(
+        "BufferedSampleReader is deprecated and will be removed in a future release.  Please migrate to the distributed sampling API in cuGraph-PyG.",
+        FutureWarning,
+    )
+    return DEPRECATED__BufferedSampleReader(*args, **kwargs)
+
+
+def DistSampleReader(*args, **kwargs):
+    warnings.warn(
+        "DistSampleReader is deprecated and will be removed in a future release.  Please migrate to the distributed sampling API in cuGraph-PyG.",
+        FutureWarning,
+    )
+    return DEPRECATED__DistSampleReader(*args, **kwargs)
+
+
+def DistSampleWriter(*args, **kwargs):
+    warnings.warn(
+        "DistSampleWriter is deprecated and will be removed in a future release.  Please migrate to the distributed sampling API in cuGraph-PyG.",
+        FutureWarning,
+    )
+    return DEPRECATED__DistSampleWriter(*args, **kwargs)

python/cugraph/cugraph/gnn/data_loading/dist_io/reader.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 
@@ -15,7 +15,7 @@
 torch = MissingModule("torch")
 
 
-class DistSampleReader:
+class DEPRECATED__DistSampleReader:
     def __init__(
         self,
         directory: str,
@@ -89,7 +89,7 @@ def __next__(self) -> Tuple[Dict[str, "torch.Tensor"], int, int]:
         raise StopIteration
 
 
-class BufferedSampleReader:
+class DEPRECATED__BufferedSampleReader:
     def __init__(
         self,
         nodes_call_groups: list["torch.Tensor"],

python/cugraph/cugraph/gnn/data_loading/dist_io/writer.py

@@ -9,7 +9,9 @@
 import cupy
 
 from cugraph.utilities.utils import MissingModule
-from cugraph.gnn.data_loading.dist_io import DistSampleReader
+from cugraph.gnn.data_loading.dist_io import (
+    DEPRECATED__DistSampleReader as DistSampleReader,
+)
 
 from cugraph.gnn.data_loading.bulk_sampler_io import create_df_from_disjoint_arrays
 
@@ -18,7 +20,7 @@
 torch = MissingModule("torch")
 
 
-class DistSampleWriter:
+class DEPRECATED__DistSampleWriter:
     def __init__(
         self,
         directory: str,

python/cugraph/cugraph/gnn/data_loading/dist_sampler.py

@@ -44,7 +44,7 @@ def verify_metadata(metadata: Optional[Dict[str, Union[str, Tuple[str, str, str]
                 )
 
 
-class DistSampler:
+class DEPRECATED__DistSampler:
     def __init__(
         self,
         graph: Union[pylibcugraph.SGGraph, pylibcugraph.MGGraph],
@@ -669,7 +669,7 @@ def _retain_original_seeds(self):
         return self.__retain_original_seeds
 
 
-class NeighborSampler(DistSampler):
+class DEPRECATED__NeighborSampler(DEPRECATED__DistSampler):
     # Number of vertices in the output minibatch, based
     # on benchmarking.
     BASE_VERTICES_PER_BYTE = 0.1107662486009992
@@ -760,7 +760,7 @@ def __calc_local_seeds_per_call(
 
         if local_seeds_per_call is None:
             if len([x for x in fanout if x <= 0]) > 0:
-                return NeighborSampler.UNKNOWN_VERTICES_DEFAULT
+                return DEPRECATED__NeighborSampler.UNKNOWN_VERTICES_DEFAULT
 
             if heterogeneous:
                 if len(fanout) % num_edge_types != 0:
@@ -774,7 +774,9 @@ def __calc_local_seeds_per_call(
             total_memory = torch.cuda.get_device_properties(0).total_memory
             fanout_prod = reduce(lambda x, y: x * y, fanout)
             return int(
-                NeighborSampler.BASE_VERTICES_PER_BYTE * total_memory / fanout_prod
+                DEPRECATED__NeighborSampler.BASE_VERTICES_PER_BYTE
+                * total_memory
+                / fanout_prod
             )
 
         return local_seeds_per_call

python/cugraph/pytest.ini

@@ -58,6 +58,11 @@ filterwarnings =
           ignore:The with_edge_properties flag is deprecated and will be removed:FutureWarning
           ignore:This function is deprecated. Batched support for multiple vertices:DeprecationWarning
           ignore:uniform_neighbor_sample is deprecated in favor of:FutureWarning
+          ignore:DistSampler is deprecated and will be removed in a future release:FutureWarning
+          ignore:NeighborSampler is deprecated and will be removed in a future release:FutureWarning
+          ignore:BufferedSampleReader is deprecated and will be removed in a future release:FutureWarning
+          ignore:DistSamplerReader is deprecated and will be removed in a future release:FutureWarning
+          ignore:DistSampleWriter is deprecated and will be removed in a future release:FutureWarning
           # Called via dask. Not obviously addressable in cugraph.
           ignore:The behavior of array concatenation with empty entries is deprecated:FutureWarning
           ignore:This method is deprecated and will no longer be supported. The symmetrization:FutureWarning