多GPU监测

相信大家在跑实验时都希望让GPU二十四小时跑，但有时候实验在半夜才结束，为了避免晚上接着跑实验需要半夜起床，同时为了不浪费计算资源，我们可以对多个GPU进行实时监测，当监测到GPU空闲时可以接着跑其他实验。

import os
import sys
import time
 
cmd0 = 'CUDA_VISIBLE_DEVICES=0 nohup bash run.sh --stage 9'  #当GPU空闲时需要跑的脚本
cmd1 = 'CUDA_VISIBLE_DEVICES=1 nohup bash run.sh --stage 9'  #当GPU空闲时需要跑的脚本 
 
def gpu_info():
    gpu_status = os.popen('nvidia-smi | grep %').read().split('|') #根据nvidia-smi命令的返回值按照'|'为分隔符建立一个列表
    '''
    结果如：
    ['', ' N/A   64C    P0    68W /  70W ', '   9959MiB / 15079MiB ', '     79%      Default ', 
    '\n', ' N/A   73C    P0   108W /  70W ', '  11055MiB / 15079MiB ', '     63%      Default ', 
    '\n', ' N/A   60C    P0    55W /  70W ', '   3243MiB / 15079MiB ', '     63%      Default ', '\n']
    '''
    gpu0_status = gpu_status[0:4]
    gpu1_status = gpu_status[4:8]
    #gpu2_status = gpu_status[8:]
    #print(gpu2_status)
    gpu0_memory = int(gpu0_status[2].split('/')[0].split('M')[0].strip())
    gpu1_memory = int(gpu1_status[2].split('/')[0].split('M')[0].strip())
    #print(gpu_memory) 
    #获取当前0号GPU功率值：提取标签为2的元素，按照'/'为分隔符后提取标签为0的元素值再按照'M'为分隔符提取标签为0的元素值，返回值为int形式 
    gpu0_power = int(gpu0_status[1].split('   ')[-1].split('/')[0].split('W')[0].strip())
    gpu1_power = int(gpu1_status[1].split('   ')[-1].split('/')[0].split('W')[0].strip())
    #print(gpu_power)
    #获取0号GPU当前显存使用量
    #gpu_util = int(gpu1_status[3].split('   ')[1].split('%')[0].strip())
    #print(gpu_util)
    #获取0号GPU显存核心利用率
    return gpu0_power, gpu0_memory, gpu1_power, gpu1_memory# gpu_util
 
 
def narrow_setup(secs=900):  #间隔15分钟检测一次
    gpu0_power, gpu0_memory, gpu1_power, gpu1_memory = gpu_info()
    i = 0
    while not ((gpu0_memory < 1000 and gpu0_power < 70) or (gpu1_memory < 1000 and gpu1_power < 70)):  # 当功率，使用量，利用率都小于特定值才去退出循环
        gpu0_power, gpu0_memory, gpu1_power, gpu1_memory = gpu_info()
        i = i % 5
        symbol = 'monitoring: ' + '>' * i + ' ' * (10 - i - 1) + '|'
        gpu0_power_str = 'NO.0 GPU power:%d W |' % gpu0_power
        gpu0_memory_str = 'NO.0 GPU memory:%d MiB |' % gpu0_memory
        gpu1_power_str = 'NO.1 GPU power:%d W |' % gpu1_power
        gpu1_memory_str = 'NO.1 GPU memory:%d MiB |' % gpu1_memory
     #   gpu_util_str = 'gpu util:%d %% |' % gpu_util
        sys.stdout.write('\r' + gpu0_memory_str + ' ' + gpu0_power_str + ' ' + symbol+'\n' + gpu1_memory_str + ' ' + gpu1_power_str + ' ' + symbol)
        #sys.stdout.write('\r' + gpu1_memory_str + ' ' + gpu1_power_str + ' ' + symbol)
        #sys.stdout.write(obj+'\n')等价于print(obj)
        sys.stdout.flush()    #刷新输出
        time.sleep(secs)  #推迟调用线程的运行，通过参数指秒数，表示进程挂起的时间。
        i += 1
    if gpu0_memory < 1000 and gpu0_power < 70:
        print('\n' + cmd0)
        os.system(cmd0) #执行脚本
    else:
        print('\n' + cmd1)
        os.system(cmd1) #执行脚本
          
 
 
if __name__ == '__main__':
    narrow_setup()