GPU Memory

gpu.py

import subprocess


def get_gpu_memory():
    # source: https://stackoverflow.com/a/59571639
    command = "nvidia-smi --query-gpu=memory.free --format=csv"
    memory_free_info = subprocess.check_output(command.split()).decode('ascii').split('\n')[:-1][1:]
    memory_free_values = [round(int(x.split()[0]) / 1000) for i, x in enumerate(memory_free_info)]
    return memory_free_values


def check_gpu_memory(min_num_gpus=3, min_memory=30):
    gpus = get_gpu_memory()

    if not len(gpus) > min_num_gpus or not all([True if gpu > min_memory else False for gpu in gpus]):
        raise Exception("There is something wrong with GPU allocation on the server side!")