Parse output of nvidia-smi to select the GPU with free memory

nvidia_gpu_tools.py

import subprocess
import re

def parse_free_memory(output):
    """
    $ nvidia-smi -q -d MEMORY -i 1

    ==============NVSMI LOG==============

    Timestamp                           : Wed Jan 31 14:13:22 2018
    Driver Version                      : 384.111

    Attached GPUs                       : 2
    GPU 00000000:06:00.0
        FB Memory Usage
            Total                       : 11172 MiB
            Used                        : 10 MiB
            Free                        : 11162 MiB
        BAR1 Memory Usage
            Total                       : 256 MiB
            Used                        : 2 MiB
            Free                        : 254 MiB
    """
    patc = re.compile(r"Free\s+:\s+(\d+)\s+MiB")

    # Get the first entry
    mo = patc.search(output)
    return int(mo.group(1))


def get_free_gpu_id(ngpus=2):
    """ nvidia-smi -q -d MEMORY -i """
    avail = [0] * ngpus
    nvidia_cmd = ["nvidia-smi", "-q", "-d", "MEMORY", "-i"]
    for i in range(ngpus):
        cmd = nvidia_cmd + [str(i)]
        output = subprocess.check_output(cmd).decode("utf-8")

        # Extract memory usage value
        avail[i] = parse_free_memory(output)

    return avail.index(max(avail))

if __name__ == "__main__":
    print(get_free_gpu_id())