Merge remote-tracking branch 'upstream/branch-25.04' into nvks-runners

ci/build_wheel.sh

@@ -35,7 +35,7 @@ cd "${package_dir}"
 sccache --zero-stats
 
 rapids-logger "Building '${package_name}' wheel"
-python -m pip wheel \
+rapids-pip-retry wheel \
     -w dist \
     -v \
     --no-deps \

ci/test_wheel_cugraph-dgl.sh

@@ -24,7 +24,7 @@ PYTORCH_URL="https://download.pytorch.org/whl/cu${PYTORCH_CUDA_VER}"
 DGL_URL="https://data.dgl.ai/wheels/torch-2.3/cu${PYTORCH_CUDA_VER}/repo.html"
 
 # echo to expand wildcard before adding `[extra]` requires for pip
-python -m pip install \
+rapids-pip-retry install \
     -v \
     --extra-index-url "${PYTORCH_URL}" \
     --find-links "${DGL_URL}" \

ci/test_wheel_cugraph-pyg.sh

@@ -22,7 +22,7 @@ else
 fi
 
 # echo to expand wildcard before adding `[extra]` requires for pip
-python -m pip install \
+rapids-pip-retry install \
     -v \
     --extra-index-url "${PYTORCH_URL}" \
     --find-links "${PYG_URL}" \

ci/test_wheel_pylibwholegraph.sh

@@ -23,7 +23,7 @@ mkdir -p "${RAPIDS_TESTS_DIR}" "${RAPIDS_COVERAGE_DIR}"
 
 # echo to expand wildcard before adding `[extra]` requires for pip
 rapids-logger "Installing Packages"
-rapids-retry python -m pip install \
+rapids-pip-retry install \
     --extra-index-url ${INDEX_URL} \
     "$(echo ./dist/pylibwholegraph*.whl)[test]" \
     'torch>=2.3'

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

@@ -14,7 +14,7 @@ dependencies:
 - cudatoolkit
 - cudf==25.4.*,>=0.0.0a0
 - cugraph==25.4.*,>=0.0.0a0
-- cupy>=12.0.0
+- cupy>=13.2.0
 - cython>=3.0.0
 - dask-cudf==25.4.*,>=0.0.0a0
 - dglteam/label/th23_cu118::dgl>=2.4.0.th23.cu*

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

@@ -17,7 +17,7 @@ dependencies:
 - cuda-version=12.1
 - cudf==25.4.*,>=0.0.0a0
 - cugraph==25.4.*,>=0.0.0a0
-- cupy>=12.0.0
+- cupy>=13.2.0
 - cython>=3.0.0
 - dask-cudf==25.4.*,>=0.0.0a0
 - dglteam/label/th23_cu121::dgl>=2.4.0.th23.cu*

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

@@ -17,7 +17,7 @@ dependencies:
 - cuda-version=12.4
 - cudf==25.4.*,>=0.0.0a0
 - cugraph==25.4.*,>=0.0.0a0
-- cupy>=12.0.0
+- cupy>=13.2.0
 - cython>=3.0.0
 - dask-cudf==25.4.*,>=0.0.0a0
 - dglteam/label/th23_cu121::dgl>=2.4.0.th23.cu*

conda/recipes/cugraph-dgl/meta.yaml

@@ -32,7 +32,7 @@ requirements:
     - tensordict >=0.1.2
     - python
     - pytorch >=2.3
-    - cupy >=12.0.0
+    - cupy >=13.2.0
 
 tests:
   imports:

conda/recipes/cugraph-pyg/meta.yaml

@@ -33,7 +33,7 @@ requirements:
     - pandas
     - python
     - pytorch >=2.3
-    - cupy >=12.0.0
+    - cupy >=13.2.0
     - cugraph ={{ minor_version }}
     - tensordict >=0.1.2
     - pytorch_geometric >=2.5,<2.7

dependencies.yaml

@@ -644,7 +644,7 @@ dependencies:
     common:
       - output_types: conda
         packages:
-          - cupy>=12.0.0
+          - cupy>=13.2.0
     # NOTE: This is intentionally not broken into groups by a 'cuda_suffixed' selector like
     #       other packages with -cu{nn}x suffixes in this file.
     #       All RAPIDS wheel builds (including in devcontainers) expect cupy to be suffixed.
@@ -653,8 +653,8 @@ dependencies:
         matrices:
           - matrix: {cuda: "12.*"}
             packages:
-              - cupy-cuda12x>=12.0.0
+              - cupy-cuda12x>=13.2.0
           - matrix: {cuda: "11.*"}
             packages: &cupy_packages_cu11
-              - cupy-cuda11x>=12.0.0
+              - cupy-cuda11x>=13.2.0
           - {matrix: null, packages: *cupy_packages_cu11}

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

@@ -1,4 +1,4 @@
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Copyright (c) 2022-2025, 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
@@ -13,14 +13,26 @@
 
 import warnings
 
-from cugraph_pyg.data.dask_graph_store import DaskGraphStore
+from cugraph_pyg.data.dask_graph_store import (
+    DaskGraphStore as DEPRECATED__DaskGraphStore,
+)
 from cugraph_pyg.data.graph_store import GraphStore
 from cugraph_pyg.data.feature_store import (
     TensorDictFeatureStore,
     WholeFeatureStore,
 )
 
 
