make cugraph-ops optional for cugraph-pyg

rapidsai · Dec 19, 2024 · 4f4c82e · 4f4c82e
1 parent 77206de
commit 4f4c82e
Show file tree

Hide file tree

Showing 11 changed files with 51 additions and 81 deletions.
diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml
@@ -36,15 +36,13 @@ dependencies:
 - pre-commit
 - pydantic
 - pydata-sphinx-theme
-- pylibcugraphops==25.2.*,>=0.0.0a0
 - pylibraft==25.2.*,>=0.0.0a0
 - pytest
 - pytest-benchmark
 - pytest-cov
 - pytest-forked
 - pytest-xdist
-- pytorch-cuda=11.8
-- pytorch>=2.3
+- pytorch-gpu>=2.3=*cuda118*
 - pytorch_geometric>=2.5,<2.6
 - raft-dask==25.2.*,>=0.0.0a0
 - rapids-build-backend>=0.3.0,<0.4.0.dev0

diff --git a/conda/environments/all_cuda-121_arch-x86_64.yaml b/conda/environments/all_cuda-121_arch-x86_64.yaml
@@ -42,15 +42,13 @@ dependencies:
 - pre-commit
 - pydantic
 - pydata-sphinx-theme
-- pylibcugraphops==25.2.*,>=0.0.0a0
 - pylibraft==25.2.*,>=0.0.0a0
 - pytest
 - pytest-benchmark
 - pytest-cov
 - pytest-forked
 - pytest-xdist
-- pytorch-cuda=12.1
-- pytorch>=2.3
+- pytorch-gpu>=2.3=*cuda120*
 - pytorch_geometric>=2.5,<2.6
 - raft-dask==25.2.*,>=0.0.0a0
 - rapids-build-backend>=0.3.0,<0.4.0.dev0

diff --git a/conda/environments/all_cuda-124_arch-x86_64.yaml b/conda/environments/all_cuda-124_arch-x86_64.yaml
@@ -42,15 +42,13 @@ dependencies:
 - pre-commit
 - pydantic
 - pydata-sphinx-theme
-- pylibcugraphops==25.2.*,>=0.0.0a0
 - pylibraft==25.2.*,>=0.0.0a0
 - pytest
 - pytest-benchmark
 - pytest-cov
 - pytest-forked
 - pytest-xdist
-- pytorch-cuda=12.4
-- pytorch>=2.3
+- pytorch-gpu>=2.3=*cuda120*
 - pytorch_geometric>=2.5,<2.6
 - raft-dask==25.2.*,>=0.0.0a0
 - rapids-build-backend>=0.3.0,<0.4.0.dev0

diff --git a/dependencies.yaml b/dependencies.yaml
@@ -20,7 +20,6 @@ files:
       - depends_on_dask_cudf
       - depends_on_pylibraft
       - depends_on_raft_dask
-      - depends_on_pylibcugraphops
       - depends_on_cupy
       - depends_on_pytorch
       - depends_on_dgl
@@ -45,7 +44,6 @@ files:
       - cuda_version
       - docs
       - py_version
-      - depends_on_pylibcugraphops
   test_cpp:
     output: none
     includes:
@@ -116,7 +114,6 @@ files:
       table: project
     includes:
       - depends_on_cugraph
-      - depends_on_pylibcugraphops
       - python_run_cugraph_dgl
   py_test_cugraph_dgl:
     output: pyproject
@@ -142,7 +139,6 @@ files:
       table: project
     includes:
       - depends_on_cugraph
-      - depends_on_pylibcugraphops
       - depends_on_pyg
       - python_run_cugraph_pyg
   py_test_cugraph_pyg:
@@ -166,7 +162,6 @@ files:
     includes:
       - checks
       - depends_on_cugraph
-      - depends_on_pylibcugraphops
       - depends_on_dgl
       - depends_on_pytorch
       - cugraph_dgl_dev
@@ -180,7 +175,6 @@ files:
       - checks
       - depends_on_cugraph
       - depends_on_pyg
-      - depends_on_pylibcugraphops
       - depends_on_pytorch
       - cugraph_pyg_dev
       - test_python_common
