Add RAPIDS Doctor Check for cuGraph-PyG and pylibwholegraph

python/cugraph-pyg/cugraph_pyg/_doctor_check.py

@@ -0,0 +1,68 @@
+# 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):
+        import warnings
+
+        warnings.warn(
+            "PyTorch is required to use cuGraph-PyG. "
+            "Please install PyTorch from PyPI or Conda-Forge."
+        )
+    else:
+        import os
+        from cugraph_pyg.data import GraphStore
+
+        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])
+
+        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,39 @@
+# 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"
+        )
+
+    from pylibwholegraph.utils import import_optional, MissingModule
+
+    torch = import_optional("torch")
+
+    if isinstance(torch, MissingModule):
+        import warnings
+
+        warnings.warn(
+            "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 = [