Parse nvidia-smi by Python

gpu_info.py

from subprocess import Popen, PIPE
from xml.etree.ElementTree import fromstring


def main():
    p = Popen(['nvidia-smi', '-q', '-x'], stdout=PIPE)
    outs, errors = p.communicate()
    xml = fromstring(outs)
    num_gpus = int(xml.getiterator('attached_gpus')[0].text)
    results = []
    for gpu_id, gpu in enumerate(xml.getiterator('gpu')):
        gpu_data = {}

        name = gpu.getiterator('product_name')[0].text
        gpu_data['name'] = name

        # get memory
        memory_usage = gpu.getiterator('fb_memory_usage')[0]
        total_memory = memory_usage.getiterator('total')[0].text
        used_memory = memory_usage.getiterator('used')[0].text
        free_memory = memory_usage.getiterator('free')[0].text
        gpu_data['memory'] = {
            'total': total_memory,
            'used_memory': used_memory,
            'free_memory': free_memory
        }

        # get utilization
        utilization = gpu.getiterator('utilization')[0]
        gpu_util = utilization.getiterator('gpu_util')[0].text
        memory_util = utilization.getiterator('memory_util')[0].text
        gpu_data['utilization'] = {
            'gpu_util': gpu_util,
            'memory_util': memory_util
        }

        # processes
        processes = gpu.getiterator('processes')[0]
        infos = []
        for info in processes.getiterator('process_info'):
            pid = info.getiterator('pid')[0].text
            process_name = info.getiterator('process_name')[0].text
            used_memory = info.getiterator('used_memory')[0].text
            infos.append({
                'pid': pid,
                'process_name': process_name,
                'used_memory': used_memory
            })
        gpu_data['processes'] = infos

        results.append(gpu_data)
    print(results)


if __name__ == '__main__':
    main()

Hi @chenbo-ponyai ,
Here is an edited version to fix the issue above

NOTE: this might not be highly efficient in terms of execution time but it worked! :)

from subprocess import Popen, PIPE
from xml.etree.ElementTree import fromstring


def main():
    p = Popen(['nvidia-smi', '-q', '-x'], stdout=PIPE)
    outs, errors = p.communicate()
    xml = fromstring(outs)
    num_gpus = int(list(xml.iter('attached_gpus'))[0].text)
    results = []
    for gpu_id, gpu in enumerate(xml.iter('gpu')):
        gpu_data = {}
        
        name = list(gpu.iter('product_name'))[0].text
        gpu_data['name'] = name
        
        # get memory
        memory_usage = list(gpu.iter('fb_memory_usage'))[0]
        total_memory = list(memory_usage.iter('total'))[0].text
        used_memory = list(memory_usage.iter('used'))[0].text
        free_memory = list(memory_usage.iter('free'))[0].text
        gpu_data['memory'] = {
            'total': total_memory,
            'used_memory': used_memory,
            'free_memory': free_memory
        }
        
        # get utilization
        utilization = list(gpu.iter('utilization'))[0]
        gpu_util = list(utilization.iter('gpu_util'))[0].text
        memory_util = list(utilization.iter('memory_util'))[0].text
        gpu_data['utilization'] = {
            'gpu_util': gpu_util,
            'memory_util': memory_util
        }
        
        # processes
        processes = list(gpu.iter('processes'))[0]
        infos = []
        for info in list(processes.iter('process_info')):
            pid = list(info.iter('pid'))[0].text
            process_name = list(info.iter('process_name'))[0].text
            used_memory = list(info.iter('used_memory'))[0].text
            infos.append({
                'pid': pid,
                'process_name': process_name,
                'used_memory': used_memory
            })
        gpu_data['processes'] = infos
        
        results.append(gpu_data)
    return results

if __name__ == '__main__':
    res = main()
    print(res[0]) # GPU #1
    print(res[1]) # GPU #2
    print(res[2]) # GPU #3
    print(res[3]) # GPU #4