File size: 9,680 Bytes
82ea528 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 |
import torch
import pynvml
import comfy.model_management
from ..core import logger
# from ctypes import *
# from pyrsmi import rocml
class CGPUInfo:
"""
This class is responsible for getting information from GPU (ONLY).
"""
cuda = False
pynvmlLoaded = False
# pyamdLoaded = False
# anygpuLoaded = False
cudaAvailable = False
torchDevice = 'cpu'
cudaDevice = 'cpu'
cudaDevicesFound = 0
switchGPU = True
switchVRAM = True
switchTemperature = True
gpus = []
gpusUtilization = []
gpusVRAM = []
gpusTemperature = []
def __init__(self):
try:
pynvml.nvmlInit()
self.pynvmlLoaded = True
logger.info('Pynvml (Nvidia) initialized.')
except Exception as e:
logger.error('Could not init pynvml (Nvidia).' + str(e))
# if not self.pynvmlLoaded:
# try:
# rocml.smi_initialize()
# self.pyamdLoaded = True
# logger.info('Pyrsmi (AMD) initialized.')
# except Exception as e:
# logger.error('Could not init pyrsmi (AMD).' + str(e))
# self.anygpuLoaded = self.pynvmlLoaded or self.pyamdLoaded
self.anygpuLoaded = self.pynvmlLoaded
try:
self.torchDevice = comfy.model_management.get_torch_device_name(comfy.model_management.get_torch_device())
except Exception as e:
logger.error('Could not pick default device.' + str(e))
# ZLUDA Check, self.torchDevice has 'ZLUDA' in it.
if 'zluda' in self.torchDevice or 'ZLUDA' in self.torchDevice or 'Zluda' in self.torchDevice:
logger.warn('ZLUDA detected. GPU monitoring will be disabled.')
self.anygpuLoaded = False
# self.pyamdLoaded = False
self.pynvmlLoaded = False
if self.anygpuLoaded and self.deviceGetCount() > 0:
self.cudaDevicesFound = self.deviceGetCount()
logger.info(f"GPU/s:")
# for simulate multiple GPUs (for testing) interchange these comments:
# for deviceIndex in range(3):
# deviceHandle = pynvml.nvmlDeviceGetHandleByIndex(0)
for deviceIndex in range(self.cudaDevicesFound):
deviceHandle = self.deviceGetHandleByIndex(deviceIndex)
gpuName = self.deviceGetName(deviceHandle, deviceIndex)
logger.info(f"{deviceIndex}) {gpuName}")
self.gpus.append({
'index': deviceIndex,
'name': gpuName,
})
# same index as gpus, with default values
self.gpusUtilization.append(True)
self.gpusVRAM.append(True)
self.gpusTemperature.append(True)
self.cuda = True
logger.info(self.systemGetDriverVersion())
else:
logger.warn('No GPU with CUDA detected.')
self.cudaDevice = 'cpu' if self.torchDevice == 'cpu' else 'cuda'
self.cudaAvailable = torch.cuda.is_available()
if self.cuda and self.cudaAvailable and self.torchDevice == 'cpu':
logger.warn('CUDA is available, but torch is using CPU.')
def getInfo(self):
logger.debug('Getting GPUs info...')
return self.gpus
def getStatus(self):
# logger.debug('CGPUInfo getStatus')
gpuUtilization = -1
gpuTemperature = -1
vramUsed = -1
vramTotal = -1
vramPercent = -1
gpuType = ''
gpus = []
if self.cudaDevice == 'cpu':
gpuType = 'cpu'
gpus.append({
'gpu_utilization': -1,
'gpu_temperature': -1,
'vram_total': -1,
'vram_used': -1,
'vram_used_percent': -1,
})
else:
gpuType = self.cudaDevice
if self.anygpuLoaded and self.cuda and self.cudaAvailable:
# for simulate multiple GPUs (for testing) interchange these comments:
# for deviceIndex in range(3):
# deviceHandle = self.deviceGetHandleByIndex(0)
for deviceIndex in range(self.cudaDevicesFound):
deviceHandle = self.deviceGetHandleByIndex(deviceIndex)
gpuUtilization = -1
vramPercent = -1
vramUsed = -1
vramTotal = -1
gpuTemperature = -1
# GPU Utilization
if self.switchGPU and self.gpusUtilization[deviceIndex]:
try:
gpuUtilization = self.deviceGetUtilizationRates(deviceHandle)
except Exception as e:
if str(e) == "Unknown Error":
logger.error('For some reason, pynvml is not working in a laptop with only battery, try to connect and turn on the monitor')
else:
logger.error('Could not get GPU utilization.' + str(e))
logger.error('Monitor of GPU is turning off (not on UI!)')
self.switchGPU = False
# VRAM
if self.switchVRAM and self.gpusVRAM[deviceIndex]:
# Torch or pynvml?, pynvml is more accurate with the system, torch is more accurate with comfyUI
memory = self.deviceGetMemoryInfo(deviceHandle)
vramUsed = memory['used']
vramTotal = memory['total']
# device = torch.device(gpuType)
# vramUsed = torch.cuda.memory_allocated(device)
# vramTotal = torch.cuda.get_device_properties(device).total_memory
# check if vramTotal is not zero or None
if vramTotal and vramTotal != 0:
vramPercent = vramUsed / vramTotal * 100
# Temperature
if self.switchTemperature and self.gpusTemperature[deviceIndex]:
try:
gpuTemperature = self.deviceGetTemperature(deviceHandle)
except Exception as e:
logger.error('Could not get GPU temperature. Turning off this feature. ' + str(e))
self.switchTemperature = False
gpus.append({
'gpu_utilization': gpuUtilization,
'gpu_temperature': gpuTemperature,
'vram_total': vramTotal,
'vram_used': vramUsed,
'vram_used_percent': vramPercent,
})
return {
'device_type': gpuType,
'gpus': gpus,
}
def deviceGetCount(self):
if self.pynvmlLoaded:
return pynvml.nvmlDeviceGetCount()
# elif self.pyamdLoaded:
# return rocml.smi_get_device_count()
else:
return 0
def deviceGetHandleByIndex(self, index):
if self.pynvmlLoaded:
return pynvml.nvmlDeviceGetHandleByIndex(index)
# elif self.pyamdLoaded:
# return index
else:
return 0
def deviceGetName(self, deviceHandle, deviceIndex):
if self.pynvmlLoaded:
gpuName = 'Unknown GPU'
try:
gpuName = pynvml.nvmlDeviceGetName(deviceHandle)
try:
gpuName = gpuName.decode('utf-8', errors='ignore')
except AttributeError as e:
pass
except UnicodeDecodeError as e:
gpuName = 'Unknown GPU (decoding error)'
print(f"UnicodeDecodeError: {e}")
return gpuName
# elif self.pyamdLoaded:
# return rocml.smi_get_device_name(deviceIndex)
else:
return ''
def systemGetDriverVersion(self):
if self.pynvmlLoaded:
return f'NVIDIA Driver: {pynvml.nvmlSystemGetDriverVersion()}'
# elif self.pyamdLoaded:
# ver_str = create_string_buffer(256)
# rocml.rocm_lib.rsmi_version_str_get(0, ver_str, 256)
# return f'AMD Driver: {ver_str.value.decode()}'
else:
return 'Driver unknown'
def deviceGetUtilizationRates(self, deviceHandle):
if self.pynvmlLoaded:
return pynvml.nvmlDeviceGetUtilizationRates(deviceHandle).gpu
# elif self.pyamdLoaded:
# return rocml.smi_get_device_utilization(deviceHandle)
else:
return 0
def deviceGetMemoryInfo(self, deviceHandle):
if self.pynvmlLoaded:
mem = pynvml.nvmlDeviceGetMemoryInfo(deviceHandle)
return {'total': mem.total, 'used': mem.used}
# elif self.pyamdLoaded:
# mem_used = rocml.smi_get_device_memory_used(deviceHandle)
# mem_total = rocml.smi_get_device_memory_total(deviceHandle)
# return {'total': mem_total, 'used': mem_used}
else:
return {'total': 1, 'used': 1}
def deviceGetTemperature(self, deviceHandle):
if self.pynvmlLoaded:
return pynvml.nvmlDeviceGetTemperature(deviceHandle, pynvml.NVML_TEMPERATURE_GPU)
# elif self.pyamdLoaded:
# temp = c_int64(0)
# rocml.rocm_lib.rsmi_dev_temp_metric_get(deviceHandle, 1, 0, byref(temp))
# return temp.value / 1000
else:
return 0
|