diff --git a/python/xoscar/nvutils.py b/python/xoscar/nvutils.py index 187fdceb..5da59fee 100644 --- a/python/xoscar/nvutils.py +++ b/python/xoscar/nvutils.py @@ -347,6 +347,10 @@ def get_device_count() -> int | None: return _gpu_count +def cuda_count() -> int: + return get_device_count() or 0 + + def _get_all_device_count() -> int | None: _init_nvml() if _nvml_lib is None: diff --git a/python/xoscar/utils.py b/python/xoscar/utils.py index 04e19ec0..d40488e1 100644 --- a/python/xoscar/utils.py +++ b/python/xoscar/utils.py @@ -238,6 +238,8 @@ def lazy_import( rename: str | None = None, placeholder: bool = False, ): + from .nvutils import cuda_count + rename = rename or name prefix_name = name.split(".", 1)[0] globals = globals or inspect.currentframe().f_back.f_globals # type: ignore @@ -267,6 +269,12 @@ def add_load_handler(self, func: Callable): return func if pkgutil.find_loader(prefix_name) is not None: + # There are cuda-related files in python env (commonly happen in a shared env), + # but that is a cpu only machine without GPU cards, just return None + if ( + "cupy" in name or "cudf" in name or "cuda" in name or "rmm" in name + ) and cuda_count() == 0: # pragma: no cover + return None return LazyModule() elif placeholder: return ModulePlaceholder(prefix_name)