Fix sphinx complains by moving cupy sync to _nccl for protected import mech

Author: tharittk

pylops_mpi/utils/_nccl.py

@@ -1,6 +1,7 @@
 __all__ = [
     "_prepare_nccl_allgather_inputs",
     "_unroll_nccl_allgather_recv",
+    "_nccl_sync",
     "initialize_nccl_comm",
     "nccl_split",
     "nccl_allgather",
@@ -19,7 +20,6 @@
 import cupy as cp
 import cupy.cuda.nccl as nccl
 
-
 cupy_to_nccl_dtype = {
     "float32": nccl.NCCL_FLOAT32,
     "float64": nccl.NCCL_FLOAT64,
@@ -63,6 +63,13 @@ def _nccl_buf_size(buf, count=None):
         return count if count else buf.size
 
 
+def _nccl_sync():
+    """A thin wrapper of CuPy's synchronization for protected import"""
+    if cp.cuda.runtime.getDeviceCount() == 0:
+        return
+    cp.cuda.runtime.deviceSynchronize()
+
+
 def _prepare_nccl_allgather_inputs(send_buf, send_buf_shapes) -> Tuple[cp.ndarray, cp.ndarray]:
     r""" Prepare send_buf and recv_buf for NCCL allgather (nccl_allgather)
 

pylops_mpi/utils/benchmark.py

@@ -3,19 +3,18 @@
 import time
 from typing import Callable, Optional, List
 from mpi4py import MPI
-from pylops.utils import deps
-
-cupy_message = deps.cupy_import("benchmark module")
-if cupy_message is None:
-    import cupy as cp
-    if cp.cuda.runtime.getDeviceCount() == 0:
-        has_cupy = False
-        print(UserWarning("CuPy is installed, but no CUDA-capable device is available."))
-    else:
-        has_cupy = True
-else:
-    has_cupy = False
 
+from pylops.utils import deps as pylops_deps  # avoid namespace crashes with pylops_mpi.utils
+from pylops_mpi.utils import deps
+
+cupy_message = pylops_deps.cupy_import("the benchmark module")
+nccl_message = deps.nccl_import("the benchmark module")
+
+if nccl_message is None and cupy_message is None:
+    from pylops_mpi.utils._nccl import _nccl_sync
+else:
+    def _nccl_sync():
+        pass
 
 # TODO (tharitt): later move to env file or something
 ENABLE_BENCHMARK = True
@@ -65,10 +64,7 @@ def _parse_output_tree(markers: List[str]):
 
 def _sync():
     """Synchronize all MPI processes or CUDA Devices"""
-    if has_cupy:
-        # this is ok to call even if CUDA runtime is not initialized
-        cp.cuda.runtime.deviceSynchronize()
-
+    _nccl_sync()
     MPI.COMM_WORLD.Barrier()