get_gpu.py

def get_gpu_with_highest_free_memory():
    """
    Returns the index of the GPU with the highest free memory.

    This function sorts the GPUs by PCI bus ID and then returns the index of the GPU with the highest free memory.

    Returns:
        list: A list containing the index and free memory of the GPU with the highest free memory.

    Notes:
        This function uses the `nvidia-smi` command to get the GPU information and the `CUDA_DEVICE_ORDER` environment variable to sort the GPUs by PCI bus ID.

    Example Output:
        ['PCI_ID', 'FREE MEMORY']
        [2, 10240]
        [0, 8192]
        [1, 4096]
        [3, 102]

    In this example, the GPU with index 2 has the highest free memory of 10240 MB. And so the output
    is [2, 10240]
    """
    import subprocess
    import os
    # sort GPUs by PCI_BUS_ID
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    # Run the nvidia-smi command and capture the output
    output = subprocess.check_output(['nvidia-smi', '--query-gpu=index,memory.free', '--format=csv'])
    # Split the output into lines
    lines = output.decode('utf-8').strip().split('\n')
    gpu_tuples = [i.strip().split(',') for i in lines[1:]]
    for tup in gpu_tuples:
        tup[0] = int(tup[0])
        tup[1] = int(tup[1].strip().split(" ")[0])
    gpu_tuples.sort(key=lambda x: x[1], reverse=True)
    print(['PCI_ID', 'FREE MEMORY'])
    for i in gpu_tuples:
        print(i)
    return gpu_tuples[0]