GPU 排队代码

import os
import sys
import time

cmd = 'nohup python -u main.py > nohup.out 2>&1 &'


def gpu_info(gpu_id: int):
    gpu_status = os.popen('nvidia-smi | grep %').read().split('|')
    gpu_memory = int(gpu_status[2+4*gpu_id].split('/')
                     [0].split('M')[0].strip())
    gpu_power = int(gpu_status[1+4*gpu_id].split('   ')
                    [-1].split('/')[0].split('W')[0].strip())
    return gpu_power, gpu_memory


def narrow_setup(interval=2):
    gpu_id = int(input("Please input which GPU you'd like to wait?"))
    gpu_power, gpu_memory = gpu_info(gpu_id)
    i = 0
    while gpu_memory > 500 or gpu_power > 150:  # set waiting condition
        gpu_power, gpu_memory = gpu_info(gpu_id)
        i = i % 5
        symbol = 'monitoring: ' + '>' * i + ' ' * (10 - i - 1) + '|'
        gpu_power_str = 'gpu power:%d W |' % gpu_power
        gpu_memory_str = 'gpu memory:%d MiB |' % gpu_memory
        sys.stdout.write('\r' + gpu_memory_str + ' ' +
                         gpu_power_str + ' ' + symbol)
        sys.stdout.flush()
        time.sleep(interval)
        i += 1
    print('\n'+str(gpu_memory)+str(gpu_power))
    print('\n' + cmd)
    os.system(cmd)


if __name__ == '__main__':
    narrow_setup()