import subprocess
import time
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# 初始化存储CUDA利用率的数据和时间戳
cuda_utilization = []
timestamps = []

def get_cuda_utilization():
    # 使用nvidia-smi命令获取CUDA利用率
    result = subprocess.run(['nvidia-smi', '--query-gpu=utilization.gpu', '--format=csv,noheader,nounits'], 
                            stdout=subprocess.PIPE, text=True)
    # 解析输出并返回利用率
    return int(result.stdout.strip())

def update_plot(frame):
    # 获取当前的CUDA利用率
    utilization = get_cuda_utilization()
    current_time = time.time()  # 获取当前时间戳
    
    # 添加当前数据和时间戳
    cuda_utilization.append(utilization)
    timestamps.append(current_time)
    
    # 移除60秒之前的数据
    while current_time - timestamps[0] > 60:
        cuda_utilization.pop(0)
        timestamps.pop(0)
    
    # 清空当前图形并重新绘制
    plt.cla()
    plt.plot(timestamps, cuda_utilization, label='CUDA Utilization (%)')
    
    # 设置X轴范围为最近60秒
    plt.xlim(current_time - 60, current_time)
    plt.ylim(0, 100)
    plt.xlabel('Time (s)')
    plt.ylabel('CUDA Utilization (%)')
    plt.title('Real-time CUDA Utilization (Last 60 seconds)')
    plt.legend()

# 创建图形
fig, ax = plt.subplots()

# 使用FuncAnimation实现实时更新
ani = animation.FuncAnimation(fig, update_plot, interval=1000)

# 显示图形
plt.show()