Skip to content

Fix how NCCL existence is checked #135

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 2, 2025
Merged

Fix how NCCL existence is checked #135

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
281be46
change from checking cupy.cuda.nccl import to nccl.available
May 31, 2025
e890875
fix flake8 lint
tharittk May 31, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 11 additions & 12 deletions pylops_mpi/utils/deps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,30 +3,29 @@
]

import os
from importlib import import_module, util
from importlib import util
from typing import Optional


# error message at import of available package
def nccl_import(message: Optional[str] = None) -> str:
nccl_test = (
# detect if nccl is available and the user is expecting it to be used
# CuPy must be checked first otherwise util.find_spec assumes it presents and check nccl immediately and lead to crash
util.find_spec("cupy") is not None and util.find_spec("cupy.cuda.nccl") is not None and int(os.getenv("NCCL_PYLOPS_MPI", 1)) == 1
# cupy.cuda.nccl comes with cupy installation so check the cupy
util.find_spec("cupy") is not None and int(os.getenv("NCCL_PYLOPS_MPI", 1)) == 1
)
if nccl_test:
# try importing it
try:
import_module("cupy.cuda.nccl") # noqa: F401

# if succesful, set the message to None
# if cupy is present, this import will not throw error. The NCCL existence is checked with nccl.avaiable
import cupy.cuda.nccl as nccl
if nccl.available:
# if succesfull, set the message to None
nccl_message = None
# if unable to import but the package is installed
except (ImportError, ModuleNotFoundError) as e:
else:
# if unable to import but the package is installed
nccl_message = (
f"Fail to import cupy.cuda.nccl, Falling back to pure MPI (error: {e})."
"cupy is installed but cupy.cuda.nccl not available, Falling back to pure MPI."
"Please ensure your CUDA NCCL environment is set up correctly "
"for more detials visit 'https://docs.cupy.dev/en/stable/install.html'"
"for more details visit 'https://docs.cupy.dev/en/stable/install.html'"
)
print(UserWarning(nccl_message))
else:
Expand Down