@@ -415,7 +409,6 @@ dependencies:
     common:
       - output_types: [conda]
         packages:
-          - pytorch>=2.3
           - torchdata
           - pydantic
     specific:
@@ -440,18 +433,16 @@ dependencies:
               - *tensordict
           - {matrix: null, packages: [*pytorch_pip, *tensordict]}
       - output_types: [conda]
+        # PyTorch will stop publishing conda packages after 2.5.
+        # Consider switching to conda-forge::pytorch-gpu.
+        # Note that the CUDA version may differ from the official PyTorch wheels.
         matrices:
-          - matrix: {cuda: "12.1"}
-            packages:
-              - pytorch-cuda=12.1
-          - matrix: {cuda: "12.4"}
+          - matrix: {cuda: "12.*"}
             packages:
-              - pytorch-cuda=12.4
-          - matrix: {cuda: "11.8"}
+            - pytorch-gpu>=2.3=*cuda120*
+          - matrix: {cuda: "11.*"}
             packages:
-              - pytorch-cuda=11.8
-          # pytorch only supports certain CUDA versions... skip
-          # adding pytorch-cuda pinning if any other CUDA version is requested
+            - pytorch-gpu>=2.3=*cuda118*
           - matrix:
             packages:
 
@@ -676,31 +667,6 @@ dependencies:
               - pylibcugraph-cu11==25.2.*,>=0.0.0a0
           - {matrix: null, packages: [*pylibcugraph_unsuffixed]}
 
-  depends_on_pylibcugraphops:
-    common:
-      - output_types: conda
-        packages:
-          - &pylibcugraphops_unsuffixed pylibcugraphops==25.2.*,>=0.0.0a0
-      - output_types: requirements
-        packages:
-          # pip recognizes the index as a global option for the requirements.txt file
-          - --extra-index-url=https://pypi.nvidia.com
-          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
-    specific:
-      - output_types: [requirements, pyproject]
-        matrices:
-          - matrix:
-              cuda: "12.*"
-              cuda_suffixed: "true"
-            packages:
-              - pylibcugraphops-cu12==25.2.*,>=0.0.0a0
-          - matrix:
-              cuda: "11.*"
-              cuda_suffixed: "true"
-            packages:
-              - pylibcugraphops-cu11==25.2.*,>=0.0.0a0
-          - {matrix: null, packages: [*pylibcugraphops_unsuffixed]}
-
   depends_on_cupy:
     common:
       - output_types: conda

diff --git a/python/cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml b/python/cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml
@@ -12,13 +12,11 @@ dependencies:
 - dglteam/label/th23_cu118::dgl>=2.4.0.th23.cu*
 - pre-commit
 - pydantic
-- pylibcugraphops==25.2.*,>=0.0.0a0
 - pytest
 - pytest-benchmark
 - pytest-cov
 - pytest-xdist
-- pytorch-cuda=11.8
-- pytorch>=2.3
+- pytorch-gpu>=2.3=*cuda118*
 - tensordict>=0.1.2
 - torchdata
 name: cugraph_dgl_dev_cuda-118
