Visual Storyteller | Navy

利用python实时监控显卡cuda使用率实现自动关机

背景

在进行3D渲染任务时，往往需要长时间占用计算资源，尤其是在夜间或无人值守的情况下。如何确保渲染任务完成后系统能够自动关闭，以节约电力并减少硬件损耗？

因为渲染占用的主要是显卡的GPU，我目前想到的解决办法是监控显卡使用率：当显卡使用率低于某个值时，就判定渲染完成。同时检测鼠标是否移动，当鼠标未移动时表明无人操作。满足条件时，脚本会自动关机。该脚本支持多显卡，目前测试两块显卡的场景可以正常监控。

环境要求

核心逻辑

以下是Python脚本的主要逻辑：

  1. 检测CUDA使用
    • 使用pycuda库查询GPU的当前使用情况，判断是否有进程占用CUDA。
  2. 检测用户活动
    • 使用pyautogui库检查鼠标和键盘是否有活动。
  3. 自动关机
    • 当检测到CUDA未使用且无用户活动时，调用os.system命令执行关机操作。

代码实现

将以下代码存储为.py格式的文件：

import time
import os
import tkinter as tk
from tkinter import messagebox
from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetUtilizationRates, nvmlDeviceGetCount
import subprocess
import ctypes
from ctypes import wintypes

# 初始化NVML库
nvmlInit()
device_count = nvmlDeviceGetCount()

# 定义变量
low_usage_threshold = 50  # CUDA使用率低于50%时判定空闲
time_threshold = 300  # 5分钟（以秒为单位）
shutdown_countdown = 60  # 关机倒计时（秒）
mouse_idle_threshold = 100  # 鼠标移动距离阈值
last_mouse_position = (0, 0)
last_mouse_time = time.time()

# 检查鼠标是否移动
def get_mouse_position():
    pt = ctypes.wintypes.POINT()
    ctypes.windll.user32.GetCursorPos(ctypes.byref(pt))
    return pt.x, pt.y

def check_mouse_idle():
    global last_mouse_position, last_mouse_time
    current_position = get_mouse_position()
    distance = ((current_position[0] - last_mouse_position[0]) ** 2 + (current_position[1] - last_mouse_position[1]) ** 2) ** 0.5
    if distance < mouse_idle_threshold:
        idle_duration = time.time() - last_mouse_time
    else:
        idle_duration = 0
        last_mouse_position = current_position
        last_mouse_time = time.time()
    return idle_duration

def check_gpu_usage_and_mouse_idle():
    global last_mouse_time
    all_below_threshold = True
    for i in range(device_count):
        handle = nvmlDeviceGetHandleByIndex(i)
        utilization = nvmlDeviceGetUtilizationRates(handle)
        gpu_usage = utilization.gpu
        print(f"当前CUDA使用率 (显卡{i}): {gpu_usage}%")
        if gpu_usage >= low_usage_threshold:
            all_below_threshold = False
            break

    if all_below_threshold:
        idle_duration = check_mouse_idle()
        if idle_duration >= 60:
            print("条件满足，系统准备关机...")
            show_shutdown_warning()
    else:
        print("CUDA使用率高于阈值，或鼠标活动中。")

def show_shutdown_warning():
    def countdown(count):
        label['text'] = f"系统将在{count}秒后自动关机。请保存你的工作。"
        if count > 0:
            root.after(1000, countdown, count - 1)
        else:
            shutdown_computer()

    root = tk.Tk()
    root.title("关机警告")
    root.geometry("400x200")
    label = tk.Label(root, text="", font=('Helvetica', 12))
    label.pack(pady=20)
    countdown(shutdown_countdown)
    root.mainloop()

def shutdown_computer():
    print("系统正在关机")
    subprocess.run(['shutdown', '/s', '/t', '1'], check=True)

while True:
    check_gpu_usage_and_mouse_idle()
    time.sleep(60)

将以下内容存储为.bat格式的文件，与上面的.py脚本放在同一目录，双击运行.bat文件，可以自动调用Python代码；也可以不用下面的.bat文件直接运行.py文件。

@echo off
echo Navigating to script directory...
cd %USERPROFILE%\Desktop
echo Running Python script...
C:\Users\admin\AppData\Local\Programs\Python\Python312\python.exe monitor_cuda_usage.py
echo Python script executed.
pause

说明

  1. 安装依赖 使用以下命令安装必要的Python库：

    pip install pycuda psutil pyautogui

  2. 调整参数

    • 修改检测时间间隔，例如将time.sleep(60)改为其他秒数。
    • 调整用户空闲时间阈值，例如将idle_duration >= 60改为其他时间。