Source code for cellvit.utils.check_cupy
# -*- coding: utf-8 -*-
# Check if CuPy is installed and working correctly.
#
# @ Fabian Hörst, fabian.hoerst@uk-essen.de
# Institute for Artifical Intelligence in Medicine,
# University Medicine Essen
from typing import Optional
import logging
from cellvit.utils.logger import NullLogger
[docs]
def check_cupy(
check_gpu: True, logger: Optional[logging.Logger] = NullLogger()
) -> bool:
"""Check if CuPy is installed and working correctly.
Args:
check_gpu (True): Check if CuPy can access the GPU.
logger (Optional[logging.Logger], optional): Logger. Defaults to None.
Raises:
SystemError: No CUDA-capable GPU detected by CuPy.
Returns:
bool: True if CuPy is installed and working correctly.
"""
try:
import cupy as cp
# Perform a simple CuPy operation to ensure it's functioning
a = cp.array([1, 2, 3])
b = cp.array([4, 5, 6])
c = a + b
if not cp.allclose(c, [5, 7, 9]):
logger.error("CuPy operation validation failed.")
return False
elif not check_gpu:
logger.info("CuPy is functioning correctly, but without GPU check")
return True
device_count = cp.cuda.runtime.getDeviceCount()
if device_count == 0:
raise SystemError("No CUDA-capable GPU detected by CuPy.")
logger.info(f"CuPy detected {device_count} CUDA device(s).")
for device_id in range(device_count):
# Set the current device
cp.cuda.Device(device_id).use()
# Get device properties
props = cp.cuda.runtime.getDeviceProperties(device_id)
logger.info(f"Device ID: {device_id}")
logger.info(f" Name: {props['name']}")
logger.info(f" Total Global Memory: {props['totalGlobalMem']} bytes")
logger.info(f" Multiprocessor Count: {props['multiProcessorCount']}")
logger.info(f" Compute Capability: {props['major']}.{props['minor']}")
logger.info("")
logger.info("CuPy is able to access the GPU and perform operations.")
return True
except ImportError as e:
logger.error(f"CuPy Import Error: {e}")
return False
except RuntimeError as e:
logger.error(f"CuPy Error: {e}")
return False
except SystemError as e:
logger.error(f"CuPy GPU Error: {e}")
return False
except Exception as e:
logger.error(f"Unexpected error during CuPy check: {e}")
return False