diff --git a/python/cugraph-dgl/pyproject.toml b/python/cugraph-dgl/pyproject.toml
@@ -27,7 +27,6 @@ dependencies = [
     "cugraph==25.2.*,>=0.0.0a0",
     "numba>=0.57",
     "numpy>=1.23,<3.0a0",
-    "pylibcugraphops==25.2.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project.optional-dependencies]

diff --git a/python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml b/python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml
@@ -10,13 +10,11 @@ dependencies:
 - cugraph==25.2.*,>=0.0.0a0
 - pre-commit
 - pydantic
-- pylibcugraphops==25.2.*,>=0.0.0a0
 - pytest
 - pytest-benchmark
 - pytest-cov
 - pytest-xdist
-- pytorch-cuda=11.8
-- pytorch>=2.3
+- pytorch-gpu>=2.3=*cuda118*
 - pytorch_geometric>=2.5,<2.6
 - tensordict>=0.1.2
 - torchdata

diff --git a/python/cugraph-pyg/cugraph_pyg/nn/conv/__init__.py b/python/cugraph-pyg/cugraph_pyg/nn/conv/__init__.py
@@ -11,18 +11,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .gat_conv import GATConv
-from .gatv2_conv import GATv2Conv
-from .hetero_gat_conv import HeteroGATConv
-from .rgcn_conv import RGCNConv
-from .sage_conv import SAGEConv
-from .transformer_conv import TransformerConv
+import warnings
 
-__all__ = [
-    "GATConv",
-    "GATv2Conv",
-    "HeteroGATConv",
-    "RGCNConv",
-    "SAGEConv",
-    "TransformerConv",
-]
+HAVE_CUGRAPH_OPS = False
+try:
+    import pylibcugraphops
+    HAVE_CUGRAPH_OPS = True
+except ImportError:
+    pass
+except Exception as e:
+    warnings.warn(f"Unexpected error while importing pylibcugraphops: {e}")
+
+if HAVE_CUGRAPH_OPS:
+    from .gat_conv import GATConv
+    from .gatv2_conv import GATv2Conv
+    from .hetero_gat_conv import HeteroGATConv
+    from .rgcn_conv import RGCNConv
+    from .sage_conv import SAGEConv
+    from .transformer_conv import TransformerConv
+
+    __all__ = [
+        "GATConv",
+        "GATv2Conv",
+        "HeteroGATConv",
+        "RGCNConv",
+        "SAGEConv",
+        "TransformerConv",
+    ]
diff --git a/python/cugraph-pyg/cugraph_pyg/tests/conftest.py b/python/cugraph-pyg/cugraph_pyg/tests/conftest.py
@@ -43,6 +43,11 @@
     gpubenchmark = pytest_benchmark.plugin.benchmark
 
 
+def pytest_ignore_collect(collection_path, config):
+    """Return True to prevent considering this path for collection."""
+    if "nn" in collection_path.name:
+        return True
+
 @pytest.fixture(scope="module")
 def dask_client():
     dask_scheduler_file = os.environ.get("SCHEDULER_FILE")

diff --git a/python/cugraph-pyg/cugraph_pyg/tests/loader/test_dask_neighbor_loader.py b/python/cugraph-pyg/cugraph_pyg/tests/loader/test_dask_neighbor_loader.py
@@ -23,7 +23,6 @@
 from cugraph_pyg.loader import DaskNeighborLoader
 from cugraph_pyg.loader import BulkSampleLoader
 from cugraph_pyg.data import DaskGraphStore
-from cugraph_pyg.nn import SAGEConv as CuGraphSAGEConv
 
 from cugraph.gnn import FeatureStore
 from cugraph.utilities.utils import import_optional, MissingModule
@@ -403,15 +402,15 @@ def test_cugraph_loader_e2e_csc(framework: str):
     )
 
     if framework == "pyg":
-        convs = [
-            torch_geometric.nn.SAGEConv(256, 64, aggr="mean").cuda(),
-            torch_geometric.nn.SAGEConv(64, 1, aggr="mean").cuda(),
-        ]
+        SAGEConv = torch_geometric.nn.SAGEConv
     else:
-        convs = [
-            CuGraphSAGEConv(256, 64, aggr="mean").cuda(),
-            CuGraphSAGEConv(64, 1, aggr="mean").cuda(),
-        ]
+        pytest.skip("Skipping tests that requires cugraph-ops")
+        # SAGEConv = cugraph_pyg.nn.SAGEConv
+
+    convs = [
+        SAGEConv(256, 64, aggr="mean").cuda(),
+        SAGEConv(64, 1, aggr="mean").cuda(),
+    ]
 
     trim = trim_to_layer.TrimToLayer()
     relu = torch.nn.functional.relu

diff --git a/python/cugraph-pyg/pyproject.toml b/python/cugraph-pyg/pyproject.toml
@@ -34,7 +34,6 @@ dependencies = [
     "numba>=0.57",
     "numpy>=1.23,<3.0a0",
     "pandas",
-    "pylibcugraphops==25.2.*,>=0.0.0a0",
     "torch-geometric>=2.5,<2.6",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.