start process if memory free

start_on_free_mem.py

import subprocess, os
import re
import click

def find_free_gpu(amount_of_space):
    #in Mb
    nvidia_ouput = subprocess.check_output(['nvidia-smi'])
    pid_memory = re.findall('(\d+)MiB\s/\s(\d+)MiB', str(nvidia_ouput), re.DOTALL)
    print(pid_memory)
    for gpu_num, (used, whole) in enumerate(pid_memory):
        if int(whole)-int(used)>amount_of_space:
            return gpu_num
    return None

@click.command()
@click.argument('mem') 
@click.argument('command') #command to run
def main(mem, command):
    ''' 
    mem - required memory value in Mb
    command - command to run
    e.g.
    python start_if_free.py 21000 '/opt/conda/bin/python train.py with batch_size=10 gpu_ids=0'
    be sure to set gpu id in you script to 0, because
    current script will change visabelity of gpus for torch
    '''
    print('Parsing command...')
    command_parts = command.split()
    print(command_parts)
    print('Looking for a memory ', mem, 'Mb')
    while True:
        gpu_id = find_free_gpu(int(mem))
        print(gpu_id)
        if gpu_id is not None:
            break
    print('Gpu id', gpu_id, ' is free')
    my_env = os.environ
    my_env['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
    subprocess.run([*command_parts], env=my_env)
if __name__=='__main__':
    main()