+def DaskGraphStore(*args, **kwargs):
+    warnings.warn(
+        "DaskGraphStore and the Dask API are deprecated."
+        " Please switch over to the new API (cugraph_pyg.data.GraphStore)",
+        FutureWarning,
+    )
+
+    return DEPRECATED__DaskGraphStore(*args, **kwargs)
+
+
 def CuGraphStore(*args, **kwargs):
     warnings.warn("CuGraphStore has been renamed to DaskGraphStore", FutureWarning)
     return DaskGraphStore(*args, **kwargs)

python/cugraph-pyg/cugraph_pyg/examples/graph_sage_mg.py

@@ -419,6 +419,12 @@ def parse_args():
 
 
 def main():
+    warnings.warn(
+        "The Dask API is used in this example is deprecated.  "
+        "Please refer to 'gcn_dist_mg' for an example that uses the new API.",
+        FutureWarning,
+    )
+
     args = parse_args()
     if args.dask_scheduler_file is None:
         warnings.warn(

python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py

@@ -15,6 +15,7 @@
 import time
 import argparse
 import gc
+import warnings
 
 import torch
 
@@ -201,6 +202,12 @@ def parse_args():
 
 
 def main():
+    warnings.warn(
+        "The Dask API is used in this example is deprecated.  "
+        "Please refer to 'gcn_dist_sg' for an example that uses the new API.",
+        FutureWarning,
+    )
+
     args = parse_args()
 
     try:

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

@@ -1,4 +1,4 @@
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Copyright (c) 2022-2025, 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
@@ -19,9 +19,35 @@
 from cugraph_pyg.loader.link_loader import LinkLoader
 from cugraph_pyg.loader.link_neighbor_loader import LinkNeighborLoader
 
-from cugraph_pyg.loader.dask_node_loader import DaskNeighborLoader
+from cugraph_pyg.loader.dask_node_loader import (
+    DaskNeighborLoader as DEPRECATED__DaskNeighborLoader,
+)
 
-from cugraph_pyg.loader.dask_node_loader import BulkSampleLoader
+from cugraph_pyg.loader.dask_node_loader import (
+    BulkSampleLoader as DEPRECATED__BulkSampleLoader,
+)
+
+
+def DaskNeighborLoader(*args, **kwargs):
+    warnings.warn(
+        "DaskNeighborLoader and the Dask API are deprecated."
+        "Consider switching to the new API"
+        " (cugraph_pyg.loader.node_loader, cugraph_pyg.loader.link_loader).",
+        FutureWarning,
+    )
+
+    return DEPRECATED__DaskNeighborLoader(*args, **kwargs)
+
+
+def BulkSampleLoader(*args, **kwargs):
+    warnings.warn(
+        "BulkSampleLoader and the Dask API are deprecated.  Unbuffered sampling"
+        " in cugraph_pyg will soon be completely deprecated and removed, including"
+        " in the new API.",
+        FutureWarning,
+    )
+
+    return DEPRECATED__BulkSampleLoader(*args, **kwargs)
 
 
 def CuGraphNeighborLoader(*args, **kwargs):

python/pylibwholegraph/examples/node_classfication.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2024, NVIDIA CORPORATION.
+# Copyright (c) 2019-2025, 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
@@ -18,6 +18,7 @@
 
 import apex
 import torch
+import pickle
 from apex.parallel import DistributedDataParallel as DDP
 import pylibwholegraph.torch as wgth
 
@@ -141,8 +142,12 @@ def main_func():
     if args.use_cpp_ext:
         wgth.compile_cpp_extension()
 
-    train_ds, valid_ds, test_ds = wgth.create_node_claffication_datasets(
-        args.pickle_data_path
+    data_and_label = dict()
+    with open(args.pickle_data_path, "rb") as f:
+        data_and_label = pickle.load(f)
+
+    train_ds, valid_ds, test_ds = wgth.create_node_classification_datasets(
+        data_and_label
     )
 
     graph_structure = wgth.GraphStructure()

python/pylibwholegraph/pylibwholegraph/torch/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2024, NVIDIA CORPORATION.
+# Copyright (c) 2019-2025, 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
@@ -76,7 +76,7 @@
 
 from .gnn_model import set_framework, create_gnn_layers, create_sub_graph, HomoGNNModel
 from .data_loader import (
-    create_node_claffication_datasets,
+    create_node_classification_datasets,
     get_train_dataloader,
     get_valid_test_dataloader,
 )

python/pylibwholegraph/pylibwholegraph/torch/data_loader.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2024, NVIDIA CORPORATION.
+# Copyright (c) 2019-2025, 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
@@ -13,7 +13,6 @@
 
 import numpy as np
 import torch
-import pickle
 from torch.utils.data import Dataset
 
 
@@ -28,10 +27,7 @@ def __len__(self):
         return len(self.dataset)
 
 
-def create_node_claffication_datasets(pickle_data_filename: str):
-    with open(pickle_data_filename, "rb") as f:
-        data_and_label = pickle.load(f)
-
+def create_node_classification_datasets(data_and_label: dict):
     train_data = {
         "idx": data_and_label["train_idx"],
         "label": data_and_label["train_label"],