108 lines
4.8 KiB
Python
108 lines
4.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
from traceback import format_exc
|
||
|
||
from GPUtil import getAvailable, getGPUs
|
||
from loguru import logger
|
||
from torch.cuda import is_available
|
||
|
||
from enums.ExceptionEnum import ExceptionType
|
||
from exception.CustomerException import ServiceException
|
||
|
||
|
||
# order- 确定返回可用 GPU 设备 ID 的顺序。order应指定为以下字符串之一：
|
||
# 'first'- 按升序排列可用的 GPU 设备 ID（默认）
|
||
# 'last'- 按 id 降序排列可用的 GPU 设备 id
|
||
# 'random'- 随机订购可用的 GPU 设备 ID
|
||
# 'load'- 按负载递增排序可用的 GPU 设备 ID
|
||
# 'memory'- 通过升序内存使用来排序可用的 GPU 设备 ID
|
||
# limit- 将返回的 GPU 设备 ID 数量限制为指定数量。必须是正整数。（默认 = 1）
|
||
# maxLoad- 被认为可用的 GPU 的最大当前相对负载。负载大于 的 GPUmaxLoad不会返回。（默认 = 0.5）
|
||
# maxMemory- 被视为可用的 GPU 的最大当前相对内存使用量。maxMemory不返回当前内存使用量大于的 GPU 。（默认 = 0.5）
|
||
# includeNan- 真/假标志，指示是否包括负载或内存使用为 NaN 的 GPU（指示无法检索使用情况）。（默认 = 假）
|
||
# excludeID- ID 列表，应从可用 GPU 列表中排除。见GPU类描述。（默认 = []）
|
||
# excludeUUIDexcludeID-除了它使用 UUID 之外，其他相同。（默认 = []）
|
||
# 输出
|
||
# deviceIDs - 所有可用 GPU 设备 ID 的列表。如果当前负载和内存使用量分别小于maxLoad和maxMemory，则认为 GPU 可用。该列表是根据 排序的order。返回的设备 ID 的最大数量由 限制limit。
|
||
def get_gpu_ids():
|
||
deviceIDs = getAvailable(maxLoad=0.80, maxMemory=0.80)
|
||
return deviceIDs
|
||
|
||
|
||
def get_all_gpu_ids():
|
||
return getGPUs()
|
||
|
||
|
||
def get_first_gpu_name():
|
||
gps = get_all_gpu_ids()
|
||
if len(gps) == 0:
|
||
raise ServiceException(ExceptionType.NO_GPU_RESOURCES.value[0],
|
||
ExceptionType.NO_GPU_RESOURCES.value[1])
|
||
return gps[0].name
|
||
|
||
|
||
def check_gpu_resource(requestId=None):
|
||
gpu_ids = get_gpu_ids()
|
||
if len(gpu_ids) == 0 or 0 not in gpu_ids:
|
||
print_gpu_status(requestId)
|
||
raise ServiceException(ExceptionType.NO_RESOURCES.value[0],
|
||
ExceptionType.NO_RESOURCES.value[1])
|
||
return gpu_ids
|
||
|
||
|
||
def print_gpu_ex_status(requestId=None):
|
||
result = False
|
||
try:
|
||
gpu_ids = get_gpu_ids()
|
||
if len(gpu_ids) == 0 or 0 not in gpu_ids:
|
||
result = True
|
||
print_gpu_status(requestId)
|
||
except Exception:
|
||
logger.error("打印gpu状态异常: {}", format_exc())
|
||
return result
|
||
def select_best_server(servers):
|
||
best_server_info = None
|
||
lowest_memory_usage = float('inf') # 初始化为无穷大
|
||
|
||
for ip, server_info in servers.items():
|
||
gpu_list = server_info['GPU']
|
||
for gpu in gpu_list:
|
||
if gpu['ID'] == 0:
|
||
memory_used = gpu['Memory Used']
|
||
memory_total = gpu['Memory Total']
|
||
memory_usage = (memory_used / memory_total) * 100 # 计算显存使用率
|
||
|
||
if memory_usage < lowest_memory_usage:
|
||
lowest_memory_usage = memory_usage
|
||
best_server_info = {
|
||
'hostname': server_info['System']['Platform Node'],
|
||
'IP': server_info['System']['Local IP Address'],
|
||
'gpuId': gpu['ID']
|
||
}
|
||
|
||
return best_server_info
|
||
|
||
def print_gpu_status(requestId=None):
|
||
try:
|
||
GPUs = get_all_gpu_ids()
|
||
if len(GPUs) == 0:
|
||
return
|
||
for gpu in GPUs:
|
||
if requestId:
|
||
logger.info("""############################################################################################
|
||
GPU ID:{}, GPU 名称:{}, 负载率:{}, 内存使用率:{}, 总内存:{}, 占用内存:{}, 空闲内存:{}, requestId:{}
|
||
############################################################################################""", gpu.id,
|
||
gpu.name, gpu.load * 100, gpu.memoryUtil * 100, gpu.memoryTotal, gpu.memoryUsed, gpu.memoryFree,
|
||
requestId)
|
||
else:
|
||
logger.info("""############################################################################################
|
||
GPU ID:{}, GPU 名称:{}, 负载率:{}, 内存使用率:{}, 总内存:{}, 占用内存:{}, 空闲内存:{}
|
||
############################################################################################""", gpu.id,
|
||
gpu.name, gpu.load * 100, gpu.memoryUtil * 100, gpu.memoryTotal, gpu.memoryUsed, gpu.memoryFree)
|
||
except Exception:
|
||
logger.error("打印gpu状态异常: {}", format_exc())
|
||
|
||
|
||
def check_cude_is_available():
|
||
if not is_available():
|
||
raise Exception("cuda不在活动状态, 请检测显卡驱动是否正常！！！！")
|