Парсинг данных из nvidia-smi

Задача состоит в том, чтобы парсить текущее энергопотребление для каждой карточки из nvidia-smi,

Сделал следующий скрипт

import subprocess
def get_gpu_info():
    sp = subprocess.Popen(['nvidia-smi', '-q'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
    inf_p1 = []
    inf_p2 = []
    two_gpu_flag = False
    out_str = sp.communicate()
    out_list = out_str[0].decode("utf-8").split("\n")
    new_line = []
    for i in range(len(out_list)):
        if out_list[i] == "":
            new_line.append(i)
    split_arr = []
    for i in range(len(new_line)-1):
        if new_line[i] - new_line[i+1] < -5:
            split_arr.append(new_line[i])
            split_arr.append(new_line[i+1])
    out_list_1 = []
    out_list_2 = []
    if len(split_arr) > 2:
            two_gpu_flag = True
    for i in range(1000):
        sp = subprocess.Popen(['nvidia-smi', '-q'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
        out_str = sp.communicate()
        out_list = out_str[0].decode("utf-8").split("\n")
       
        if two_gpu_flag:
            two_gpu_flag = True
            out_list_1 = out_list[split_arr[0]:split_arr[1]]
            out_list_2 = out_list[split_arr[2]:split_arr[3]]


            out_dict_gpu1 = {}
            out_dict_gpu2 = {}

            for item in out_list_1:
                try:
                    key, val = item.split(':')
                    key, val = key.strip(), val.strip()
                    out_dict_gpu1[key] = val
                except:
                    pass
            for item in out_list_2:
                try:
                    key, val = item.split(':')
                    key, val = key.strip(), val.strip()
                    out_dict_gpu2[key] = val
                except:
                    pass

            device_id_gpu1 = out_dict_gpu1['Minor Number']
            power_gpu1 = float(out_dict_gpu1[ 'Power Draw'].split(" ")[0])
            inf_p1.append(power_gpu1)
            device_id_gpu2 = out_dict_gpu2['Minor Number']
            power_gpu2 = float(out_dict_gpu2[ 'Power Draw'].split(" ")[0])
            
            inf_p2.append(power_gpu2)

        else:
            out_dict_gpu1 = {}
            for item in out_list:
                try:
                    key, val = item.split(':')
                    key, val = key.strip(), val.strip()
                    out_dict_gpu1[key] = val
                except:
                    pass

            device_id_gpu1 = out_dict_gpu1['Minor Number']
            power_gpu1 = float(out_dict_gpu1[ 'Power Draw'].split(" ")[0])
            inf_p1.append(power_gpu1)
    if two_gpu_flag:
        power_gpu1 = sum(inf_p1) / len(inf_p1)
        power_gpu2 = sum(inf_p2) / len(inf_p2)
        return {device_id_gpu1 : power_gpu1 , device_id_gpu2 : power_gpu2}
    else:
        power_gpu1 = sum(inf_p1) / len(inf_p1)
        return {device_id_gpu1 : power_gpu1 }

print(get_gpu_info())
    
    

Он хоть и работает, но строка out_str = sp.communicate() отрабатывает очень долго

Подскажите, пожалуйста, какой-нибудь удобный и более быстрый способ парсинга данных через python из nvidia-smi


Ответы (1 шт):

Автор решения: andrew
from subprocess import Popen, PIPE
from xml.etree.ElementTree import fromstring

import time

def get_gpu_info(*args, **kwargs):
    p = Popen(["nvidia-smi", "-q", "-x"], stdout=PIPE)
    outs, errors = p.communicate()
    xml = fromstring(outs)
    datas = []
    driver_version = xml.findall("driver_version")[0].text
    cuda_version = xml.findall("cuda_version")[0].text

    for gpu_id, gpu in enumerate(xml.getiterator("gpu")):
        gpu_data = {}

        minor_number = gpu.findall("minor_number")[0].text
        power_R = gpu.findall("power_readings")[0]
        power_draw = power_R.findall("power_draw")[0].text


        gpu_data["minor_number"] = minor_number
        gpu_data["power"] = power_draw
        datas.append(gpu_data)
    return datas 


if __name__ == '__main__':
    c = time.time()
    get_gpu_info()
    print(time.time() - c )
→ Ссылка