Python如何每秒获取NVIDIA GPU使用率以统计平均/峰值占用

问题根因 #

• 原代码中nvidia-smi -l 1的-l参数会让命令进入循环输出模式，不会自动终止，sp.check_output会一直阻塞等待命令结束，因此无法拿到返回数据
• 线程逻辑缺少终止控制，没有办法在模型运行结束后主动停止采集，也无法汇总过程中所有时间点的GPU数据

实现代码 #

import subprocess as sp
import time
from threading import Thread, Event

class GPUMonitor:
    def __init__(self, interval=1):
        self.interval = interval
        self._stop_event = Event()
        self.gpu_usage_data = []
        # 若需要采集显存占用，将utilization.gpu改为memory.used即可
        self.command = "nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits"
    
    def start(self):
        self.thread = Thread(target=self._monitor)
        self.thread.start()
    
    def stop(self):
        self._stop_event.set()
        self.thread.join()
        # 返回采集到的所有数据
        return self.gpu_usage_data
    
    def _monitor(self):
        while not self._stop_event.is_set():
            try:
                output = sp.check_output(self.command.split(), stderr=sp.STDOUT, text=True)
                # 多卡的话会返回多行，这里默认取第一张卡的数据，多卡可自行扩展
                usage = int(output.strip().split('\n')[0])
                self.gpu_usage_data.append(usage)
            except Exception as e:
                print(f"采集GPU数据失败: {e}")
            time.sleep(self.interval)

# 你的模型运行函数
def run_model():
    # 这里替换成你实际的模型运行代码，以下为10秒等待模拟模型运行
    time.sleep(10)

if __name__ == "__main__":
    # 初始化监控，间隔1秒采集一次
    monitor = GPUMonitor(interval=1)
    # 启动监控
    monitor.start()
    # 运行模型
    run_model()
    # 模型运行结束后停止监控，拿到所有采集数据
    gpu_data = monitor.stop()
    
    # 计算平均和峰值
    avg_usage = sum(gpu_data) / len(gpu_data) if gpu_data else 0
    peak_usage = max(gpu_data) if gpu_data else 0
    print(f"GPU平均使用率: {avg_usage:.2f}%")
    print(f"GPU峰值使用率: {peak_usage}%")
    print(f"所有采集数据: {gpu_data}")

逻辑说明 #

• 去掉了nvidia-smi的循环参数，改为每秒主动调用一次单次查询命令，避免永久阻塞
• 通过Event做线程终止控制，模型运行结束后主动通知监控线程停止
• 过程中所有采集到的GPU数据会存在列表里，结束后直接汇总计算平均和峰值即可
• 若需要采集显存占用，只需要将--query-gpu后的参数换成memory.used即可，返回数值单位为MB

内容的提问来源于stack exchange，提问作者haojie