diff --git a/python/cugraph-pyg/cugraph_pyg/_doctor_check.py b/python/cugraph-pyg/cugraph_pyg/_doctor_check.py new file mode 100644 index 00000000..72642dff --- /dev/null +++ b/python/cugraph-pyg/cugraph_pyg/_doctor_check.py @@ -0,0 +1,90 @@ +# 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 + + env_vars = [ + "MASTER_ADDR", + "MASTER_PORT", + "LOCAL_RANK", + "WORLD_SIZE", + "LOCAL_WORLD_SIZE", + "RANK", + ] + env_values = {var: os.environ.get(var, None) for var in env_vars} + + initialized = False + 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") + initialized = True + + 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: + for key, value in env_values.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + if initialized: + torch.distributed.destroy_process_group() diff --git a/python/cugraph-pyg/pyproject.toml b/python/cugraph-pyg/pyproject.toml index 1013f7af..5b102513 100644 --- a/python/cugraph-pyg/pyproject.toml +++ b/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/" diff --git a/python/pylibwholegraph/pylibwholegraph/_doctor_check.py b/python/pylibwholegraph/pylibwholegraph/_doctor_check.py new file mode 100644 index 00000000..c6a653f1 --- /dev/null +++ b/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, AssertionError): + 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." + ) diff --git a/python/pylibwholegraph/pyproject.toml b/python/pylibwholegraph/pyproject.toml index dd4dda49..78bd8dbc 100644 --- a/python/pylibwholegraph/pyproject.toml +++ b/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 = [