GPU检测显卡是否空闲排队程序

GPU检测显卡是否空闲排队程序

本程序特有地加入了检测部分显卡空闲时，可以使用部分显卡直接运行程序，更加实用
测试GPU为3090，不同型号可能略有差别

import os
import sys
import time
from IPython import embed
CUDA_cmd = 'CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7'
cmd = 'echo hello world'
NUM_GPU = 8

def gpu_info(num_gpu=NUM_GPU):
    info = os.popen('nvidia-smi|grep %').read().split('\n')
    power = []
    memory = []
    for i in range(num_gpu):
        memory.append(int(info[i].split('|')[2].split('/')[0][:-4]))
        power.append(int(info[i].split('|')[1].split('/')[0].split(' ')[-2][:-1]))
    # print('memory', memory)
    # print('power', power)
    return power, memory

def narrow_setup(interval=2, num_gpu=NUM_GPU):
    gpu_power_list, gpu_memory_list = gpu_info()
    k = 0

    # check if there is empty gpu initially
    empty_gpu = []
    for i in range(num_gpu):
        if gpu_memory_list[i] < 1000 and gpu_power_list[i] < 30:
            gpu_index = i
            print('gpu %d is empty' % gpu_index)
            empty_gpu.append(gpu_index)
    if len(empty_gpu) >=2:
        print('gpu', empty_gpu, 'is empty, ready to run')
        return
    else:
        print('no enough empty gpus now')
        while(1):
            available_gpu = []
            gpu_power_list, gpu_memory_list = gpu_info()
            for j in range(num_gpu):
                if gpu_memory_list[j] < 1000 and gpu_power_list[j] < 30:
                    gpu_index = j
                    print('gpu %d is empty' % gpu_index)
                    available_gpu.append(gpu_index)
                
                k = k % 5
                symbol = 'monitoring: ' + '>' * k + ' ' * (10 - k - 1) + '|'
                gpu_power_str = f'gpu {j} power:{gpu_power_list[j]} W |'
                gpu_memory_str = f'gpu {j} memory:{gpu_memory_list[j]} MiB |'
        
                sys.stdout.write('\r' + gpu_memory_str + ' ' + gpu_power_str + ' ' + symbol)
                sys.stdout.flush()
                time.sleep(interval)
                k += 1
            if len(available_gpu) >= 2:
                break
        
        CUDA_cmd = generate_cuda_visible_devices_string(available_gpu)
        print('\n' + cmd)
        os.system(CUDA_cmd)
        os.system(cmd)

def generate_cuda_visible_devices_string(int_list):
    cuda_visible_devices = 'CUDA_VISIBLE_DEVICES='
    elements = ','.join(str(element) for element in int_list)
    return cuda_visible_devices + elements

if __name__ == '__main__':
    narrow_setup()
    # gpu_info()