Add RAPIDS Doctor Check for cuGraph-PyG and pylibwholegraph

python/cugraph-pyg/cugraph_pyg/_doctor_check.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Smoke check for `rapids doctor` (RAPIDS CLI).
+
+See: https://github.com/rapidsai/rapids-cli#check-plugins
+"""
+
+
+def cugraph_pyg_smoke_check(**kwargs):
+    """
+    A quick check to ensure cugraph-pyg can be imported and its core
+    submodules are loadable.
+    """
+    try:
+        import cugraph_pyg
+
+        # Ensure core submodules load (touches pylibwholegraph, torch-geometric, etc.)
+        import cugraph_pyg.data
+        import cugraph_pyg.tensor
+
+    except ImportError as e:
+        raise ImportError(
+            "cugraph-pyg or its dependencies could not be imported. "
+            "Tip: install with `pip install cugraph-pyg` or use a RAPIDS conda environment."
+        ) from e
+
+    if not hasattr(cugraph_pyg, "__version__") or not cugraph_pyg.__version__:
+        raise AssertionError(
+            "cugraph-pyg smoke check failed: __version__ not found or empty"
+        )
+
+    from cugraph_pyg.utils.imports import import_optional, MissingModule
+
+    torch = import_optional("torch")
+
+    if isinstance(torch, MissingModule) or not torch.cuda.is_available():
+        import warnings
+
+        warnings.warn(
+            "PyTorch with CUDA support is required to use cuGraph-PyG. "
+            "Please install PyTorch from PyPI or Conda-Forge."
+        )
+    else:
+        import os
+        from cugraph_pyg.data import GraphStore
+
+        addr = os.environ.get("MASTER_ADDR", "")
+        port = os.environ.get("MASTER_PORT", "")
+        local_rank = os.environ.get("LOCAL_RANK", "")
+        world_size = os.environ.get("WORLD_SIZE", "")
+        local_world_size = os.environ.get("LOCAL_WORLD_SIZE", "")
+        rank = os.environ.get("RANK", "")
+
+        try:
+            os.environ["MASTER_ADDR"] = "localhost"
+            os.environ["MASTER_PORT"] = "29505"
+            os.environ["LOCAL_RANK"] = "0"
+            os.environ["WORLD_SIZE"] = "1"
+            os.environ["LOCAL_WORLD_SIZE"] = "1"
+            os.environ["RANK"] = "0"
+            torch.distributed.init_process_group("nccl")
+
+            graph_store = GraphStore()
+            graph_store.put_edge_index(
+                torch.tensor([[0, 1], [1, 2]]),
+                ("person", "knows", "person"),
+                "coo",
+                False,
+                (3, 3),
+            )
+            edge_index = graph_store.get_edge_index(
+                ("person", "knows", "person"), "coo"
+            )
+            assert edge_index.shape == torch.Size([2, 2])
+        finally:
+            os.environ["MASTER_ADDR"] = addr
+            os.environ["MASTER_PORT"] = port
+            os.environ["LOCAL_RANK"] = local_rank
+            os.environ["WORLD_SIZE"] = world_size
+            os.environ["LOCAL_WORLD_SIZE"] = local_world_size
+            os.environ["RANK"] = rank
+            torch.distributed.destroy_process_group()

python/cugraph-pyg/pyproject.toml

@@ -41,6 +41,9 @@ dependencies = [
     "torch-geometric>=2.5,<2.8",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
+[project.entry-points."rapids_doctor_check"]
+cugraph_pyg_smoke_check = "cugraph_pyg._doctor_check:cugraph_pyg_smoke_check"
+
 [project.urls]
 Homepage = "https://github.com/rapidsai/cugraph-gnn"
 Documentation = "https://docs.rapids.ai/api/cugraph/stable/"

python/pylibwholegraph/pylibwholegraph/_doctor_check.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Smoke check for `rapids doctor` (RAPIDS CLI).
+
+See: https://github.com/rapidsai/rapids-cli#check-plugins
+"""
+
+
+def pylibwholegraph_smoke_check(**kwargs):
+    """
+    A quick check to ensure pylibwholegraph can be imported and the
+    native library loads correctly.
+    """
+    try:
+        import pylibwholegraph
+    except ImportError as e:
+        raise ImportError(
+            "pylibwholegraph or its dependencies could not be imported. "
+            "Tip: install with `pip install pylibwholegraph` or use a RAPIDS conda environment."
+        ) from e
+
+    if not hasattr(pylibwholegraph, "__version__") or not pylibwholegraph.__version__:
+        raise AssertionError(
+            "pylibwholegraph smoke check failed: __version__ not found or empty"
+        )
+
+    try:
+        import torch
+
+        assert torch.cuda.is_available()
+
+    except ImportError:
+        import warnings
+
+        warnings.warn(
+            "PyTorch with CUDA or its dependencies could not be imported. "
+            "PyTorch is required to use pylibwholegraph. "
+            "Please install PyTorch from PyPI or Conda-Forge."
+        )

python/pylibwholegraph/pyproject.toml

@@ -30,6 +30,8 @@ dependencies = [
     "numpy>=1.23,<3.0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
+[project.entry-points."rapids_doctor_check"]
+pylibwholegraph_smoke_check = "pylibwholegraph._doctor_check:pylibwholegraph_smoke_check"
 
 [project.optional-dependencies]
 test = [