from typing import List import cupy as cp class CuPyDeviceInfo: """ Represents picklable CuPy device info """ def __init__(self, index=None, name=None, total_memory=None): self._index : int = index self._name : str = name self._total_memory : int = total_memory def __getstate__(self): return self.__dict__.copy() def __setstate__(self, d): self.__init__() self.__dict__.update(d) def is_cpu(self) -> bool: return self._index == -1 def get_index(self) -> int: return self._index def get_name(self) -> str: return self._name def get_total_memory(self) -> int: return self._total_memory def __eq__(self, other): if self is not None and other is not None and isinstance(self, CuPyDeviceInfo) and isinstance(other, CuPyDeviceInfo): return self._index == other._index return False def __hash__(self): return self._index def __str__(self): if self.is_cpu(): return "CPU" else: return f"[{self._index}] {self._name} [{(self._total_memory / 1024**3) :.3}Gb]" def __repr__(self): return f'{self.__class__.__name__} object: ' + self.__str__() _cupy_devices = None def get_available_devices() -> List[CuPyDeviceInfo]: """ returns a list of available CuPyDeviceInfo """ global _cupy_devices if _cupy_devices is None: devices = [] for i in range (cp.cuda.runtime.getDeviceCount()): device_props = cp.cuda.runtime.getDeviceProperties(i) devices.append ( CuPyDeviceInfo(index=i, name=device_props['name'].decode('utf-8'), total_memory=device_props['totalGlobalMem'])) _cupy_devices = devices return _cupy_devices