GPU自动排队脚本

源代码来自知乎木盏老师的帖子https://zhuanlan.zhihu.com/p/90818278。

用于自动抢实验室的gpu，不用自己一直查看gpu是否空闲。

改成了多gpu版本，轮询所有gpu查看有无空闲，如果gpu空闲则指定这个gpu跑代码。

import os
import sys
import time

gpu_id = 'CUDA_VISIBLE_DEVICES='
cmd = 'python /home/kunyuan/codes/imFedSemi/train_main.py'
gpu_num = 4  #服务器上有几块卡就改成几
'''
def gpu_info():
    gpu_status = os.popen('nvidia-smi | grep %').read().split('|')
    gpu_memory = int(gpu_status[2].split('/')[0].split('M')[0].strip())
    gpu_power = int(gpu_status[1].split('   ')[-1].split('/')[0].split('W')[0].strip())
    return gpu_power, gpu_memory
'''

def gpu_info(gpu_index=0):
    info = os.popen('nvidia-smi|grep %').read().split('\n')[gpu_index].split('|')
    power = int(info[1].split()[-3][:-1])
    memory = int(info[2].split('/')[0].strip()[:-3])
    return power, memory

def narrow_setup(interval=2):
    gpu_power, gpu_memory = gpu_info()
    i = gpu_num
    while gpu_memory > 1000 or gpu_power > 20:  # set waiting condition
        i = i % gpu_num
        gpu_index = i
        gpu_power, gpu_memory = gpu_info(gpu_index)
        # i = i % 4
        symbol = 'monitoring: ' + '>' * i + ' ' * (10 - i - 1) + '|'
        gpu_index_str = 'gpu index %d |' % gpu_index
        gpu_power_str = 'gpu power:%d W |' % gpu_power
        gpu_memory_str = 'gpu memory:%d MiB |' % gpu_memory
        sys.stdout.write('\r' + gpu_index_str + ' ' + gpu_memory_str + ' ' + gpu_power_str + ' ' + symbol)
        sys.stdout.flush()
        time.sleep(interval)
        i += 1
    print('\n' + gpu_id + str(gpu_index) + ' ' + cmd)
    os.system(gpu_id + str(gpu_index) + ' ' + cmd)


if __name__ == '__main__':
    narrow_setup()

TODO:

改成自动计算gpu数量

可以指定多块gpu跑代码（当前只抢一块）